"""Modal GPU backend for Whisperkey guardians (MiniCPM4-8B and/or Nemotron-Mini-4B).

The HF Space stays a thin Gradio frontend; the heavy model runs here on a Modal L4 GPU
(~$0.80/hr). Set MODAL_MIN_CONTAINERS=1 to keep a warm replica for demos; default is 0 (scale-to-zero).

The web endpoint is a METHOD on the Guardian class, so each HTTP request hits a container that
already has the model loaded (via @modal.enter) - no cross-container hop.

Deploy (MiniCPM - default):
    pip install modal
    modal token new
    modal secret create jailbreak-dojo MODAL_API_KEY=<random>
    MODAL_MIN_CONTAINERS=1 modal deploy modal_app.py
    # copy URL -> MODAL_ENDPOINT on the HF Space

Deploy (Nemotron - NVIDIA prize path):
    MODAL_APP_NAME=jailbreak-dojo-nemotron \\
    GUARDIAN_HF_MODEL=nvidia/Nemotron-Mini-4B-Instruct \\
    MODAL_VOLUME=guardian-nemotron-cache \\
    MODAL_MIN_CONTAINERS=1 modal deploy modal_app.py
    # copy URL -> NEMOTRON_MODAL_ENDPOINT on the HF Space
"""

from __future__ import annotations

import hmac
import os
import time
from collections import deque

import modal
from fastapi import HTTPException, Request

# Guardian model (OpenBMB MiniCPM). MiniCPM4-8B: the warm replica (min_containers=1) kills the cold
# start, and at ~2–4s/turn warm it's snappy enough for a turn-based game while staying COHERENT - the
# 0.5B is far faster but too garbled to roleplay/refuse, and the 3-4B sizes fight the transformers
# version. Override with GUARDIAN_HF_MODEL; re-run the difficulty battery if you change the size.
MODEL_ID = os.environ.get("GUARDIAN_HF_MODEL", "openbmb/MiniCPM4-8B")
APP_NAME = os.environ.get("MODAL_APP_NAME", "jailbreak-dojo-guardian")
VOLUME_NAME = os.environ.get("MODAL_VOLUME", "guardian-model-cache")

# Per-container rate limit (best-effort; not shared across replicas).
_RATE_WINDOW_SEC = 60
_RATE_MAX_REQUESTS = 30
_rate_times: deque[float] = deque()

MAX_MESSAGES = 20
MAX_CONTENT_LEN = 4000
MAX_NEW_TOKENS = 256
_VALID_ROLES = frozenset({"user", "assistant", "system"})

app = modal.App(APP_NAME)

# transformers pin is MODEL-SPECIFIC. MiniCPM3-4B's trust_remote_code uses DynamicCache.seen_tokens,
# removed in transformers 4.41 → it needs <4.41. (MiniCPM4-8B instead needs >=4.44,<5 - if you swap
# back to the 8B via GUARDIAN_HF_MODEL, change this pin to "transformers>=4.44,<5".)
# transformers MUST stay <5 for the MiniCPM4 family: their trust_remote_code modeling imports
# `is_torch_fx_available`, which transformers 5.x removed.
image = modal.Image.debian_slim().pip_install(
    "transformers>=4.44,<5", "torch>=2.2", "accelerate>=0.30", "sentencepiece>=0.2",
    "fastapi[standard]",  # the web endpoint is a method on this GPU container
)
cache = modal.Volume.from_name(VOLUME_NAME, create_if_missing=True)


def _check_rate_limit() -> None:
    now = time.time()
    while _rate_times and _rate_times[0] < now - _RATE_WINDOW_SEC:
        _rate_times.popleft()
    if len(_rate_times) >= _RATE_MAX_REQUESTS:
        raise HTTPException(status_code=429, detail="Rate limit exceeded - try again shortly.")
    _rate_times.append(now)


def _check_api_key(request: Request) -> None:
    expected = os.environ.get("MODAL_API_KEY")
    if not expected:
        return  # open endpoint when unset (local dev only - set a secret in production)
    auth = request.headers.get("Authorization", "")
    provided = request.headers.get("X-API-Key") or auth.removeprefix("Bearer ").strip()
    if not hmac.compare_digest(provided, expected):
        raise HTTPException(status_code=401, detail="Unauthorized")


def _validate_messages(payload: dict) -> list[dict]:
    messages = payload.get("messages")
    if not isinstance(messages, list) or not messages:
        raise HTTPException(status_code=400, detail="messages must be a non-empty list")
    if len(messages) > MAX_MESSAGES:
        raise HTTPException(status_code=400, detail=f"too many messages (max {MAX_MESSAGES})")
    for msg in messages:
        if not isinstance(msg, dict) or "role" not in msg or "content" not in msg:
            raise HTTPException(status_code=400, detail="each message needs role and content")
        if msg["role"] not in _VALID_ROLES:
            raise HTTPException(status_code=400, detail=f"invalid role (must be one of {sorted(_VALID_ROLES)})")
        if len(str(msg["content"])) > MAX_CONTENT_LEN:
            raise HTTPException(status_code=400, detail=f"message too long (max {MAX_CONTENT_LEN} chars)")
    return messages


@app.cls(
    gpu="L4",                 # 24 GB - fits MiniCPM4-8B in bf16; ~$0.80/hr
    image=image,
    volumes={"/cache": cache},
    secrets=[modal.Secret.from_name("jailbreak-dojo")],
    scaledown_window=300,     # stay warm 5 min after the last request
    timeout=600,
    min_containers=int(os.environ.get("MODAL_MIN_CONTAINERS", "0")),  # 1 = demo warm start; 0 = scale-to-zero
)
class Guardian:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer

        self.tok = AutoTokenizer.from_pretrained(
            MODEL_ID, cache_dir="/cache", trust_remote_code=True
        )
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, cache_dir="/cache", dtype=torch.bfloat16,
            device_map="auto", trust_remote_code=True,
        )

    def _generate(self, messages: list[dict], max_new_tokens: int, temperature: float) -> dict:
        import torch

        max_new_tokens = min(max(1, max_new_tokens), MAX_NEW_TOKENS)
        inputs = self.tok.apply_chat_template(
            messages, add_generation_prompt=True, return_tensors="pt", return_dict=True
        ).to(self.model.device)
        in_len = inputs["input_ids"].shape[1]
        with torch.no_grad():
            out = self.model.generate(
                **inputs, max_new_tokens=max_new_tokens,
                do_sample=temperature > 0, temperature=temperature, top_p=0.9,
            )
        gen = out[0, in_len:]
        return {
            "text": self.tok.decode(gen, skip_special_tokens=True),
            "model": MODEL_ID,
            "tokens": {
                "input": int(in_len),
                "output": int(gen.shape[0]),
                "total": int(in_len + gen.shape[0]),
            },
        }

    @modal.method()
    def generate(self, messages: list[dict], max_new_tokens: int = 64, temperature: float = 0.3) -> dict:
        return self._generate(messages, max_new_tokens, temperature)

    @modal.fastapi_endpoint(method="POST")
    def web_generate(self, payload: dict, request: Request) -> dict:
        """HTTP endpoint the Gradio Space calls (runs in this GPU container, model preloaded).

        Request:  {"messages": [{"role": ..., "content": ...}, ...], "max_new_tokens": 64}
        Response: {"text": str, "tokens": {"input": int, "output": int, "total": int}}
        """
        _check_api_key(request)
        _check_rate_limit()
        messages = _validate_messages(payload)
        return self._generate(
            messages,
            int(payload.get("max_new_tokens", 64)),
            float(payload.get("temperature", 0.3)),
        )