import time
import uuid
import json
from google import genai
from google.genai import types

from xulcan.core.primitives import ExternalID
from xulcan.llm.base import BaseLLMAdapter
from xulcan.protocol.message import UnifiedMessage, SystemMessage, UserMessage, AssistantMessage, ToolMessage
from xulcan.protocol.io import UnifiedResponse, FinishReason
from xulcan.protocol.tools import ToolCall, ToolDefinition
from xulcan.core.economics import UsageStats

class GeminiAdapter(BaseLLMAdapter):
    """Adapter for Google's Gemini models using the new genai SDK."""
    
    def __init__(self, api_key: str, **kwargs):
        super().__init__(**kwargs)
        self.client = genai.Client(api_key=api_key)

    async def _generate_impl(
        self,
        model: ExternalID, # Tipado estricto
        messages: list[UnifiedMessage],
        tools: list[ToolDefinition] | None = None,
        temperature: float = 0.0,
        max_tokens: int | None = None,
        **kwargs
    ) -> UnifiedResponse:
        
        start_time = time.time()
        
        # Variables para la API de Gemini
        gemini_messages = []
        system_instruction = None

        # ════════════════════════════════════════════════════════════════════════
        # 1. TRADUCCIÓN DE ENTRADA (Xulcan -> Gemini)
        # ════════════════════════════════════════════════════════════════════════
        for msg in messages:
            # A) SYSTEM MESSAGE (Va separado en Gemini)
            if isinstance(msg, SystemMessage):
                system_instruction = msg.content
                
            # B) USER MESSAGE
            elif isinstance(msg, UserMessage):
                # Asumimos texto plano por ahora (multimodalidad iría aquí)
                gemini_messages.append(types.Content(
                    role="user", 
                    parts=[types.Part.from_text(text=str(msg.content))]
                ))
            
            # C) ASSISTANT MESSAGE
            elif isinstance(msg, AssistantMessage):
                if msg.tool_calls:
                    parts = [
                        types.Part.from_function_call(name=tc.name, args=tc.arguments) 
                        for tc in msg.tool_calls
                    ]
                    gemini_messages.append(types.Content(role="model", parts=parts))
                elif msg.content:
                    gemini_messages.append(types.Content(
                        role="model", 
                        parts=[types.Part.from_text(text=msg.content)]
                    ))
            
            elif isinstance(msg, ToolMessage):
                try:
                    parsed_content = json.loads(msg.content) if isinstance(msg.content, str) else msg.content
                    # Si el parseo fue exitoso pero el resultado es un string/lista/int en lugar de dict
                    if not isinstance(parsed_content, dict):
                        resp_dict = {"result": parsed_content}
                    else:
                        resp_dict = parsed_content
                except (json.JSONDecodeError, TypeError):
                    # Si el contenido era texto plano o basura in-parseable
                    resp_dict = {"result": str(msg.content)}
                    
                gemini_messages.append(types.Content(
                    role="user", 
                    parts=[types.Part.from_function_response(
                        name=msg.name, 
                        response=resp_dict 
                    )]
                ))

        # ════════════════════════════════════════════════════════════════════════
        # 2. DEFINICIÓN DINÁMICA DE HERRAMIENTAS
        # ════════════════════════════════════════════════════════════════════════
        gemini_tools = None
        
        if tools:
            declarations = []
            for tool_def in tools:
                func = tool_def.function
                declarations.append(types.FunctionDeclaration(
                    name=func.name,
                    description=func.description or "",
                    # El SDK de Python de genai acepta diccionarios JSON Schema directamente
                    parameters=func.parameters 
                ))
            
            gemini_tools = [types.Tool(function_declarations=declarations)]

        # ════════════════════════════════════════════════════════════════════════
        # 3. LLAMADA A LA NUBE
        # ════════════════════════════════════════════════════════════════════════
        config = types.GenerateContentConfig(
            temperature=temperature,
            tools=gemini_tools,
            system_instruction=system_instruction
        )
        if max_tokens:
            config.max_output_tokens = max_tokens

        response = await self.client.aio.models.generate_content(
            model=model,
            contents=gemini_messages,
            config=config
        )

        # ════════════════════════════════════════════════════════════════════════
        # 4. TRADUCCIÓN DE SALIDA (Gemini -> Xulcan)
        # ════════════════════════════════════════════════════════════════════════
        tool_calls_out =[]
        text_content = ""
        finish_reason = FinishReason.STOP # Por defecto

        if response.candidates:
            candidate = response.candidates[0]
            
            # 🛡️ DEFENSA: Verificamos si Google censuró o vació el contenido
            if candidate.content and candidate.content.parts:
                for part in candidate.content.parts:
                    if part.function_call:
                        tool_calls_out.append(ToolCall(
                            id=f"call_{uuid.uuid4().hex[:8]}",
                            name=part.function_call.name,
                            arguments=dict(part.function_call.args)
                        ))
                    elif part.text:
                        text_content += part.text
            else:
                # Si 'content' es None, atrapamos el motivo real (Safety, Recitation, etc.)
                motivo = getattr(candidate, "finish_reason", "UNKNOWN")
                text_content = f"[XULCAN KERNEL WARNING] La API de Google bloqueó o devolvió una respuesta vacía. Motivo (finish_reason): {motivo}"
                finish_reason = FinishReason.CONTENT_FILTER

        if tool_calls_out:
            finish_reason = FinishReason.TOOL_CALLS

        in_tokens = getattr(response.usage_metadata, 'prompt_token_count', 0) or 0
        out_tokens = getattr(response.usage_metadata, 'candidates_token_count', 0) or 0
        latency = int((time.time() - start_time) * 1000)

        return UnifiedResponse(
            content=text_content if text_content else None,
            tool_calls=tool_calls_out if tool_calls_out else None,
            finish_reason=finish_reason,
            usage=UsageStats(
                input_tokens=in_tokens,
                output_tokens=out_tokens,
                total_tokens=in_tokens + out_tokens,
                latency_ms=latency
            )
        )