# app/xulcan/llm/adapters/openai_protocol.py

import time
import json
from typing import List, Optional
from openai import AsyncOpenAI

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 OpenAICompatibleAdapter(BaseLLMAdapter):
    """Universal adapter for any LLM provider that implements the OpenAI API standard 
    (OpenAI, Groq, DeepSeek, xAI, OpenRouter, vLLM, etc).
    """
    
    def __init__(self, api_key: str, base_url: Optional[str] = None, **kwargs):
        super().__init__(**kwargs)
        # El cliente AsyncOpenAI maneja todo el protocolo de red subyacente
        self.client = AsyncOpenAI(api_key=api_key, base_url=base_url)

    async def _generate_impl(
        self,
        model: ExternalID, 
        messages: List[UnifiedMessage],
        tools: Optional[List[ToolDefinition]] = None,
        temperature: float = 0.0,
        max_tokens: Optional[int] = None,
        **kwargs
    ) -> UnifiedResponse:
        
        start_time = time.time()
        
        # ════════════════════════════════════════════════════════════════════════
        # 1. TRADUCCIÓN DE ENTRADA (Xulcan -> OpenAI Standard)
        # ════════════════════════════════════════════════════════════════════════
        oai_messages =[]
        
        for msg in messages:
            if isinstance(msg, SystemMessage):
                oai_messages.append({"role": "system", "content": str(msg.content)})
                
            elif isinstance(msg, UserMessage):
                oai_messages.append({"role": "user", "content": str(msg.content)})
                
            elif isinstance(msg, AssistantMessage):
                oai_msg = {"role": "assistant", "content": msg.content or ""}
                
                # Si el asistente había decidido usar herramientas en el pasado
                if msg.tool_calls:
                    oai_tool_calls =[]
                    for tc in msg.tool_calls:
                        oai_tool_calls.append({
                            "id": tc.id, # El ID es crucial en el estándar OpenAI
                            "type": "function",
                            "function": {
                                "name": tc.name,
                                # OpenAI exige que los argumentos sean un string JSON, no un dict
                                "arguments": json.dumps(tc.arguments) if isinstance(tc.arguments, dict) else tc.arguments
                            }
                        })
                    oai_msg["tool_calls"] = oai_tool_calls
                    
                oai_messages.append(oai_msg)
                
            elif isinstance(msg, ToolMessage):
                # En el protocolo OpenAI, el mensaje de la tool TIENE que tener el tool_call_id
                # al que está respondiendo, para que el LLM no pierda el hilo.
                oai_messages.append({
                    "role": "tool",
                    "tool_call_id": getattr(msg, 'tool_call_id', 'unknown'),
                    "content": str(msg.content)
                })

        # ════════════════════════════════════════════════════════════════════════
        # 2. DEFINICIÓN DINÁMICA DE HERRAMIENTAS
        # ════════════════════════════════════════════════════════════════════════
        oai_tools = None
        if tools:
            oai_tools =[]
            for tool_def in tools:
                func = tool_def.function
                oai_tools.append({
                    "type": "function",
                    "function": {
                        "name": func.name,
                        "description": func.description or "",
                        "parameters": func.parameters # JSON Schema Dict
                    }
                })

        # ════════════════════════════════════════════════════════════════════════
        # 3. LLAMADA A LA NUBE (O LOCAL)
        # ════════════════════════════════════════════════════════════════════════
        api_kwargs = {
            "model": model,
            "messages": oai_messages,
            "temperature": temperature,
        }
        if oai_tools:
            api_kwargs["tools"] = oai_tools
        if max_tokens:
            api_kwargs["max_tokens"] = max_tokens

        response = await self.client.chat.completions.create(**api_kwargs)

        # ════════════════════════════════════════════════════════════════════════
        # 4. TRADUCCIÓN DE SALIDA (OpenAI Standard -> Xulcan)
        # ════════════════════════════════════════════════════════════════════════
        choice = response.choices[0]
        msg_out = choice.message
        
        tool_calls_out =[]
        if msg_out.tool_calls:
            for tc in msg_out.tool_calls:
                try:
                    # OpenAI estándar devuelve los argumentos como string JSON
                    args_dict = json.loads(tc.function.arguments)
                except json.JSONDecodeError:
                    # 🛡️ DEFENSA: Si el modelo alucina un JSON inválido
                    args_dict = {"raw_arguments": tc.function.arguments}
                    
                tool_calls_out.append(ToolCall(
                    id=tc.id,
                    name=tc.function.name,
                    arguments=args_dict
                ))

        # Mapear el motivo de finalización
        finish_reason = FinishReason.STOP
        if choice.finish_reason == "tool_calls":
            finish_reason = FinishReason.TOOL_CALLS
        elif choice.finish_reason == "length":
            finish_reason = FinishReason.LENGTH
        elif choice.finish_reason == "content_filter":
            finish_reason = FinishReason.CONTENT_FILTER

        # Métricas
        usage = response.usage
        in_tokens = getattr(usage, 'prompt_tokens', 0)
        out_tokens = getattr(usage, 'completion_tokens', 0)
        latency = int((time.time() - start_time) * 1000)

        return UnifiedResponse(
            content=msg_out.content,
            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
            )
        )