import json
import os

import pydantic
import structlog
from google import genai
import google.genai.types as genai_types
from google.genai.errors import APIError

from talemate.client.base import (
    ClientBase,
    ErrorAction,
    ExtraField,
    ParameterReroute,
    CommonDefaults,
)
from talemate.client.registry import register
from talemate.client.remote import (
    RemoteServiceMixin,
    EndpointOverride,
    EndpointOverrideMixin,
    endpoint_override_extra_fields,
)
from talemate.config.schema import Client as BaseClientConfig
from talemate.emit import emit
from talemate.util import count_tokens

__all__ = [
    "GoogleClient",
]
log = structlog.get_logger("talemate")

# Edit this to add new models / remove old models
SUPPORTED_MODELS = [
    "gemini-1.0-pro",
    "gemini-1.5-pro-preview-0409",
    "gemini-1.5-flash",
    "gemini-1.5-flash-8b",
    "gemini-1.5-pro",
    "gemini-2.0-flash",
    "gemini-2.0-flash-lite",
    "gemini-2.5-flash-lite-preview-06-17",
    "gemini-2.5-flash-preview-05-20",
    "gemini-2.5-flash",
    "gemini-2.5-pro-preview-06-05",
    "gemini-2.5-pro",
    "gemini-3-pro-preview",
]


class Defaults(EndpointOverride, CommonDefaults, pydantic.BaseModel):
    max_token_length: int = 16384
    model: str = "gemini-2.0-flash"
    disable_safety_settings: bool = False
    double_coercion: str = None


class ClientConfig(EndpointOverride, BaseClientConfig):
    disable_safety_settings: bool = False


MIN_THINKING_TOKENS = 512


