# Copyright 2023 LiveKit, Inc.
#

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations

import json
import os
from dataclasses import dataclass
from typing import Any, cast

import google.auth.credentials
from google.auth._default_async import default_async
from google.genai import Client, types
from google.genai.errors import APIError, ClientError, ServerError
from livekit.agents import APIConnectionError, APIStatusError, llm, utils
from livekit.agents.llm import ToolChoice, utils as llm_utils
from livekit.agents.types import (
    DEFAULT_API_CONNECT_OPTIONS,
    NOT_GIVEN,
    APIConnectOptions,
    NotGivenOr,
)
from livekit.agents.utils import is_given

from .log import logger
from .models import ChatModels
from .utils import create_tools_config, to_response_format
from .version import __version__


def _is_gemini_3_model(model: str) -> bool:
    """Check if model is Gemini 3 series"""
    return "gemini-3" in model.lower() or model.lower().startswith("gemini-3")


def _is_gemini_3_flash_model(model: str) -> bool:
    """Check if model is Gemini 3 Flash"""
    m = model.lower()
    return m.startswith("gemini-3") and "flash" in m


def _requires_thought_signatures(model: str) -> bool:
    """Check if model requires thought_signature handling for multi-turn function calling.

    Gemini 2.5+ models require thought signatures to be stored from responses and
    passed back in subsequent requests for proper multi-turn function calling.
    """
    if _is_gemini_3_model(model):
        return True
    model_lower = model.lower()
    return "gemini-2.5" in model_lower or model_lower.startswith("gemini-2.5")


@dataclass
class _LLMOptions:
    model: ChatModels | str
    temperature: NotGivenOr[float]
    tool_choice: NotGivenOr[ToolChoice]
    vertexai: NotGivenOr[bool]
    project: NotGivenOr[str]
    location: NotGivenOr[str]
    max_output_tokens: NotGivenOr[int]
    top_p: NotGivenOr[float]
    top_k: NotGivenOr[float]
    presence_penalty: NotGivenOr[float]
    frequency_penalty: NotGivenOr[float]
    thinking_config: NotGivenOr[types.ThinkingConfigOrDict]
    retrieval_config: NotGivenOr[types.RetrievalConfigOrDict]
    automatic_function_calling_config: NotGivenOr[types.AutomaticFunctionCallingConfigOrDict]
    http_options: NotGivenOr[types.HttpOptions]
    seed: NotGivenOr[int]
    safety_settings: NotGivenOr[list[types.SafetySettingOrDict]]
    service_tier: NotGivenOr[types.ServiceTier]
    cached_content: NotGivenOr[str]
    media_resolution: NotGivenOr[types.MediaResolution]


BLOCKED_REASONS = [
    types.FinishReason.SAFETY,
    types.FinishReason.SPII,
    types.FinishReason.PROHIBITED_CONTENT,
    types.FinishReason.BLOCKLIST,
    types.FinishReason.LANGUAGE,
    types.FinishReason.RECITATION,
]


