# llm_router.py — enruta llamadas a Spaces remotos según config.yaml
from __future__ import annotations
from typing import Any, Dict, List, Optional
from pathlib import Path
import os
import yaml

from remote_clients import InstructClient, VisionClient, ToolsClient, ASRClient

def load_yaml(path: str) -> Dict[str, Any]:
    p = Path(path)
    if not p.exists():
        return {}
    return yaml.safe_load(p.read_text(encoding="utf-8")) or {}

class LLMRouter:
    def __init__(self, cfg: Dict[str, Any]):
        self.cfg = cfg
        self.rem = cfg.get("models", {}).get("routing", {}).get("use_remote_for", [])
        base_user = cfg.get("remote_spaces", {}).get("user", "veureu")
        eps = cfg.get("remote_spaces", {}).get("endpoints", {})
        token_enabled = cfg.get("security", {}).get("use_hf_token", False)
        hf_token = os.getenv(cfg.get("security", {}).get("hf_token_env", "HF_TOKEN")) if token_enabled else None

        def mk(endpoint_key: str, cls):
            info = eps.get(endpoint_key, {})
            base_url = info.get("base_url") or f"https://{base_user}-{info.get('space')}.hf.space"
            use_gradio = (info.get("client", "gradio") == "gradio")
            timeout = int(cfg.get("remote_spaces", {}).get("http", {}).get("timeout_seconds", 180))
            return cls(base_url=base_url, use_gradio=use_gradio, hf_token=hf_token, timeout=timeout)

        self.clients = {
            "salamandra-instruct": mk("salamandra-instruct", InstructClient),
            "salamandra-vision": mk("salamandra-vision", VisionClient),
            "salamandra-tools": mk("salamandra-tools", ToolsClient),
            "whisper-catalan": mk("whisper-catalan", ASRClient),
        }

    # ---- INSTRUCT ----
    def instruct(self, prompt: str, system: Optional[str] = None, model: str = "salamandra-instruct", **kwargs) -> str:
        if model in self.rem:
            return self.clients[model].generate(prompt, system=system, **kwargs)  # type: ignore
        raise RuntimeError(f"Modelo local no implementado para: {model}")

    # ---- VISION ----
    def vision_describe(self, image_paths: List[str], context: Optional[Dict[str, Any]] = None, model: str = "salamandra-vision", **kwargs) -> List[str]:
        if model in self.rem:
            return self.clients[model].describe(image_paths, context=context, **kwargs)  # type: ignore
        raise RuntimeError(f"Modelo local no implementado para: {model}")

    # ---- TOOLS ----
    def chat_with_tools(self, messages: List[Dict[str, str]], tools: Optional[List[Dict[str, Any]]] = None, model: str = "salamandra-tools", **kwargs) -> Dict[str, Any]:
        if model in self.rem:
            return self.clients[model].chat(messages, tools=tools, **kwargs)  # type: ignore
        raise RuntimeError(f"Modelo local no implementado para: {model}")

    # ---- ASR ----
    def asr_transcribe(self, audio_path: str, model: str = "whisper-catalan", **kwargs) -> Dict[str, Any]:
        if model in self.rem:
            return self.clients[model].transcribe(audio_path, **kwargs)  # type: ignore
        raise RuntimeError(f"Modelo local no implementado para: {model}")