"""Pluggable LLM providers for Chalchitra.

The whole point: develop against whatever model is "in front of us" (LM Studio
running Gemma locally, today) and swap to the Hugging Face / ZeroGPU model later
by changing config only — never the frontend, never the prompt, never the JSON
contract.

Both LM Studio and HF's TGI/vLLM serving expose the OpenAI chat-completions API,
so a single OpenAI-compatible provider covers both "in front of you" cases. The
hf_local provider is a stub for the eventual in-process transformers path on
ZeroGPU (with @spaces.GPU) if we choose to run the model inside the Space rather
than behind an OpenAI-compatible server.

There is no hosted/cloud API in the loop — the model answering is always the one
we are pointing at.
"""

from __future__ import annotations

import os
from typing import Any, Protocol


class Provider(Protocol):
    """Minimal contract every provider satisfies.

    Given a system prompt and OpenAI-format content blocks (text + image_url),
    return the model's raw text response. Parsing/validation is the oracle's job.
    """

    name: str
    model: str

    def complete(
        self,
        system: str,
        content: list[dict[str, Any]],
        max_tokens: int,
        temperature: float | None = None,
    ) -> str: ...


class OpenAICompatibleProvider:
    """Talks to any OpenAI chat-completions endpoint.

    Covers LM Studio (local Gemma, dev) and HF TGI/vLLM (deploy) alike.
    """

    def __init__(
        self,
        base_url: str,
        model: str,
        api_key: str = "not-needed",
        temperature: float = 0.8,
        timeout: float = 120.0,
    ) -> None:
        # Imported lazily so the module loads even before deps are installed.
        from openai import OpenAI

        self.name = "openai_compatible"
        self.base_url = base_url
        self.model = model
        self.temperature = temperature
        # A bounded timeout so a down/slow model fails fast instead of hanging
        # the request forever. max_retries=0 — we own retries in the oracle.
        self._client = OpenAI(
            base_url=base_url, api_key=api_key, timeout=timeout, max_retries=0
        )

    def complete(
        self,
        system: str,
        content: list[dict[str, Any]],
        max_tokens: int,
        temperature: float | None = None,
    ) -> str:
        resp = self._client.chat.completions.create(
            model=self.model,
            max_tokens=max_tokens,
            temperature=self.temperature if temperature is None else temperature,
            messages=[
                {"role": "system", "content": system},
                {"role": "user", "content": content},
            ],
        )
        return resp.choices[0].message.content or ""


# ── ZeroGPU plumbing for the in-process model path ────────────────────────────
# `_gpu` applies @spaces.GPU only when running on a Space (where `spaces` is
# installed); locally it's a no-op so the same code imports without a GPU. The
# decorated `_generate` is defined at module level so it's registered for
# ZeroGPU's startup scan.
try:
    import spaces as _spaces

    def _gpu(fn):
        return _spaces.GPU(duration=180)(fn)
except ImportError:
    def _gpu(fn):
        return fn

GPU_MAX_NEW_TOKENS = 2048


def _to_qwen_messages(system: str, content: list[dict[str, Any]]):
    """Translate our OpenAI-format content into Qwen2.5-VL chat messages.

    image_url blocks (data URLs) become {"type":"image","image": <data url>},
    which qwen-vl-utils decodes; text blocks pass through.
    """
    user_content: list[dict[str, Any]] = []
    for block in content:
        if block.get("type") == "image_url":
            user_content.append({"type": "image", "image": block["image_url"]["url"]})
        elif block.get("type") == "text":
            user_content.append({"type": "text", "text": block["text"]})
    return [
        {"role": "system", "content": system},
        {"role": "user", "content": user_content},
    ]


@_gpu
def _generate(model, processor, messages, max_new_tokens: int, temperature: float) -> str:
    """The GPU-bound forward pass. Runs inside a ZeroGPU allocation on the Space."""
    import torch
    from qwen_vl_utils import process_vision_info

    text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
    image_inputs, video_inputs = process_vision_info(messages)
    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt",
    )
    model.to("cuda")
    inputs = inputs.to("cuda")
    with torch.no_grad():
        generated = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=temperature > 0,
            temperature=temperature if temperature > 0 else None,
            top_p=0.9,
        )
    trimmed = [out[len(inp):] for inp, out in zip(inputs.input_ids, generated)]
    decoded = processor.batch_decode(
        trimmed, skip_special_tokens=True, clean_up_tokenization_spaces=False
    )
    return decoded[0] if decoded else ""


class HFLocalProvider:
    """Run Qwen2.5-VL in-process on ZeroGPU via transformers (the HF deploy path).

    The Space loads the model itself and runs it on a ZeroGPU-allocated GPU. The
    seam matches OpenAICompatibleProvider exactly, so app code never branches on
    which path we took — only the env config changes. The model + processor load
    lazily to CPU on first request; the forward pass moves to cuda inside the
    `_gpu`-wrapped `_generate` (ZeroGPU only grants the GPU there).
    """

    def __init__(self, model: str, temperature: float = 0.8, **_: Any) -> None:
        self.name = "hf_local"
        self.model = model
        self.temperature = temperature
        self._model = None
        self._processor = None

    def _ensure_loaded(self) -> None:
        if self._model is not None:
            return
        import torch
        from transformers import AutoProcessor, Qwen2_5_VLForConditionalGeneration

        self._processor = AutoProcessor.from_pretrained(self.model)
        self._model = Qwen2_5_VLForConditionalGeneration.from_pretrained(
            self.model, torch_dtype=torch.bfloat16
        )

    def complete(
        self,
        system: str,
        content: list[dict[str, Any]],
        max_tokens: int,
        temperature: float | None = None,
    ) -> str:
        self._ensure_loaded()
        messages = _to_qwen_messages(system, content)
        temp = self.temperature if temperature is None else temperature
        return _generate(self._model, self._processor, messages, max_tokens, temp)


def get_provider() -> Provider:
    """Build the configured provider from environment variables.

    CHALCHITRA_PROVIDER   openai_compatible (default) | hf_local
    CHALCHITRA_BASE_URL   default http://localhost:1234/v1  (LM Studio)
    CHALCHITRA_MODEL      default google/gemma-4-e4b
    CHALCHITRA_API_KEY    default "lm-studio" (LM Studio ignores it)
    CHALCHITRA_TEMPERATURE default 0.8
    """
    kind = os.environ.get("CHALCHITRA_PROVIDER", "openai_compatible").strip()
    model = os.environ.get("CHALCHITRA_MODEL", "google/gemma-4-e4b").strip()

    if kind == "openai_compatible":
        return OpenAICompatibleProvider(
            base_url=os.environ.get("CHALCHITRA_BASE_URL", "http://localhost:1234/v1").strip(),
            model=model,
            api_key=os.environ.get("CHALCHITRA_API_KEY", "lm-studio").strip(),
            temperature=float(os.environ.get("CHALCHITRA_TEMPERATURE", "0.8")),
            timeout=float(os.environ.get("CHALCHITRA_TIMEOUT", "120")),
        )
    if kind == "hf_local":
        return HFLocalProvider(
            model=model,
            temperature=float(os.environ.get("CHALCHITRA_TEMPERATURE", "0.8")),
        )

    raise ValueError(
        f"Unknown CHALCHITRA_PROVIDER={kind!r}. Use 'openai_compatible' or 'hf_local'."
    )