import time
import uuid
import json
from ollama import AsyncClient, ResponseError

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

# Importamos a nuestros "Secuaces" locales
from xulcan.llm.adapters.common.strategies import NativeToolStrategy, PromptInjectionStrategy
from xulcan.llm.adapters.common.extractors import OutputExtractor

class OllamaAdapter(BaseLLMAdapter):
    """Clean, Strategy-based Adapter for local models on Ollama."""
    
    def __init__(self, host: str = "http://localhost:11434", **kwargs):
        super().__init__(**kwargs)
        self.client = AsyncClient(host=host)
        self._fallback_cache = set()

    def _translate_messages(self, messages: list[UnifiedMessage]) -> list[dict]:
        """Traduce de Xulcan UnifiedMessage a Ollama Dict format."""
        ollama_messages =[]
        for msg in messages:
            if isinstance(msg, SystemMessage):
                ollama_messages.append({"role": "system", "content": str(msg.content)})
            elif isinstance(msg, UserMessage):
                ollama_messages.append({"role": "user", "content": str(msg.content)})
            elif isinstance(msg, AssistantMessage):
                msg_dict = {"role": "assistant", "content": msg.content or ""}
                if msg.tool_calls:
                    msg_dict["tool_calls"] =[{"function": {"name": tc.name, "arguments": tc.arguments}} for tc in msg.tool_calls]
                ollama_messages.append(msg_dict)
            elif isinstance(msg, ToolMessage):
                ollama_messages.append({
                    "role": "tool", 
                    "content": json.dumps(msg.content) if isinstance(msg.content, dict) else str(msg.content), 
                    "name": msg.name
                })
        return ollama_messages

    async def _generate_impl(
        self, model: ExternalID, messages: list[UnifiedMessage], 
        tools: list[ToolDefinition] | None = None, temperature: float = 0.0, 
        max_tokens: int | None = None, **kwargs
    ) -> UnifiedResponse:
        
        start_time = time.time()
        ollama_messages = self._translate_messages(messages)
        options = {"temperature": temperature}
        if max_tokens: options["num_predict"] = max_tokens

        # 1. ELEGIR ESTRATEGIA (Native vs Inyectada)
        requires_fallback = model in self._fallback_cache
        strategy = PromptInjectionStrategy() if requires_fallback else NativeToolStrategy()
        final_msgs, api_tools = strategy.apply(ollama_messages, tools)

        # 2. LLAMADA A OLLAMA
        try:
            response = await self.client.chat(model=model, messages=final_msgs, tools=api_tools, options=options)
        except ResponseError as e:
            if not requires_fallback and ("does not support tools" in str(e).lower() or e.status_code == 400):
                print(f"\n[Adaptador Ollama] Modelo '{model}' no soporta herramientas nativas. Guardando en caché y activando Fallback...")
                self._fallback_cache.add(model)
                strategy = PromptInjectionStrategy()
                final_msgs, api_tools = strategy.apply(ollama_messages, tools)
                response = await self.client.chat(model=model, messages=final_msgs, options=options)
            else:
                raise e

        # 3. EXTRACCIÓN Y LIMPIEZA
        raw_msg = response.get("message", {})
        content = raw_msg.get("content", "")
        
        content, reasoning = OutputExtractor.strip_think_tags(content)
        
        tool_calls_out =[]
        if "tool_calls" in raw_msg and raw_msg["tool_calls"]:
            for tc in raw_msg["tool_calls"]:
                func = tc.get("function", {})
                tool_calls_out.append(ToolCall(
                    id=f"call_{uuid.uuid4().hex[:8]}",
                    name=func.get("name", ""), arguments=func.get("arguments", {})
                ))
        elif tools:
            # Si no hubo nativas, usamos al Caza-Rebeldes
            rebel_tools, content = OutputExtractor.extract_rebel_json(content)
            tool_calls_out.extend(rebel_tools)

        in_tokens = response.get("prompt_eval_count", 0) or 0
        out_tokens = response.get("eval_count", 0) or 0

        # 4. DEVOLVER EL ESTÁNDAR XULCAN
        return UnifiedResponse(
            content=content if content else None,
            reasoning_content=reasoning,
            tool_calls=tool_calls_out if tool_calls_out else None,
            finish_reason=FinishReason.TOOL_CALLS if tool_calls_out else FinishReason.STOP,
            usage=UsageStats(
                input_tokens=in_tokens, output_tokens=out_tokens, 
                total_tokens=in_tokens + out_tokens,
                latency_ms=int((time.time() - start_time) * 1000)
            )
        )