"""Modal deployment for Pinch's reasoning model (Mellum 2).

Vision (MiniCPM-V) uses OpenBMB's free API and the image uses HF Inference, so
Modal only serves Mellum 2 — which has no public API and must run on a GPU.

Cost design:
  - min_containers=0 → $0 when idle (scales to zero).
  - Weights are cached in a Modal Volume, so a cold start mounts them (~30-60s)
    instead of re-downloading 24GB from HF (~2-3 min) — ~6x cheaper per cold start.
  - scaledown_window=120 → only ~2 min of warm idle after the last request.

Deploy:  modal deploy modal_app.py
Endpoint URL is unchanged, so the Space's MODAL_REASON_URL stays valid.
"""

import modal

REASON_MODEL = "JetBrains/Mellum2-12B-A2.5B-Instruct"
CACHE_DIR = "/cache"

# Persists HF weights across cold starts: downloaded once, mounted thereafter.
hf_cache = modal.Volume.from_name("pinch-hf-cache", create_if_missing=True)

image = (
    modal.Image.debian_slim(python_version="3.12")
    .uv_pip_install("torch", "transformers", "accelerate", "fastapi[standard]")
    .env({"HF_HOME": CACHE_DIR})  # transformers caches into the mounted Volume
)

app = modal.App("epicurean-simmer")


@app.cls(gpu="A100-40GB", image=image, volumes={CACHE_DIR: hf_cache},
         scaledown_window=120, min_containers=0, timeout=600)
class Reasoner:
    """JetBrains Mellum 2 — the flavour-planning brain (12B MoE, bf16 ≈ 24GB)."""

    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.tokenizer = AutoTokenizer.from_pretrained(REASON_MODEL)
        self.model = AutoModelForCausalLM.from_pretrained(
            REASON_MODEL, dtype=torch.bfloat16, device_map="cuda"
        )
        hf_cache.commit()  # persist freshly-downloaded weights for future cold starts

    @modal.fastapi_endpoint(method="POST")
    def generate(self, item: dict):
        """POST {"messages": [...], "max_tokens": int} -> {"text": str}."""
        enc = self.tokenizer.apply_chat_template(
            item["messages"], add_generation_prompt=True,
            return_tensors="pt", return_dict=True,
        ).to("cuda")
        input_len = enc["input_ids"].shape[-1]
        output = self.model.generate(
            **enc, max_new_tokens=item.get("max_tokens", 400),
            temperature=0.3, do_sample=True,
            pad_token_id=self.tokenizer.eos_token_id,
        )
        text = self.tokenizer.decode(output[0][input_len:], skip_special_tokens=True)
        return {"text": text}