src/llm/gemma_client.py

from __future__ import annotations

import os
from dataclasses import dataclass
from pathlib import Path

import spaces
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

from src.ppt.deck_generator import parse_slides_payload


SYSTEM_PROMPT = (
    "You are a presentation planner. Return ONLY a strict JSON array, no markdown and no commentary. "
    "Each array element must be an object with keys: "
    "'title' (string), 'bullets' (array of strings), and "
    "'layout_type' (one of: 'title_slide', 'content_slide', 'split_slide'). "
    "Use 5-8 slides, concise business-friendly bullets, and valid JSON syntax."
)


@dataclass
class GemmaClientConfig:
    model_id: str = "unsloth/gemma-4-12b-it-GGUF"
    transformers_model_id: str = "google/gemma-4-12b-it"
    temperature: float = 0.4
    max_new_tokens: int = 900
    backend: str = "transformers"  # transformers | llama.cpp | vllm


class GemmaClient:
    def __init__(self, config: GemmaClientConfig | None = None):
        self.config = config or GemmaClientConfig()
        self._tokenizer = None
        self._model = None

    def _load_transformers_model(self):
        if self._model is not None and self._tokenizer is not None:
            return

        # Patch transformers bug: Gemma-4 tokenizer sends extra_special_tokens as
        # a list but transformers calls .keys() on it expecting a dict.
        try:
            from transformers import tokenization_utils_base as _tub
            _orig_set = _tub.PreTrainedTokenizerBase._set_model_specific_special_tokens
            def _safe_set(self_tok, special_tokens):
                if isinstance(special_tokens, dict):
                    _orig_set(self_tok, special_tokens)
            _tub.PreTrainedTokenizerBase._set_model_specific_special_tokens = _safe_set
        except Exception:
            pass

        model_id = self.config.transformers_model_id
        hf_token = os.getenv("HF_TOKEN")
        self._tokenizer = AutoTokenizer.from_pretrained(
            model_id, trust_remote_code=True, token=hf_token
        )

        kwargs = {
            "device_map": "auto",
            "torch_dtype": torch.float16 if torch.cuda.is_available() else torch.float32,
            "trust_remote_code": True,
            "token": hf_token,
        }

        try:
            from transformers import BitsAndBytesConfig

            kwargs["quantization_config"] = BitsAndBytesConfig(
                load_in_4bit=True,
                bnb_4bit_quant_type="nf4",
                bnb_4bit_use_double_quant=True,
                bnb_4bit_compute_dtype=torch.float16,
            )
        except Exception:
            # On unsupported environments (e.g., CPU/macOS without bitsandbytes),
            # fallback to non-quantized loading.
            pass

        self._model = AutoModelForCausalLM.from_pretrained(model_id, **kwargs)

    def _generate_transformers(self, user_prompt: str) -> str:
        self._load_transformers_model()

        messages = [
            {"role": "system", "content": SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ]

        raw = self._tokenizer.apply_chat_template(
            messages,
            tokenize=True,
            add_generation_prompt=True,
            return_tensors="pt",
        )
        # Newer transformers may return a BatchEncoding; older returns a plain tensor
        if hasattr(raw, "input_ids"):
            input_ids = raw.input_ids.to(self._model.device)
            gen_kwargs = {"input_ids": input_ids}
            if hasattr(raw, "attention_mask"):
                gen_kwargs["attention_mask"] = raw.attention_mask.to(self._model.device)
        else:
            input_ids = raw.to(self._model.device)
            gen_kwargs = {"input_ids": input_ids}

        with torch.no_grad():
            outputs = self._model.generate(
                **gen_kwargs,
                max_new_tokens=self.config.max_new_tokens,
                temperature=self.config.temperature,
                do_sample=True,
                pad_token_id=self._tokenizer.eos_token_id,
            )

        prompt_len = input_ids.shape[1]
        generated_ids = outputs[0][prompt_len:]
        return self._tokenizer.decode(generated_ids, skip_special_tokens=True).strip()

    def _generate_llama_cpp(self, user_prompt: str) -> str:
        try:
            from llama_cpp import Llama
        except Exception as exc:
            raise RuntimeError(
                f"llama.cpp backend import failed: {exc}"
            ) from exc

        model_path = self._resolve_llama_cpp_model_path()

        llm = Llama(
            model_path=model_path,
            n_ctx=4096,
            n_gpu_layers=-1,
            verbose=False,
        )
        completion = llm.create_chat_completion(
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=self.config.temperature,
            max_tokens=self.config.max_new_tokens,
        )
        return completion["choices"][0]["message"]["content"].strip()

    def _resolve_llama_cpp_model_path(self) -> str:
        env_model_path = os.getenv("LLAMA_CPP_MODEL_PATH")
        if env_model_path:
            return env_model_path

        preferred = [
            "Q4_K_M.gguf",
            "q4_k_m.gguf",
            "Q4_K_S.gguf",
            "q4_k_s.gguf",
            "Q5_K_M.gguf",
            "q5_k_m.gguf",
        ]

        try:
            from huggingface_hub import HfApi, hf_hub_download
        except Exception as exc:
            raise RuntimeError(
                "llama.cpp backend requires either LLAMA_CPP_MODEL_PATH or huggingface_hub to auto-download GGUF."
            ) from exc

        repo_id = os.getenv("GGUF_MODEL_REPO", self.config.model_id)
        files = HfApi(token=os.getenv("HF_TOKEN")).list_repo_files(repo_id=repo_id)
        gguf_files = [f for f in files if f.lower().endswith(".gguf")]
        if not gguf_files:
            raise RuntimeError(f"No GGUF files found in repo: {repo_id}")

        selected = None
        for suffix in preferred:
            selected = next((f for f in gguf_files if f.endswith(suffix)), None)
            if selected:
                break
        if selected is None:
            selected = gguf_files[0]

        local_dir = Path(os.getenv("GGUF_CACHE_DIR", "models"))
        local_dir.mkdir(parents=True, exist_ok=True)
        downloaded = hf_hub_download(
            repo_id=repo_id,
            filename=selected,
            local_dir=str(local_dir),
            token=os.getenv("HF_TOKEN"),
        )
        os.environ["LLAMA_CPP_MODEL_PATH"] = downloaded
        return downloaded

    def _generate_vllm(self, user_prompt: str) -> str:
        try:
            from openai import OpenAI
        except Exception as exc:
            raise RuntimeError("vLLM backend requested but openai package is unavailable.") from exc

        base_url = os.getenv("VLLM_BASE_URL", "http://localhost:8000/v1")
        model_name = os.getenv("VLLM_MODEL_NAME", self.config.model_id)
        client = OpenAI(base_url=base_url, api_key=os.getenv("VLLM_API_KEY", "EMPTY"))

        response = client.chat.completions.create(
            model=model_name,
            messages=[
                {"role": "system", "content": SYSTEM_PROMPT},
                {"role": "user", "content": user_prompt},
            ],
            temperature=self.config.temperature,
            max_tokens=self.config.max_new_tokens,
        )
        return response.choices[0].message.content.strip()

    @spaces.GPU
    def generate_json(self, user_prompt: str) -> str:
        backend = self.config.backend.lower().strip()
        if backend == "llama.cpp":
            return self._generate_llama_cpp(user_prompt)
        if backend == "vllm":
            return self._generate_vllm(user_prompt)
        return self._generate_transformers(user_prompt)

    def generate_slides(self, user_prompt: str) -> list[dict]:
        raw_output = self.generate_json(user_prompt)
        return parse_slides_payload(raw_output)


_DEFAULT_CLIENT: GemmaClient | None = None


def load_model(backend: str | None = None) -> GemmaClient:
    global _DEFAULT_CLIENT
    if _DEFAULT_CLIENT is None or (backend and _DEFAULT_CLIENT.config.backend != backend):
        cfg = GemmaClientConfig(backend=backend or os.getenv("MODEL_BACKEND", "transformers"))
        _DEFAULT_CLIENT = GemmaClient(cfg)
    return _DEFAULT_CLIENT


def should_startup_warmup() -> bool:
    policy = os.getenv("ENABLE_STARTUP_WARMUP", "auto").strip().lower()
    if policy in {"1", "true", "yes", "on"}:
        return True
    if policy in {"0", "false", "no", "off"}:
        return False

    # Auto mode: enable warmup on hosted Space environments.
    return bool(
        os.getenv("SPACE_ID")
        or os.getenv("HF_SPACE_ID")
        or os.getenv("SPACE_AUTHOR_NAME")
    )


def warmup_model_cache(backend: str | None = None) -> str:
    client = load_model(backend=backend)
    chosen_backend = (backend or client.config.backend).lower().strip()

    if chosen_backend == "llama.cpp":
        model_path = client._resolve_llama_cpp_model_path()
        return f"llama.cpp cache ready: {model_path}"

    if chosen_backend == "transformers":
        client._load_transformers_model()
        return f"transformers model ready: {client.config.transformers_model_id}"

    if chosen_backend == "vllm":
        return "vllm warmup skipped: remote endpoint assumed"

    return f"warmup skipped: unsupported backend '{chosen_backend}'"


def generate_slides(user_input: str, backend: str | None = None):
    client = load_model(backend=backend)
    return client.generate_slides(user_input)


def generate_slides_json(user_input: str, backend: str | None = None) -> str:
    client = load_model(backend=backend)
    return client.generate_json(user_input)


def safe_parse_generated_json(raw_output: str) -> list[dict]:
    return parse_slides_payload(raw_output)