class LLM(llm.LLM):
    def __init__(
        self,
        *,
        model: ChatModels | str = "gemini-2.5-flash",
        api_key: NotGivenOr[str] = NOT_GIVEN,
        vertexai: NotGivenOr[bool] = NOT_GIVEN,
        project: NotGivenOr[str] = NOT_GIVEN,
        location: NotGivenOr[str] = NOT_GIVEN,
        temperature: NotGivenOr[float] = NOT_GIVEN,
        max_output_tokens: NotGivenOr[int] = NOT_GIVEN,
        top_p: NotGivenOr[float] = NOT_GIVEN,
        top_k: NotGivenOr[float] = NOT_GIVEN,
        presence_penalty: NotGivenOr[float] = NOT_GIVEN,
        frequency_penalty: NotGivenOr[float] = NOT_GIVEN,
        tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
        thinking_config: NotGivenOr[types.ThinkingConfigOrDict] = NOT_GIVEN,
        retrieval_config: NotGivenOr[types.RetrievalConfigOrDict] = NOT_GIVEN,
        automatic_function_calling_config: NotGivenOr[
            types.AutomaticFunctionCallingConfigOrDict
        ] = NOT_GIVEN,
        http_options: NotGivenOr[types.HttpOptions] = NOT_GIVEN,
        seed: NotGivenOr[int] = NOT_GIVEN,
        safety_settings: NotGivenOr[list[types.SafetySettingOrDict]] = NOT_GIVEN,
        service_tier: NotGivenOr[types.ServiceTier] = NOT_GIVEN,
        cached_content: NotGivenOr[str] = NOT_GIVEN,
        media_resolution: NotGivenOr[types.MediaResolution] = NOT_GIVEN,
        credentials: google.auth.credentials.Credentials | None = None,
    ) -> None:
        """
        Create a new instance of Google GenAI LLM.

        Environment Requirements:
        - For VertexAI: Set the `GOOGLE_APPLICATION_CREDENTIALS` environment variable to the path of the service account key file or use any of the other Google Cloud auth methods.
        The Google Cloud project and location can be set via `project` and `location` arguments or the environment variables
        `GOOGLE_CLOUD_PROJECT` and `GOOGLE_CLOUD_LOCATION`. By default, the project is inferred from the service account key file,
        and the location defaults to "us-central1".
        - For Google Gemini API: Set the `api_key` argument or the `GOOGLE_API_KEY` environment variable.

        Args:
            model (ChatModels | str, optional): The model name to use. Defaults to "gemini-2.0-flash-001".
            api_key (str, optional): The API key for Google Gemini. If not provided, it attempts to read from the `GOOGLE_API_KEY` environment variable.
            vertexai (bool, optional): Whether to use VertexAI. If not provided, it attempts to read from the `GOOGLE_GENAI_USE_VERTEXAI` environment variable. Defaults to False.
                project (str, optional): The Google Cloud project to use (only for VertexAI). Defaults to None.
                location (str, optional): The location to use for VertexAI API requests. Defaults value is "us-central1".
            temperature (float, optional): Sampling temperature for response generation. Defaults to 0.8.
            max_output_tokens (int, optional): Maximum number of tokens to generate in the output. Defaults to None.
            top_p (float, optional): The nucleus sampling probability for response generation. Defaults to None.
            top_k (int, optional): The top-k sampling value for response generation. Defaults to None.
            presence_penalty (float, optional): Penalizes the model for generating previously mentioned concepts. Defaults to None.
            frequency_penalty (float, optional): Penalizes the model for repeating words. Defaults to None.
            tool_choice (ToolChoice, optional): Specifies whether to use tools during response generation. Defaults to "auto".
            thinking_config (ThinkingConfigOrDict, optional): The thinking configuration for response generation. Defaults to None.
            retrieval_config (RetrievalConfigOrDict, optional): The retrieval configuration for response generation. Defaults to None.
            automatic_function_calling_config (AutomaticFunctionCallingConfigOrDict, optional): The automatic function calling configuration for response generation. Defaults to None.
            http_options (HttpOptions, optional): The HTTP options to use for the session.
            seed (int, optional): Random seed for reproducible generation. Defaults to None.
            safety_settings (list[SafetySettingOrDict], optional): Safety settings for content filtering. Defaults to None.
            service_tier (types.ServiceTier, optional): The service tier for the request (e.g. types.ServiceTier.PRIORITY). Defaults to None.
            cached_content (str, optional): Resource name of an explicit context cache to attach to every request from this LLM instance, e.g. ``"cachedContents/abc123"`` for the Gemini API or ``"projects/<project>/locations/<location>/cachedContents/abc123"`` for VertexAI. The cache must already exist — create it via ``client.caches.create(...)`` and pass the returned ``name``. Gemini rejects ``generateContent`` requests that combine ``cached_content`` with ``system_instruction``, ``tools``, or ``tool_config``, so when this option is set the plugin bakes those fields out of every outgoing request; the cache resource itself must contain whichever of them the model needs (typically the system prompt and the tool schemas). Useful for long-lived static prefixes where implicit caching is unreliable. See https://ai.google.dev/gemini-api/docs/caching for details and minimum prefix-token requirements. Defaults to None.
            media_resolution (types.MediaResolution, optional): The media resolution for the request. Defaults to None.
        """  # noqa: E501
        super().__init__()
        gcp_project = project if is_given(project) else os.environ.get("GOOGLE_CLOUD_PROJECT")
        gcp_location: str | None = (
            location
            if is_given(location)
            else os.environ.get("GOOGLE_CLOUD_LOCATION") or "us-central1"
        )
        use_vertexai = (
            vertexai
            if is_given(vertexai)
            else os.environ.get("GOOGLE_GENAI_USE_VERTEXAI", "0").lower() in ["true", "1"]
        )
        gemini_api_key = api_key if is_given(api_key) else os.environ.get("GOOGLE_API_KEY")

        if use_vertexai:
            if not gcp_project:
                _, gcp_project = default_async(  # type: ignore
                    scopes=["https://www.googleapis.com/auth/cloud-platform"]
                )
            if not gcp_project or not gcp_location:
                raise ValueError(
                    "Project is required for VertexAI via project kwarg or GOOGLE_CLOUD_PROJECT environment variable"  # noqa: E501
                )
            gemini_api_key = None  # VertexAI does not require an API key

        else:
            gcp_project = None
            gcp_location = None
            if credentials is not None:
                logger.warning(
                    "'credentials' is only applicable to VertexAI and will be ignored for the Gemini API"
                )
                credentials = None
            if not gemini_api_key:
                raise ValueError(
                    "API key is required for Google API either via api_key or GOOGLE_API_KEY environment variable"  # noqa: E501
                )

        # Validate thinking_config
        if is_given(thinking_config):
            _thinking_budget = None
            _thinking_level = None
            if isinstance(thinking_config, dict):
                _thinking_budget = thinking_config.get("thinking_budget")
                _thinking_level = thinking_config.get("thinking_level")
            elif isinstance(thinking_config, types.ThinkingConfig):
                _thinking_budget = thinking_config.thinking_budget
                _thinking_level = getattr(thinking_config, "thinking_level", None)

            if _thinking_budget is not None:
                if not isinstance(_thinking_budget, int):
                    raise ValueError("thinking_budget inside thinking_config must be an integer")

        self._opts = _LLMOptions(
            model=model,
            temperature=temperature,
            tool_choice=tool_choice,
            vertexai=use_vertexai,
            project=project,
            location=location,
            max_output_tokens=max_output_tokens,
            top_p=top_p,
            top_k=top_k,
            presence_penalty=presence_penalty,
            frequency_penalty=frequency_penalty,
            thinking_config=thinking_config,
            retrieval_config=retrieval_config,
            automatic_function_calling_config=automatic_function_calling_config,
            http_options=http_options,
            seed=seed,
            safety_settings=safety_settings,
            service_tier=service_tier,
            cached_content=cached_content,
            media_resolution=media_resolution,
        )
        self._client = Client(
            api_key=gemini_api_key,
            vertexai=use_vertexai,
            project=gcp_project,
            location=gcp_location,
            credentials=credentials,
        )
        # Store thought_signatures for Gemini 2.5+ multi-turn function calling
        self._thought_signatures: dict[str, bytes] = {}

    @property
    def model(self) -> str:
        return self._opts.model

    @property
    def provider(self) -> str:
        if self._client.vertexai:
            return "Vertex AI"
        else:
            return "Gemini"

    def chat(
        self,
        *,
        chat_ctx: llm.ChatContext,
        tools: list[llm.Tool] | None = None,
        conn_options: APIConnectOptions = DEFAULT_API_CONNECT_OPTIONS,
        parallel_tool_calls: NotGivenOr[bool] = NOT_GIVEN,
        tool_choice: NotGivenOr[ToolChoice] = NOT_GIVEN,
        response_format: NotGivenOr[
            types.SchemaUnion | type[llm_utils.ResponseFormatT]
        ] = NOT_GIVEN,
        extra_kwargs: NotGivenOr[dict[str, Any]] = NOT_GIVEN,
    ) -> LLMStream:
        extra = {}

        if is_given(extra_kwargs):
            extra.update(extra_kwargs)

        tool_choice = tool_choice if is_given(tool_choice) else self._opts.tool_choice
        retrieval_config = (
            self._opts.retrieval_config if is_given(self._opts.retrieval_config) else None
        )
        if isinstance(retrieval_config, dict):
            retrieval_config = types.RetrievalConfig.model_validate(retrieval_config)

        if is_given(tool_choice):
            gemini_tool_choice: types.ToolConfig
            if isinstance(tool_choice, dict) and tool_choice.get("type") == "function":
                gemini_tool_choice = types.ToolConfig(
                    function_calling_config=types.FunctionCallingConfig(
                        mode=types.FunctionCallingConfigMode.ANY,
                        allowed_function_names=[tool_choice["function"]["name"]],
                    ),
                    retrieval_config=retrieval_config,
                )
                extra["tool_config"] = gemini_tool_choice
            elif tool_choice == "required":
                tool_names = []
                for tool in tools or []:
                    if isinstance(tool, (llm.FunctionTool, llm.RawFunctionTool)):
                        tool_names.append(tool.info.name)

                gemini_tool_choice = types.ToolConfig(
                    function_calling_config=types.FunctionCallingConfig(
                        mode=types.FunctionCallingConfigMode.ANY,
                        allowed_function_names=tool_names or None,
                    ),
                    retrieval_config=retrieval_config,
                )
                extra["tool_config"] = gemini_tool_choice
            elif tool_choice == "auto":
                gemini_tool_choice = types.ToolConfig(
                    function_calling_config=types.FunctionCallingConfig(
                        mode=types.FunctionCallingConfigMode.AUTO,
                    ),
                    retrieval_config=retrieval_config,
                )
                extra["tool_config"] = gemini_tool_choice
            elif tool_choice == "none":
                gemini_tool_choice = types.ToolConfig(
                    function_calling_config=types.FunctionCallingConfig(
                        mode=types.FunctionCallingConfigMode.NONE,
                    ),
                    retrieval_config=retrieval_config,
                )
                extra["tool_config"] = gemini_tool_choice
        elif retrieval_config:
            extra["tool_config"] = types.ToolConfig(
                retrieval_config=retrieval_config,
            )

        if is_given(response_format):
            extra["response_schema"] = to_response_format(response_format)
            extra["response_mime_type"] = "application/json"

        if is_given(self._opts.temperature):
            extra["temperature"] = self._opts.temperature
        if is_given(self._opts.max_output_tokens):
            extra["max_output_tokens"] = self._opts.max_output_tokens
        if is_given(self._opts.top_p):
            extra["top_p"] = self._opts.top_p
        if is_given(self._opts.top_k):
            extra["top_k"] = self._opts.top_k
        if is_given(self._opts.presence_penalty):
            extra["presence_penalty"] = self._opts.presence_penalty
        if is_given(self._opts.frequency_penalty):
            extra["frequency_penalty"] = self._opts.frequency_penalty
        if is_given(self._opts.seed):
            extra["seed"] = self._opts.seed

        # Handle thinking_config based on model version
        if is_given(self._opts.thinking_config):
            is_gemini_3 = _is_gemini_3_model(self._opts.model)
            is_gemini_3_flash = _is_gemini_3_flash_model(self._opts.model)
            thinking_cfg = self._opts.thinking_config

            # Extract both parameters
            _budget = None
            _level: str | types.ThinkingLevel | None = None
            if isinstance(thinking_cfg, dict):
                _budget = thinking_cfg.get("thinking_budget")
                _level = thinking_cfg.get("thinking_level")
            elif isinstance(thinking_cfg, types.ThinkingConfig):
                _budget = thinking_cfg.thinking_budget
                _level = getattr(thinking_cfg, "thinking_level", None)

            if is_gemini_3:
                # Gemini 3: only support thinking_level
                if _budget is not None and _level is None:
                    logger.warning(
                        f"Model {self._opts.model} is Gemini 3 which does not support thinking_budget. "
                        "Please use thinking_level ('low' or 'high') instead. Ignoring thinking_budget."
                    )
                if _level is None:
                    # If no thinking_level is provided, use the fastest thinking level
                    if is_gemini_3_flash:
                        _level = "minimal"
                    else:
                        _level = "low"
                # Use thinking_level only (pass as dict since SDK may not have this field yet)
                extra["thinking_config"] = {"thinking_level": _level}

            else:
                # Gemini 2.5 and earlier: only support thinking_budget
                if _level is not None and _budget is None:
                    raise ValueError(
                        f"Model {self._opts.model} does not support thinking_level. "
                        "Please use thinking_budget (int) instead for Gemini 2.5 and earlier models."
                    )
                if _budget is not None:
                    # Use thinking_budget only
                    extra["thinking_config"] = types.ThinkingConfig(thinking_budget=_budget)
                else:
                    # Pass through original config if no specific handling needed
                    extra["thinking_config"] = self._opts.thinking_config

        if is_given(self._opts.automatic_function_calling_config):
            extra["automatic_function_calling"] = self._opts.automatic_function_calling_config

        if is_given(self._opts.safety_settings):
            extra["safety_settings"] = self._opts.safety_settings

        if is_given(self._opts.service_tier):
            extra["service_tier"] = self._opts.service_tier

        if is_given(self._opts.cached_content):
            extra["cached_content"] = self._opts.cached_content

        if is_given(self._opts.media_resolution):
            extra["media_resolution"] = self._opts.media_resolution

        return LLMStream(
            self,
            client=self._client,
            model=self._opts.model,
            chat_ctx=chat_ctx,
            tools=tools or [],
            conn_options=conn_options,
            extra_kwargs=extra,
        )


