"""Modal backend: serves MiniCPM4.1-8B on vLLM with an OpenAI-compatible API.

Deploy:
  modal secret create storycode-api MODAL_API_KEY=<shared-secret>
  modal deploy modal_app.py

The deploy prints a URL; the OpenAI base_url is that URL + "/v1". Put it in the
HF Space secret MODAL_ENDPOINT_URL, and set MODAL_API_KEY to the same shared
secret. The Space (app.py / llm.py) then calls this as a normal OpenAI endpoint.

We rely on vLLM's structured outputs (`guided_json`, xgrammar backend) so the
per-file summaries and the project story are schema-valid by construction — the
Space never has to repair model output.

IMPORTANT before deploying:
  * Pin VLLM_VERSION to whatever the MiniCPM4.1-8B model card lists as supported.
    MiniCPM4.1 uses custom code, hence --trust-remote-code.
  * MiniCPM4.1-8B is a hybrid-reasoning model. We disable the deep "thinking"
    mode for this task (we want fast, grounded narration, not chain-of-thought).
"""
from __future__ import annotations

import modal

MODEL_ID = "openbmb/MiniCPM4.1-8B"
# The MiniCPM4.1-8B model card currently says to install the latest pre-release
# vLLM wheel from the nightly index for standard vLLM inference.
VLLM_PACKAGE = "vllm"
VLLM_EXTRA_INDEX_URL = "https://wheels.vllm.ai/nightly"
GPU = "L4"                      # 8B in bf16 ~16GB; L4 (24GB) / A10G are comfortable
PORT = 8000
MAX_MODEL_LEN = 16384           # plenty for map-reduce; raise toward 32k if needed


def _download():
    from huggingface_hub import snapshot_download
    snapshot_download(MODEL_ID)


image = (
    modal.Image.from_registry("nvidia/cuda:13.0.2-devel-ubuntu22.04", add_python="3.11")
    .pip_install("huggingface_hub[hf_transfer]")
    .run_commands(
        f"python -m pip install -U {VLLM_PACKAGE} --pre --extra-index-url {VLLM_EXTRA_INDEX_URL}"
    )
    .env({"HF_XET_HIGH_PERFORMANCE": "1"})
    .run_function(_download)  # bake weights into the image -> no per-request download
)

app = modal.App("storycode-minicpm", image=image)


@app.function(
    gpu=GPU,
    scaledown_window=600,
    timeout=900,
    max_containers=1,
    secrets=[modal.Secret.from_name("storycode-api")],
)
@modal.concurrent(max_inputs=20)
@modal.web_server(port=PORT, startup_timeout=900)
def serve():
    import os
    import subprocess

    api_key = os.environ["MODAL_API_KEY"]
    cmd = [
        "vllm", "serve", MODEL_ID,
        "--host", "0.0.0.0", "--port", str(PORT),
        "--trust-remote-code",
        "--dtype", "bfloat16",
        "--max-model-len", str(MAX_MODEL_LEN),
        "--api-key", api_key,
    ]
    subprocess.Popen(cmd)


# Quick local smoke test:  modal run modal_app.py
@app.local_entrypoint()
def main():
    print("Deploy with:  modal deploy modal_app.py")
    print("Then set MODAL_ENDPOINT_URL = <printed-url>/v1 and MODAL_API_KEY to your shared secret.")