"""LLM client — talks to an OpenAI-compatible endpoint.

Works with either serving engine (both expose an OpenAI-compatible /v1 API):

  vLLM (int4 w4a16, best throughput/watt on the GPU):
    python -m vllm.entrypoints.openai.api_server \
        --model outputs/model-int4 --max-model-len 4096 --port 8001

  Ollama (GGUF Q4_0; simplest, great for the E-series):
    ollama run gemma4:e4b-it-qat        # or your fine-tuned GGUF via a Modelfile
    # Ollama serves OpenAI-compat at http://localhost:11434/v1

Point serving.base_url / serving.served_model in config.yaml at whichever you run.
Generation is capped (max_tokens, stop sequences) to avoid wasted compute.
"""
from __future__ import annotations

from dataclasses import dataclass

from .settings import Config


@dataclass
class LLMResult:
    text: str
    prompt_tokens: int = 0
    completion_tokens: int = 0

    @property
    def total_tokens(self) -> int:
        return self.prompt_tokens + self.completion_tokens


def _resolve_secret(value: str) -> str:
    """Resolve an 'env:VAR_NAME' reference to the environment variable's value.

    Keeps real API keys out of config files: config holds 'env:OPENROUTER_API_KEY', the actual
    key lives only in the environment. Plain (non 'env:') values pass through unchanged.
    """
    import os

    if isinstance(value, str) and value.startswith("env:"):
        name = value[4:]
        secret = os.environ.get(name)
        if not secret:
            raise SystemExit(
                f"serving.api_key is '{value}' but ${name} is not set. "
                f"Set it first, e.g.  setx {name} \"<your-key>\"  (then open a new shell)."
            )
        return secret
    return value


class LLMClient:
    def __init__(self, cfg: Config):
        from openai import OpenAI

        self.model = cfg.get("serving.served_model") or cfg.get("model.base_id")
        self.temperature = float(cfg.get("generation.temperature", 0.2))
        self.top_p = float(cfg.get("generation.top_p", 0.9))
        self.max_tokens = int(cfg.get("generation.max_tokens", 512))
        self.stop = cfg.get("generation.stop", None)
        # Optional extra headers (OpenRouter uses HTTP-Referer / X-Title for attribution).
        headers = cfg.get("serving.headers") or None
        self._client = OpenAI(
            base_url=cfg.get("serving.base_url", "http://localhost:8001/v1"),
            api_key=_resolve_secret(cfg.get("serving.api_key", "not-needed")),
            default_headers=headers,
        )

    def generate(self, messages: list[dict]) -> LLMResult:
        resp = self._client.chat.completions.create(
            model=self.model,
            messages=messages,
            temperature=self.temperature,
            top_p=self.top_p,
            max_tokens=self.max_tokens,
            stop=self.stop,
        )
        usage = getattr(resp, "usage", None)
        return LLMResult(
            text=(resp.choices[0].message.content or "").strip(),
            prompt_tokens=getattr(usage, "prompt_tokens", 0) or 0,
            completion_tokens=getattr(usage, "completion_tokens", 0) or 0,
        )