class LLMStream(llm.LLMStream):
    def __init__(
        self,
        llm_v: LLM,
        *,
        client: Client,
        model: str | ChatModels,
        chat_ctx: llm.ChatContext,
        conn_options: APIConnectOptions,
        tools: list[llm.Tool],
        extra_kwargs: dict[str, Any],
    ) -> None:
        super().__init__(llm_v, chat_ctx=chat_ctx, tools=tools, conn_options=conn_options)
        self._client = client
        self._model = model
        self._llm: LLM = llm_v
        self._extra_kwargs = extra_kwargs
        self._tool_ctx = llm.ToolContext(tools)

    async def _run(self) -> None:
        retryable = True
        request_id = utils.shortuuid()

        try:
            # Pass thought_signatures for Gemini 2.5+ multi-turn function calling
            thought_sigs = (
                self._llm._thought_signatures if _requires_thought_signatures(self._model) else None
            )
            turns_dict, extra_data = self._chat_ctx.to_provider_format(
                format="google", thought_signatures=thought_sigs
            )

            turns = [types.Content.model_validate(turn) for turn in turns_dict]
            tool_context = llm.ToolContext(self._tools)
            tools_config = create_tools_config(tool_context, _only_single_type=True)
            # Gemini's API rejects `generateContent` requests that pass
            # `cached_content` together with `system_instruction`, `tools`,
            # or `tool_config` — those fields must live INSIDE the
            # CachedContent resource, not on the request. The application
            # bakes them into the cache via `client.caches.create(...)`;
            # here we just suppress the duplicates on the outgoing request
            # whenever a cache is attached.
            using_cache = "cached_content" in self._extra_kwargs
            if tools_config and not using_cache:
                self._extra_kwargs["tools"] = tools_config
            elif using_cache:
                dropped = [k for k in ("tools", "tool_config") if k in self._extra_kwargs]
                if tools_config and "tools" not in dropped:
                    dropped.append("tools")
                if extra_data.system_messages:
                    dropped.append("system_instruction")
                if dropped:
                    logger.warning(
                        "dropping %s from Gemini request because cached_content=%r is set; "
                        "these fields must be baked into the CachedContent resource",
                        dropped,
                        self._extra_kwargs.get("cached_content"),
                    )
                self._extra_kwargs.pop("tools", None)
                self._extra_kwargs.pop("tool_config", None)
            if is_given(self._llm._opts.http_options):
                http_options = self._llm._opts.http_options.model_copy()
                if http_options.timeout is None:
                    http_options.timeout = int(self._conn_options.timeout * 1000)
            else:
                http_options = types.HttpOptions(timeout=int(self._conn_options.timeout * 1000))

            headers = dict(http_options.headers or {})
            headers["x-goog-api-client"] = f"livekit-agents/{__version__}"
            http_options.headers = headers
            config = types.GenerateContentConfig(
                system_instruction=(
                    None
                    if using_cache
                    else (
                        [types.Part(text=content) for content in extra_data.system_messages]
                        if extra_data.system_messages
                        else None
                    )
                ),
                http_options=http_options,
                **self._extra_kwargs,
            )

            stream = await self._client.aio.models.generate_content_stream(
                model=self._model,
                contents=cast(types.ContentListUnion, turns),
                config=config,
            )

            response_generated = False
            finish_reason: types.FinishReason | None = None
            async for response in stream:
                if response.prompt_feedback:
                    raise APIStatusError(
                        response.prompt_feedback.model_dump_json(),
                        retryable=False,
                        request_id=request_id,
                    )

                if response.usage_metadata is not None:
                    usage = response.usage_metadata
                    self._event_ch.send_nowait(
                        llm.ChatChunk(
                            id=request_id,
                            usage=llm.CompletionUsage(
                                completion_tokens=usage.candidates_token_count or 0,
                                prompt_tokens=usage.prompt_token_count or 0,
                                prompt_cached_tokens=usage.cached_content_token_count or 0,
                                total_tokens=usage.total_token_count or 0,
                            ),
                        )
                    )

                if not response.candidates:
                    continue

                if len(response.candidates) > 1:
                    logger.warning(
                        "gemini llm: there are multiple candidates in the response, returning response from the first one."  # noqa: E501
                    )

                candidate = response.candidates[0]

                if candidate.finish_reason is not None:
                    finish_reason = candidate.finish_reason
                    if candidate.finish_reason in BLOCKED_REASONS:
                        raise APIStatusError(
                            f"generation blocked by gemini: {candidate.finish_reason}",
                            retryable=False,
                            request_id=request_id,
                        )

                if not candidate.content or not candidate.content.parts:
                    continue

                for part in candidate.content.parts:
                    chat_chunk = self._parse_part(request_id, part)
                    response_generated = True
                    if chat_chunk is not None:
                        retryable = False
                        self._event_ch.send_nowait(chat_chunk)

            if not response_generated:
                raise APIStatusError(
                    "no response generated",
                    retryable=retryable,
                    request_id=request_id,
                    body=f"finish reason: {finish_reason}",
                )

        except ClientError as e:
            raise APIStatusError(
                "gemini llm: client error",
                status_code=e.code,
                body=f"{e.message} {e.status}",
                request_id=request_id,
                retryable=True if e.code in {429, 499} else False,
            ) from e
        except ServerError as e:
            raise APIStatusError(
                "gemini llm: server error",
                status_code=e.code,
                body=f"{e.message} {e.status}",
                request_id=request_id,
                retryable=retryable,
            ) from e
        except APIError as e:
            raise APIStatusError(
                "gemini llm: api error",
                status_code=e.code,
                body=f"{e.message} {e.status}",
                request_id=request_id,
                retryable=retryable,
            ) from e
        except (APIStatusError, APIConnectionError):
            raise
        except Exception as e:
            raise APIConnectionError(
                f"gemini llm: error generating content {str(e)}",
                retryable=retryable,
            ) from e

    def _parse_part(self, id: str, part: types.Part) -> llm.ChatChunk | None:
        if part.function_call:
            tool_call = llm.FunctionToolCall(
                arguments=json.dumps(part.function_call.args),
                name=part.function_call.name,
                call_id=part.function_call.id or utils.shortuuid("function_call_"),
            )

            # Store thought_signature for Gemini 2.5+ multi-turn function calling
            if (
                _requires_thought_signatures(self._model)
                and hasattr(part, "thought_signature")
                and part.thought_signature
            ):
                self._llm._thought_signatures[tool_call.call_id] = part.thought_signature

            chat_chunk = llm.ChatChunk(
                id=id,
                delta=llm.ChoiceDelta(
                    role="assistant",
                    tool_calls=[tool_call],
                    content=None,
                ),
            )
            return chat_chunk

        if not part.text:
            return None

        return llm.ChatChunk(
            id=id,
            delta=llm.ChoiceDelta(content=part.text, role="assistant"),
        )