dcc-rpg/src/engine/providers.py

from __future__ import annotations

import json
import os
from abc import ABC, abstractmethod
from typing import Any

import httpx

from src.engine.config import ProviderConfig, load_from_env


class LLMProvider(ABC):
    @abstractmethod
    def ask(self, prompt: str, system: str = "") -> str:
        ...


class OllamaProvider(LLMProvider):
    def __init__(self, base_url: str, model: str, timeout: int = 60):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.timeout = timeout

    def ask(self, prompt: str, system: str = "") -> str:
        body = {"model": self.model, "prompt": prompt, "stream": False}
        if system:
            body["system"] = system
        resp = httpx.post(
            f"{self.base_url}/api/generate",
            json=body,
            timeout=self.timeout,
        )
        resp.raise_for_status()
        return resp.json().get("response", "")


class OpenRouterProvider(LLMProvider):
    _FREE_FALLBACK = [
        "openrouter/free",
        "google/gemma-4-31b-it:free",
        "qwen/qwen3-coder:free",
    ]

    def __init__(self, api_key: str, model: str, timeout: int = 60, free_fallback: bool = True):
        self.api_key = api_key
        self.model = model
        self.timeout = timeout
        self.free_fallback = free_fallback

    def _call(self, model: str, prompt: str, system: str) -> str:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json",
        }
        body = {"model": model, "messages": messages, "stream": False}

        resp = httpx.post(
            "https://openrouter.ai/api/v1/chat/completions",
            json=body,
            headers=headers,
            timeout=self.timeout,
        )
        resp.raise_for_status()
        return resp.json()["choices"][0]["message"]["content"]

    def ask(self, prompt: str, system: str = "") -> str:
        models = [self.model]
        if self.free_fallback:
            models.extend(self._FREE_FALLBACK)
        last_err = None
        for m in models:
            try:
                return self._call(m, prompt, system)
            except Exception as e:
                last_err = e
                continue
        raise last_err  # type: ignore


class LocalProvider(LLMProvider):
    def __init__(self, base_url: str, model: str, api_key: str = "",
                 timeout: int = 60, fallbacks: list[tuple[str, str]] | None = None):
        self.base_url = base_url.rstrip("/")
        self.model = model
        self.api_key = api_key
        self.timeout = timeout
        self.fallbacks = fallbacks or []

    def _call(self, base_url: str, model: str, prompt: str, system: str) -> str:
        messages = []
        if system:
            messages.append({"role": "system", "content": system})
        messages.append({"role": "user", "content": prompt})

        headers = {"Content-Type": "application/json"}
        if self.api_key:
            headers["Authorization"] = f"Bearer {self.api_key}"

        body = {"model": model, "messages": messages, "stream": False}

        resp = httpx.post(
            f"{base_url}/chat/completions",
            json=body,
            headers=headers,
            timeout=self.timeout,
        )
        resp.raise_for_status()
        data = resp.json()
        return data["choices"][0]["message"]["content"]

    def ask(self, prompt: str, system: str = "") -> str:
        candidates = [(self.base_url, self.model)] + self.fallbacks
        last_err = None
        for base_url, model in candidates:
            try:
                text = self._call(base_url, model, prompt, system)
                # Rileva risposte di censura del modello gratuito
                if "User Safety" in text and "safe" in text.lower():
                    raise ValueError(f"Modello '{model}' ha censurato la risposta: {text[:60]}")
                return text
            except Exception as e:
                last_err = e
                continue
        raise last_err  # type: ignore


_registry: dict[str, Any] = {
    "ollama": OllamaProvider,
    "openrouter": OpenRouterProvider,
    "local": LocalProvider,
}


def create_provider(cfg: ProviderConfig | None = None) -> LLMProvider:
    if cfg is None:
        cfg = load_from_env()

    cls = _registry.get(cfg.kind)
    if not cls:
        raise ValueError(f"Provider sconosciuto: {cfg.kind} (disponibili: {list(_registry.keys())})")

    if cfg.kind == "ollama":
        return cls(cfg.ollama_base_url, cfg.ollama_model, cfg.timeout)
    elif cfg.kind == "openrouter":
        key = cfg.openrouter_api_key or os.getenv("DCC_OPENROUTER_KEY", os.getenv("OPENROUTER_API_KEY", ""))
        if not key:
            raise ValueError("DCC_OPENROUTER_KEY non configurata")
        return cls(key, cfg.openrouter_model, cfg.timeout, cfg.openrouter_free_fallback)
    elif cfg.kind == "local":
        fallbacks = []
        if cfg.local_fallback_url and cfg.local_fallback_model:
            fallbacks.append((cfg.local_fallback_url, cfg.local_fallback_model))
        return cls(cfg.local_base_url, cfg.local_model, cfg.local_api_key, cfg.timeout, fallbacks=fallbacks)
    else:
        raise ValueError(f"Provider non implementato: {cfg.kind}")