import asyncio
import json
from typing import Awaitable, Callable

from google import genai
from google.genai import types as genai_types

from .base import BaseModel


class GoogleGeminiModel(BaseModel):
    def __init__(self, model_name: str, api_key: str):
        super().__init__()
        self.model_name = model_name
        self.client = genai.Client(api_key=api_key)

    def get_response(self, input_data: str) -> str:
        """
        Gets a response from the Google Gemini model.

        Args:
            input_data (str): The input data for the model.

        Returns:
            str: The response generated by the model.
        """
        try:
            response = self.client.models.generate_content(
                model=self.model_name,
                contents=input_data,
            )
            return response.text if response and hasattr(response, "text") else "No response generated."
        except Exception as e:
            return f"Error: {str(e)}"

    def get_response_with_tools(self, messages: list, tools: list) -> dict:
        """
        Call Gemini with native function calling.

        Converts OpenAI-format messages and tools to Gemini format, then
        normalises the response back to the shared dict schema.

        Returns:
            dict with keys:
              text: str | None
              tool_calls: list[{id, name, args}] | None
              usage: dict | None
        """
        gemini_tools, system_instruction, contents = self._prepare_gemini_request(
            messages, tools
        )

        config_kwargs = {}
        if gemini_tools[0].function_declarations:
            config_kwargs["tools"] = gemini_tools
        if system_instruction:
            config_kwargs["system_instruction"] = system_instruction

        response = self.client.models.generate_content(
            model=self.model_name,
            contents=contents,
            config=genai_types.GenerateContentConfig(**config_kwargs),
        )
        return self._parse_response(response)

    async def stream_response_with_tools(
        self,
        messages: list,
        tools: list,
        on_token: Callable[[str], Awaitable[None]],
    ) -> dict:
        """
        Stream text tokens via on_token, then return the assembled result dict.

        The google-genai SDK exposes only a synchronous streaming iterator
        (generate_content_stream). We collect chunks in a background thread via
        asyncio.to_thread, then emit on_token calls from the event loop thread
        for each text chunk. Function-call parts are collected from the final
        chunk that carries a complete candidate.

        Returns:
            dict with keys:
              text: str | None
              tool_calls: list[{id, name, args}] | None
              usage: dict | None
        """
        gemini_tools, system_instruction, contents = self._prepare_gemini_request(
            messages, tools
        )

        config_kwargs = {}
        if gemini_tools[0].function_declarations:
            config_kwargs["tools"] = gemini_tools
        if system_instruction:
            config_kwargs["system_instruction"] = system_instruction

        # Stream chunks from a background thread, emitting each text token
        # immediately via on_token so the SSE stream isn't held up until the
        # full response is collected.
        loop = asyncio.get_running_loop()
        text_parts: list[str] = []
        final_chunk_holder: list = []

        def _stream_chunks():
            last_chunk = None
            for chunk in self.client.models.generate_content_stream(
                model=self.model_name,
                contents=contents,
                config=genai_types.GenerateContentConfig(**config_kwargs),
            ):
                last_chunk = chunk
                if hasattr(chunk, "text") and chunk.text:
                    token = chunk.text
                    text_parts.append(token)
                    asyncio.run_coroutine_threadsafe(on_token(token), loop).result()
            if last_chunk is not None:
                final_chunk_holder.append(last_chunk)

        await asyncio.to_thread(_stream_chunks)

        # The final chunk carries the fully assembled candidate with function calls.
        if final_chunk_holder:
            return self._parse_response(final_chunk_holder[0])

        return {"text": "".join(text_parts) or None, "tool_calls": None, "usage": None}

    # ------------------------------------------------------------------
    # Private helpers
    # ------------------------------------------------------------------

    def _prepare_gemini_request(self, messages: list, tools: list) -> tuple:
        """Convert OpenAI-format messages/tools to Gemini format.

        Returns:
            (gemini_tools, system_instruction, contents)
        """
        # --- Convert OpenAI tool schema → Gemini function declarations ---
        function_declarations = []
        for t in tools:
            fn = t["function"]
            params = fn.get("parameters", {"type": "object", "properties": {}})
            function_declarations.append(
                genai_types.FunctionDeclaration(
                    name=fn["name"],
                    description=fn.get("description", ""),
                    parameters=params,
                )
            )
        gemini_tools = [genai_types.Tool(function_declarations=function_declarations)]

        # --- Convert OpenAI messages → Gemini contents ---
        system_instruction = None
        contents = []

        for msg in messages:
            role = msg["role"]
            content = msg.get("content") or ""

            if role == "system":
                system_instruction = content
                continue

            if role == "user":
                contents.append(genai_types.Content(role="user", parts=[genai_types.Part(text=content)]))

            elif role == "assistant":
                if msg.get("tool_calls"):
                    # Assistant requested tool calls
                    parts = []
                    if content:
                        parts.append(genai_types.Part(text=content))
                    for tc in msg["tool_calls"]:
                        args = tc["function"]["arguments"]
                        parts.append(
                            genai_types.Part(
                                function_call=genai_types.FunctionCall(
                                    name=tc["function"]["name"],
                                    args=json.loads(args) if isinstance(args, str) else args,
                                )
                            )
                        )
                    contents.append(genai_types.Content(role="model", parts=parts))
                else:
                    contents.append(genai_types.Content(role="model", parts=[genai_types.Part(text=content)]))

            elif role == "tool":
                # Tool result — Gemini expects function_response as a user-role turn
                contents.append(
                    genai_types.Content(
                        role="user",
                        parts=[
                            genai_types.Part(
                                function_response=genai_types.FunctionResponse(
                                    name=msg.get("name", "tool"),
                                    response={"result": content},
                                )
                            )
                        ],
                    )
                )

        return gemini_tools, system_instruction, contents

    def _parse_response(self, response) -> dict:
        """Normalise a Gemini GenerateContentResponse into the shared dict schema."""
        usage = None
        if hasattr(response, "usage_metadata") and response.usage_metadata:
            usage = {
                "input_tokens": getattr(response.usage_metadata, "prompt_token_count", None),
                "output_tokens": getattr(response.usage_metadata, "candidates_token_count", None),
            }

        # Check for function calls in the response
        candidate = response.candidates[0] if (hasattr(response, "candidates") and response.candidates) else None
        if candidate and candidate.content and candidate.content.parts:
            tool_calls = []
            text_parts = []
            for part in candidate.content.parts:
                if part.function_call:
                    fc = part.function_call
                    tool_calls.append({
                        "id": f"gemini-{fc.name}-{len(tool_calls)}",
                        "name": fc.name,
                        "args": dict(fc.args) if fc.args else {},
                    })
                elif hasattr(part, "text") and part.text:
                    text_parts.append(part.text)

            if tool_calls:
                return {"text": None, "tool_calls": tool_calls, "usage": usage}
            return {"text": "".join(text_parts), "tool_calls": None, "usage": usage}

        text = response.text if hasattr(response, "text") and response.text else ""
        return {"text": text, "tool_calls": None, "usage": usage}
