import asyncio
import json
from typing import Dict, Any, List, Optional

import vertexai
from google.api_core.exceptions import NotFound
from vertexai.generative_models import FunctionDeclaration, GenerativeModel, Tool

from app.core.settings import settings
from app.models.tool_model import ToolDefinition


# Essa classe encapsula a integracao com o Vertex AI:
# inicializacao, cache de modelos e serializacao das tools.
class LLMService:
    _vertex_initialized = False
    _models: dict[str, GenerativeModel] = {}
    _vertex_tools_cache: dict[str, Optional[List[Tool]]] = {}

    def __init__(self):
        """Inicializa o cliente Vertex AI e define modelos de fallback."""
        if not LLMService._vertex_initialized:
            vertexai.init(
                project=settings.google_project_id,
                location=settings.google_location,
            )
            LLMService._vertex_initialized = True

        configured = settings.vertex_model_name.strip()
        fallback_models = ["gemini-2.5-pro", "gemini-2.5-flash", "gemini-2.0-flash-001"]
        self.model_names = [configured] + [m for m in fallback_models if m != configured]

    def build_vertex_tools(self, tools: List[ToolDefinition]) -> Optional[List[Tool]]:
        """Converte tools internas para o formato esperado pelo Vertex AI."""
        # Vertex espera uma lista de Tool, com function_declarations agrupadas em um unico Tool.
        if not tools:
            return None

        cache_key = json.dumps(
            [
                {
                    "name": tool.name,
                    "description": tool.description,
                    "parameters": tool.parameters,
                }
                for tool in tools
            ],
            sort_keys=True,
            ensure_ascii=True,
            separators=(",", ":"),
        )
        cached = LLMService._vertex_tools_cache.get(cache_key)
        if cached is not None:
            return cached

        function_declarations = [
            FunctionDeclaration(
                name=tool.name,
                description=tool.description,
                parameters=tool.parameters,
            )
            for tool in tools
        ]

        vertex_tools = [Tool(function_declarations=function_declarations)]
        LLMService._vertex_tools_cache[cache_key] = vertex_tools
        return vertex_tools

    def _get_model(self, model_name: str) -> GenerativeModel:
        model = LLMService._models.get(model_name)
        if model is None:
            model = GenerativeModel(model_name)
            LLMService._models[model_name] = model
        return model

    def _extract_response_payload(self, response) -> Dict[str, Any]:
        candidate = response.candidates[0] if getattr(response, "candidates", None) else None
        content = getattr(candidate, "content", None)
        parts = list(getattr(content, "parts", None) or [])

        tool_call = None
        text_parts: list[str] = []
        for part in parts:
            function_call = getattr(part, "function_call", None)
            if function_call is not None and tool_call is None:
                tool_call = {
                    "name": function_call.name,
                    "arguments": dict(function_call.args),
                }
            text_value = getattr(part, "text", None)
            if isinstance(text_value, str) and text_value.strip():
                text_parts.append(text_value)

        response_text = "\n".join(text_parts).strip() or None
        return {
            "response": response_text,
            "tool_call": tool_call,
        }

    async def generate_response(
        self,
        message: str,
        tools: List[ToolDefinition],
        history: List[Dict[str, Any]] = None,
    ) -> Dict[str, Any]:
        """Gera resposta textual ou chamada de tool a partir da mensagem do usuario."""
        vertex_tools = self.build_vertex_tools(tools)

        response = None
        last_error = None

        # Tenta o modelo configurado e cai para nomes alternativos
        # quando o principal nao estiver disponivel no projeto/regiao.
        for model_name in self.model_names:
            try:
                model = self._get_model(model_name)
                chat = model.start_chat(history=history or [])
                send_kwargs = {"tools": vertex_tools} if vertex_tools else {}
                response = await asyncio.to_thread(chat.send_message, message, **send_kwargs)
                break
            except NotFound as err:
                last_error = err
                LLMService._models.pop(model_name, None)
                continue

        if response is None:
            if last_error:
                raise RuntimeError(
                    f"Nenhum modelo Vertex disponivel. Verifique VERTEX_MODEL_NAME e acesso no projeto. Erro: {last_error}"
                ) from last_error
            raise RuntimeError("Falha ao gerar resposta no Vertex AI.")

        return self._extract_response_payload(response)

    async def warmup(self) -> None:
        """Preaquece conexao/modelo para reduzir latencia da primeira requisicao real."""
        try:
            await self.generate_response(
                message="Responda apenas: ok",
                tools=[],
            )
        except Exception:
            # Warmup e melhor esforco; falhas nao devem bloquear inicializacao.
            return