"""Modal service: MiniCPM4.1-8B (the reasoning model) in its OWN environment.

The 8B's trust_remote_code is written for transformers ~4.56 and RuntimeErrors
on the transformers 5.12 that MiniCPM-V-4.6 forces on the Space. Modal isolates
it: this container pins transformers <5.0, loads the 8B on a warm GPU, and
exposes `Engine.generate` which the Space calls via the Modal SDK.

Deploy:   modal deploy modal_app.py
Test:     modal run modal_app.py::smoke
"""

import re

import modal

APP_NAME = "budgetbuddy-8b"
MODEL_ID = "openbmb/MiniCPM4.1-8B"

image = (
    modal.Image.debian_slim(python_version="3.11")
    .pip_install(
        "torch",
        "transformers>=4.56,<5.0",
        "accelerate>=0.30",
        "sentencepiece",
        "numpy<2",
    )
)

app = modal.App(APP_NAME)
hf_cache = modal.Volume.from_name("bb-hf-cache", create_if_missing=True)

with image.imports():
    import torch
    from transformers import AutoModelForCausalLM, AutoTokenizer

_THINK_RE = re.compile(r"<think>.*?</think>", re.DOTALL | re.IGNORECASE)


@app.cls(
    gpu="A10G",
    image=image,
    volumes={"/root/.cache/huggingface": hf_cache},
    scaledown_window=1200,         # stay warm 20 min after last call
    timeout=900,
    enable_memory_snapshot=True,   # snapshot the CPU-loaded model -> fast cold starts
)
class Engine:
    @modal.enter(snap=True)
    def load_cpu(self):
        # Runs during snapshot creation (CPU only). The 16GB weight load is
        # captured in the snapshot, so later cold starts restore instead of reload.
        self.tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, torch_dtype="bfloat16", trust_remote_code=True
        ).eval()

    @modal.enter(snap=False)
    def to_gpu(self):
        self.model = self.model.to("cuda")

    @modal.method()
    def generate(self, messages, max_new_tokens: int = 512,
                 enable_thinking: bool = False) -> str:
        try:
            inputs = self.tokenizer.apply_chat_template(
                messages, tokenize=True, add_generation_prompt=True,
                enable_thinking=enable_thinking, return_dict=True, return_tensors="pt",
            ).to("cuda")
        except TypeError:
            inputs = self.tokenizer.apply_chat_template(
                messages, tokenize=True, add_generation_prompt=True,
                return_dict=True, return_tensors="pt",
            ).to("cuda")
        with torch.no_grad():
            out = self.model.generate(
                **inputs, max_new_tokens=max_new_tokens, do_sample=False
            )
        text = self.tokenizer.decode(
            out[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True
        )
        return _THINK_RE.sub("", str(text)).strip()


@app.local_entrypoint()
def smoke():
    r = Engine().generate.remote(
        messages=[{"role": "user", "content": "Reply with exactly: ok"}],
        max_new_tokens=12,
    )
    print("SMOKE RESULT:", repr(r))