"""Async Google Generative AI (Gemini) driver."""

from __future__ import annotations

import logging
import os
import uuid
from collections.abc import AsyncIterator
from typing import Any

from google import genai
from google.genai import types

from .async_base import AsyncDriver
from ..infra.cost_mixin import CostMixin
from .google_driver import GoogleDriver

logger = logging.getLogger(__name__)


class AsyncGoogleDriver(CostMixin, AsyncDriver):
    """Async driver for Google's Generative AI API (Gemini)."""

    supports_json_mode = True
    supports_json_schema = True
    supports_vision = True
    supports_tool_use = True
    supports_streaming = True

    MODEL_PRICING = GoogleDriver.MODEL_PRICING
    _PRICING_UNIT = 1_000_000

    def __init__(self, api_key: str | None = None, model: str = "gemini-1.5-pro"):
        self.api_key = api_key or os.getenv("GOOGLE_API_KEY")
        if not self.api_key:
            raise ValueError("Google API key not found. Set GOOGLE_API_KEY env var or pass api_key to constructor")
        self.model = model
        self._client = genai.Client(api_key=self.api_key)
        self.options: dict[str, Any] = {}

    def _calculate_cost_chars(self, prompt_chars: int, completion_chars: int) -> float:
        """Calculate cost from character counts (same logic as sync GoogleDriver)."""
        from ..infra.model_rates import get_model_rates

        live_rates = get_model_rates("google", self.model)
        if live_rates:
            est_prompt_tokens = prompt_chars / 4
            est_completion_tokens = completion_chars / 4
            prompt_cost = (est_prompt_tokens / 1_000_000) * live_rates["input"]
            completion_cost = (est_completion_tokens / 1_000_000) * live_rates["output"]
        else:
            model_pricing = self.MODEL_PRICING.get(self.model, {"prompt": 0, "completion": 0})
            prompt_cost = (prompt_chars / 1_000_000) * model_pricing["prompt"]
            completion_cost = (completion_chars / 1_000_000) * model_pricing["completion"]
        return round(prompt_cost + completion_cost, 6)

    def _extract_usage_metadata(self, response: Any, messages: list[dict[str, Any]]) -> dict[str, Any]:
        """Extract token counts from response, falling back to character estimation."""
        usage = getattr(response, "usage_metadata", None)
        if usage:
            prompt_tokens = getattr(usage, "prompt_token_count", 0) or 0
            completion_tokens = getattr(usage, "candidates_token_count", 0) or 0
            total_tokens = getattr(usage, "total_token_count", 0) or (prompt_tokens + completion_tokens)
            cost = self._calculate_cost("google", self.model, prompt_tokens, completion_tokens)
        else:
            # Fallback: estimate from character counts
            total_prompt_chars = 0
            for msg in messages:
                c = msg.get("content", "")
                if isinstance(c, str):
                    total_prompt_chars += len(c)
                elif isinstance(c, list):
                    for part in c:
                        if isinstance(part, str):
                            total_prompt_chars += len(part)
                        elif isinstance(part, dict) and "text" in part:
                            total_prompt_chars += len(part["text"])
            completion_chars = len(response.text) if response.text else 0
            prompt_tokens = total_prompt_chars // 4
            completion_tokens = completion_chars // 4
            total_tokens = prompt_tokens + completion_tokens
            cost = self._calculate_cost_chars(total_prompt_chars, completion_chars)

        return {
            "prompt_tokens": prompt_tokens,
            "completion_tokens": completion_tokens,
            "total_tokens": total_tokens,
            "cost": round(cost, 6),
        }

    supports_messages = True

    def _prepare_messages(self, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
        from .vision_helpers import _prepare_google_vision_messages

        return _prepare_google_vision_messages(messages)

    def _build_generation_args(
        self, messages: list[dict[str, Any]], options: dict[str, Any] | None = None
    ) -> tuple[Any, dict[str, Any]]:
        """Parse messages and options into (contents, config_dict)."""
        merged_options = self.options.copy()
        if options:
            merged_options.update(options)

        config_dict: dict[str, Any] = {}

        if "temperature" in merged_options:
            config_dict.setdefault("temperature", merged_options["temperature"])
        if "max_tokens" in merged_options:
            config_dict.setdefault("max_output_tokens", merged_options["max_tokens"])
        if "top_p" in merged_options:
            config_dict.setdefault("top_p", merged_options["top_p"])
        if "top_k" in merged_options:
            config_dict.setdefault("top_k", merged_options["top_k"])

        # Forward explicit generation_config entries
        for k, v in merged_options.get("generation_config", {}).items():
            config_dict.setdefault(k, v)

        # Forward safety_settings
        safety_settings = merged_options.get("safety_settings")
        if safety_settings:
            config_dict["safety_settings"] = safety_settings

        # Native JSON mode support
        if merged_options.get("json_mode"):
            config_dict["response_mime_type"] = "application/json"
            json_schema = merged_options.get("json_schema")
            if json_schema:
                config_dict["response_schema"] = json_schema

        # Convert messages to Gemini format
        system_instruction = None
        contents: list[dict[str, Any]] = []
        for msg in messages:
            role = msg.get("role", "user")
            content = msg.get("content", "")
            if role == "system":
                system_instruction = content if isinstance(content, str) else str(content)
            else:
                gemini_role = "model" if role == "assistant" else "user"
                if msg.get("_vision_parts"):
                    contents.append({"role": gemini_role, "parts": content})
                else:
                    contents.append({"role": gemini_role, "parts": [content]})

        # For a single message, unwrap only if it has exactly one string part
        if len(contents) == 1:
            parts = contents[0]["parts"]
            if len(parts) == 1 and isinstance(parts[0], str):
                gen_input = parts[0]
            else:
                gen_input = contents
        else:
            gen_input = contents

        if system_instruction:
            config_dict["system_instruction"] = system_instruction

        return gen_input, config_dict

    async def generate(self, prompt: str, options: dict[str, Any] | None = None) -> dict[str, Any]:
        messages = [{"role": "user", "content": prompt}]
        return await self._do_generate(messages, options)

    async def generate_messages(self, messages: list[dict[str, str]], options: dict[str, Any]) -> dict[str, Any]:
        return await self._do_generate(self._prepare_messages(messages), options)

    async def _do_generate(
        self, messages: list[dict[str, str]], options: dict[str, Any] | None = None
    ) -> dict[str, Any]:
        gen_input, config_dict = self._build_generation_args(messages, options)

        # Validate capabilities against models.dev metadata
        self._validate_model_capabilities(
            "google",
            self.model,
            using_json_schema=bool((options or {}).get("json_schema")),
        )

        try:
            config = types.GenerateContentConfig(**config_dict)
            response = await self._client.aio.models.generate_content(
                model=self.model,
                contents=gen_input,
                config=config,
            )

            if not response.text:
                raise ValueError("Empty response from model")

            usage_meta = self._extract_usage_metadata(response, messages)

            meta = {
                **usage_meta,
                "raw_response": response.prompt_feedback if hasattr(response, "prompt_feedback") else None,
                "model_name": self.model,
            }

            return {"text": response.text, "meta": meta}

        except Exception as e:
            logger.error(f"Google API request failed: {e}")
            raise RuntimeError(f"Google API request failed: {e}") from e

    # ------------------------------------------------------------------
    # Tool use
    # ------------------------------------------------------------------

    async def generate_messages_with_tools(
        self,
        messages: list[dict[str, Any]],
        tools: list[dict[str, Any]],
        options: dict[str, Any],
    ) -> dict[str, Any]:
        """Generate a response that may include tool/function calls (async)."""
        model = options.get("model", self.model)
        self._validate_model_capabilities("google", model, using_tool_use=True)

        gen_input, config_dict = self._build_generation_args(self._prepare_messages(messages), options)

        # Convert tools from OpenAI format to Gemini function declarations
        function_declarations = []
        for t in tools:
            if "type" in t and t["type"] == "function":
                fn = t["function"]
                decl = types.FunctionDeclaration(
                    name=fn["name"],
                    description=fn.get("description", ""),
                )
                params = fn.get("parameters")
                if params:
                    decl = types.FunctionDeclaration(
                        name=fn["name"],
                        description=fn.get("description", ""),
                        parameters_json_schema=params,
                    )
                function_declarations.append(decl)
            elif "name" in t:
                params = t.get("parameters") or t.get("input_schema")
                if params:
                    decl = types.FunctionDeclaration(
                        name=t["name"],
                        description=t.get("description", ""),
                        parameters_json_schema=params,
                    )
                else:
                    decl = types.FunctionDeclaration(
                        name=t["name"],
                        description=t.get("description", ""),
                    )
                function_declarations.append(decl)

        config_dict["tools"] = [types.Tool(function_declarations=function_declarations)]

        try:
            config = types.GenerateContentConfig(**config_dict)
            response = await self._client.aio.models.generate_content(
                model=self.model,
                contents=gen_input,
                config=config,
            )

            usage_meta = self._extract_usage_metadata(response, messages)
            meta = {
                **usage_meta,
                "raw_response": response.prompt_feedback if hasattr(response, "prompt_feedback") else None,
                "model_name": self.model,
            }

            text = ""
            tool_calls_out: list[dict[str, Any]] = []
            stop_reason = "stop"

            for candidate in response.candidates:
                for part in candidate.content.parts:
                    if hasattr(part, "text") and part.text:
                        text += part.text
                    if hasattr(part, "function_call") and part.function_call.name:
                        fc = part.function_call
                        tool_calls_out.append(
                            {
                                "id": str(uuid.uuid4()),
                                "name": fc.name,
                                "arguments": dict(fc.args) if fc.args else {},
                            }
                        )

                finish_reason = getattr(candidate, "finish_reason", None)
                if finish_reason is not None:
                    reason_map = {1: "stop", 2: "max_tokens", 3: "safety", 4: "recitation", 5: "other"}
                    stop_reason = reason_map.get(finish_reason, "stop")

            if tool_calls_out:
                stop_reason = "tool_use"

            return {
                "text": text,
                "meta": meta,
                "tool_calls": tool_calls_out,
                "stop_reason": stop_reason,
            }

        except Exception as e:
            logger.error(f"Google API tool call request failed: {e}")
            raise RuntimeError(f"Google API tool call request failed: {e}") from e

    # ------------------------------------------------------------------
    # Streaming
    # ------------------------------------------------------------------

    async def generate_messages_stream(
        self,
        messages: list[dict[str, Any]],
        options: dict[str, Any],
    ) -> AsyncIterator[dict[str, Any]]:
        """Yield response chunks via Gemini async streaming API."""
        gen_input, config_dict = self._build_generation_args(self._prepare_messages(messages), options)

        try:
            config = types.GenerateContentConfig(**config_dict)
            response = self._client.aio.models.generate_content_stream(
                model=self.model,
                contents=gen_input,
                config=config,
            )

            full_text = ""
            async for chunk in response:
                chunk_text = getattr(chunk, "text", None) or ""
                if chunk_text:
                    full_text += chunk_text
                    yield {"type": "delta", "text": chunk_text}

            # After iteration completes, usage_metadata should be available
            usage_meta = self._extract_usage_metadata(response, messages)

            yield {
                "type": "done",
                "text": full_text,
                "meta": {
                    **usage_meta,
                    "raw_response": {},
                    "model_name": self.model,
                },
            }

        except Exception as e:
            logger.error(f"Google API streaming request failed: {e}")
            raise RuntimeError(f"Google API streaming request failed: {e}") from e
