"""Modal endpoint for the fine-tuned MiniCPM4.1-8B recipe planner.

Runs in its OWN container because MiniCPM4.1's custom code requires
transformers 4.x (CacheLayerMixin + is_torch_fx_available), which conflicts
with the MiniCPM-V-4.6 vision model in the main app (needs transformers 5.x).

Deploy:
    modal deploy modal_app/planner_endpoint.py

The Gradio app calls it via modal.Cls.from_name("cook-with-me-planner",
"Planner").infer.remote(prompt, ...).
"""
from __future__ import annotations

import os

import modal

app = modal.App("cook-with-me-planner")

# 8B bf16 weights cached on a volume so cold starts don't re-download ~16GB.
hf_cache = modal.Volume.from_name("cook-with-me-planner-cache", create_if_missing=True)
hf_secret = modal.Secret.from_name("huggingface-secret")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "torch==2.4.0",
        # MiniCPM4.1 custom code needs BOTH CacheLayerMixin (>=4.54) and
        # is_torch_fx_available (removed in 5.0) — only 4.54..4.x has both.
        "transformers>=4.54,<5.0",
        "huggingface_hub>=0.26,<1.0",
        "accelerate",
        "sentencepiece",
        "safetensors",
    )
    .env({"HF_HOME": "/cache/hf"})
)

# Fine-tuned weights; tokenizer pulled from base (FT tokenizer_config was saved
# by transformers 5.x and is not readable by 4.x).
PLANNER_REPO = os.environ.get("COOK_WITH_ME_PLANNER_FT_REPO", "eldinosaur/cook-with-me-planner-8b")
BASE_REPO = "openbmb/MiniCPM4.1-8B"


@app.cls(
    image=image,
    gpu="L4",
    volumes={"/cache": hf_cache},
    secrets=[hf_secret],
    scaledown_window=240,
    timeout=600,
)
class Planner:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        print(f"Loading planner weights from {PLANNER_REPO}...")
        self.tokenizer = AutoTokenizer.from_pretrained(BASE_REPO, trust_remote_code=True)
        if self.tokenizer.pad_token is None:
            self.tokenizer.pad_token = self.tokenizer.eos_token
        self.model = AutoModelForCausalLM.from_pretrained(
            PLANNER_REPO,
            torch_dtype=torch.bfloat16,
            trust_remote_code=True,
            device_map="cuda",
        ).eval()
        print("Planner ready.")

    @modal.method()
    def infer(self, prompt: str, max_new_tokens: int = 1024, temperature: float = 0.0) -> str:
        import torch

        messages = [{"role": "user", "content": prompt}]
        # enable_thinking=False -> direct JSON, no <think> reasoning preamble
        try:
            enc = self.tokenizer.apply_chat_template(
                messages,
                add_generation_prompt=True,
                tokenize=True,
                return_tensors="pt",
                return_dict=True,
                enable_thinking=False,
            )
        except TypeError:
            enc = self.tokenizer.apply_chat_template(
                messages, add_generation_prompt=True, tokenize=True,
                return_tensors="pt", return_dict=True,
            )

        input_ids = enc["input_ids"].to(self.model.device)
        input_len = input_ids.shape[1]
        gen_inputs = {"input_ids": input_ids}
        if enc.get("attention_mask") is not None:
            gen_inputs["attention_mask"] = enc["attention_mask"].to(self.model.device)

        gen_kwargs = dict(max_new_tokens=max_new_tokens, repetition_penalty=1.05)
        if temperature and temperature > 0:
            gen_kwargs.update(do_sample=True, temperature=temperature, top_p=0.9)
        else:
            gen_kwargs.update(do_sample=False)

        with torch.no_grad():
            out = self.model.generate(**gen_inputs, **gen_kwargs)
        return self.tokenizer.decode(out[0][input_len:], skip_special_tokens=True)


@app.local_entrypoint()
def test():
    prompt = (
        "You are a creative chef. Available ingredients: tomato, onion, garlic, pasta, olive oil.\n"
        'Respond ONLY with JSON: {"options": [{"name": "...", "why": "..."}, {"name": "...", "why": "..."}, {"name": "...", "why": "..."}]}'
    )
    out = Planner().infer.remote(prompt, max_new_tokens=400)
    print("OUTPUT:\n", out)