@register()
class GoogleClient(EndpointOverrideMixin, RemoteServiceMixin, ClientBase):
    """
    Google client for generating text.
    """

    client_type = "google"
    conversation_retries = 0
    decensor_enabled = True
    config_cls = ClientConfig

    class Meta(ClientBase.Meta):
        name_prefix: str = "Google"
        title: str = "Google"
        manual_model: bool = True
        manual_model_choices: list[str] = SUPPORTED_MODELS
        requires_prompt_template: bool = False
        defaults: Defaults = Defaults()
        unified_api_key_config_path: str = "google.api_key"
        extra_fields: dict[str, ExtraField] = {
            "disable_safety_settings": ExtraField(
                name="disable_safety_settings",
                type="bool",
                label="Disable Safety Settings",
                required=False,
                description="Disable Google's safety settings for responses generated by the model.",
            ),
        }
        extra_fields.update(endpoint_override_extra_fields())

    def __init__(self, model="gemini-2.0-flash", **kwargs):
        self.setup_status = None
        self.model_instance = None
        self.google_credentials_read = False
        self.google_project_id = None
        super().__init__(**kwargs)

    @property
    def disable_safety_settings(self):
        return self.client_config.disable_safety_settings

    @property
    def min_reason_tokens(self) -> int:
        return MIN_THINKING_TOKENS

    @property
    def can_be_coerced(self) -> bool:
        return not self.reason_enabled

    @property
    def google_credentials(self):
        path = self.google_credentials_path
        if not path:
            return None
        with open(path) as f:
            return json.load(f)

    @property
    def google_credentials_path(self):
        return self.config.google.gcloud_credentials_path

    @property
    def google_location(self):
        return self.config.google.gcloud_location

    @property
    def google_api_key(self):
        return self.config.google.api_key

    @property
    def vertexai_ready(self) -> bool:
        return all(
            [
                self.google_credentials_path,
                self.google_location,
            ]
        )

    @property
    def developer_api_ready(self) -> bool:
        return all(
            [
                self.google_api_key,
            ]
        )

    @property
    def using(self) -> str:
        if self.developer_api_ready:
            return "API"
        if self.vertexai_ready:
            return "VertexAI"
        return "Unknown"

    @property
    def ready(self):
        # all google settings must be set
        return (
            self.vertexai_ready
            or self.developer_api_ready
            or self.endpoint_override_base_url_configured
        )

    @property
    def safety_settings(self):
        if not self.disable_safety_settings:
            return None

        safety_settings = [
            genai_types.SafetySetting(
                category="HARM_CATEGORY_SEXUALLY_EXPLICIT",
                threshold="BLOCK_NONE",
            ),
            genai_types.SafetySetting(
                category="HARM_CATEGORY_DANGEROUS_CONTENT",
                threshold="BLOCK_NONE",
            ),
            genai_types.SafetySetting(
                category="HARM_CATEGORY_HARASSMENT",
                threshold="BLOCK_NONE",
            ),
            genai_types.SafetySetting(
                category="HARM_CATEGORY_HATE_SPEECH",
                threshold="BLOCK_NONE",
            ),
            genai_types.SafetySetting(
                category="HARM_CATEGORY_CIVIC_INTEGRITY",
                threshold="BLOCK_NONE",
            ),
        ]

        return safety_settings

    @property
    def http_options(self) -> genai_types.HttpOptions | None:
        if not self.endpoint_override_base_url_configured:
            return None

        return genai_types.HttpOptions(base_url=self.base_url)

    @property
    def thinking_config(self) -> genai_types.ThinkingConfig | None:
        if not self.reason_enabled:
            return None

        return genai_types.ThinkingConfig(
            thinking_budget=self.validated_reason_tokens,
            include_thoughts=True,
        )

    @property
    def supported_parameters(self):
        return [
            "temperature",
            "top_p",
            "top_k",
            ParameterReroute(
                talemate_parameter="max_tokens", client_parameter="max_output_tokens"
            ),
            ParameterReroute(
                talemate_parameter="stopping_strings", client_parameter="stop_sequences"
            ),
        ]

    @property
    def requires_reasoning_pattern(self) -> bool:
        return False

    def emit_status(self, processing: bool = None):
        error_action = None
        if processing is not None:
            self.processing = processing

        if self.ready:
            status = "busy" if self.processing else "idle"
            model_name = self.model_name
        else:
            status = "error"
            model_name = "Setup incomplete"
            error_action = ErrorAction(
                title="Setup Google API credentials",
                action_name="openAppConfig",
                icon="mdi-key-variant",
                arguments=[
                    "application",
                    "google_api",
                ],
            )

        if not self.model_name:
            status = "error"
            model_name = "No model loaded"

        self.current_status = status
        data = {
            "double_coercion": self.double_coercion,
            "error_action": error_action.model_dump() if error_action else None,
            "meta": self.Meta().model_dump(),
            "enabled": self.enabled,
        }
        data.update(self._common_status_data())
        self.populate_extra_fields(data)

        if self.using == "VertexAI":
            details = f"{model_name} (VertexAI)"
        else:
            details = model_name

        emit(
            "client_status",
            message=self.client_type,
            id=self.name,
            details=details,
            status=status if self.enabled else "disabled",
            data=data,
        )

    def set_client_base_url(self, base_url: str | None):
        if getattr(self, "client", None):
            try:
                self.client.http_options.base_url = base_url
            except Exception as e:
                log.error(
                    "Error setting client base URL", error=e, client=self.client_type
                )

    def make_client(self) -> genai.Client:
        if self.google_credentials_path:
            os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = self.google_credentials_path
        if self.vertexai_ready and not self.developer_api_ready:
            return genai.Client(
                vertexai=True,
                project=self.google_project_id,
                location=self.google_location,
            )
        else:
            return genai.Client(
                api_key=self.api_key or None, http_options=self.http_options
            )

    def response_tokens(self, response: str):
        """Return token count for a response which may be a string or SDK object."""
        return count_tokens(response)

    def prompt_tokens(self, prompt: str):
        return count_tokens(prompt)

    def clean_prompt_parameters(self, parameters: dict):
        super().clean_prompt_parameters(parameters)

        # if top_k is 0, remove it
        if "top_k" in parameters and parameters["top_k"] == 0:
            del parameters["top_k"]

    async def generate(self, prompt: str, parameters: dict, kind: str):
        """
        Generates text from the given prompt and parameters.
        """

        if not self.ready:
            raise Exception("Google setup incomplete")

        client = self.make_client()

        if self.can_be_coerced:
            prompt, coercion_prompt = self.split_prompt_for_coercion(prompt)
        else:
            coercion_prompt = None

        human_message = prompt.strip()
        system_message = self.get_system_message(kind)

        contents = [
            genai_types.Content(
                role="user",
                parts=[
                    genai_types.Part.from_text(
                        text=human_message,
                    )
                ],
            )
        ]

        if coercion_prompt:
            log.debug("Adding coercion pre-fill", coercion_prompt=coercion_prompt)
            contents.append(
                genai_types.Content(
                    role="model",
                    parts=[
                        genai_types.Part.from_text(
                            text=coercion_prompt,
                        )
                    ],
                )
            )

        self.log.debug(
            "generate",
            model=self.model_name,
            base_url=self.base_url,
            prompt=prompt[:128] + " ...",
            parameters=parameters,
            system_message=system_message,
            disable_safety_settings=self.disable_safety_settings,
            safety_settings=self.safety_settings,
            thinking_config=self.thinking_config,
        )

        try:
            # Use streaming so we can update_Request_tokens incrementally
            stream = await client.aio.models.generate_content_stream(
                model=self.model_name,
                contents=contents,
                config=genai_types.GenerateContentConfig(
                    safety_settings=self.safety_settings,
                    http_options=self.http_options,
                    thinking_config=self.thinking_config,
                    **parameters,
                ),
            )

            response = ""
            reasoning = ""
            # https://ai.google.dev/gemini-api/docs/thinking#summaries
            async for chunk in stream:
                try:
                    if not chunk:
                        continue

                    if not chunk.candidates:
                        continue

                    if not chunk.candidates[0].content.parts:
                        continue

                    for part in chunk.candidates[0].content.parts:
                        if not part.text:
                            continue
                        if part.thought:
                            reasoning += part.text
                        else:
                            response += part.text
                        self.update_request_tokens(count_tokens(part.text))
                except Exception as e:
                    log.error("error processing chunk", e=e, chunk=chunk)
                    continue

            if reasoning:
                self._reasoning_response = reasoning

            # Store total token accounting for prompt/response
            self._returned_prompt_tokens = self.prompt_tokens(prompt)
            self._returned_response_tokens = self.response_tokens(response)

            log.debug("generated response", response=response)

            return response

        except APIError as e:
            self.log.error("generate error", e=e)
            emit("status", message="google API: API Error", status="error")
            return ""
        except Exception:
            raise
