"""Modal vLLM serving layer for Thousand Token Wood.

Serves a small chat model (default Qwen2.5-3B-Instruct, <=4B for Tiny Titan
eligibility) behind a batched `chat_batch` method, so all creatures decide in
one GPU call per turn. The model is cached in a Modal Volume so cold starts
don't re-download it.

Deploy:  python -m modal deploy serve.py
Smoke:   python -m modal run serve.py
"""

import os

import modal

MODEL = os.environ.get("TTW_MODEL", "Qwen/Qwen2.5-3B-Instruct")
MODEL_REVISION = os.environ.get("TTW_MODEL_REVISION", "main")

app = modal.App("ttw-serve")

image = (
    modal.Image.debian_slim(python_version="3.12")
    .pip_install(
        "vllm==0.6.6",
        "huggingface_hub[hf_transfer]==0.26.2",
    )
    # Bake the chosen model into the image env so the CONTAINER sees it. A local
    # `TTW_MODEL=... modal deploy serve.py` env var does NOT propagate into the
    # Modal container on its own (it re-imports fresh), so capture it here.
    .env(
        {
            "HF_HUB_ENABLE_HF_TRANSFER": "1",
            "VLLM_DO_NOT_TRACK": "1",
            "TTW_MODEL": MODEL,
            "TTW_MODEL_REVISION": MODEL_REVISION,
        }
    )
)

# Persist the HF cache across cold starts so we download the model once.
hf_cache = modal.Volume.from_name("ttw-hf-cache", create_if_missing=True)
CACHE_DIR = "/root/.cache/huggingface"


@app.cls(
    gpu="L4",
    image=image,
    volumes={CACHE_DIR: hf_cache},
    scaledown_window=300,  # keep warm 5 min after last call (good for demos)
    timeout=600,
)
class Engine:
    @modal.enter()
    def load(self):
        from vllm import LLM, SamplingParams

        self.SamplingParams = SamplingParams
        self.llm = LLM(
            model=MODEL,
            revision=MODEL_REVISION,
            dtype="bfloat16",
            max_model_len=4096,
            gpu_memory_utilization=0.90,
            enforce_eager=True,  # faster cold start; we don't need CUDA graphs here
        )

    @modal.method()
    def chat_batch(
        self,
        conversations: list[list[dict]],
        max_tokens: int = 256,
        temperature: float = 0.7,
    ) -> list[str]:
        """Run a batch of chat conversations. Returns one completion per input."""
        params = self.SamplingParams(
            temperature=temperature, top_p=0.9, max_tokens=max_tokens
        )
        outputs = self.llm.chat(conversations, params)
        return [o.outputs[0].text for o in outputs]


@app.local_entrypoint()
def main():
    """Quick smoke test of the batched endpoint with two creature prompts."""
    convos = [
        [
            {"role": "system", "content": "You are Fenn, a sly fox speculator."},
            {"role": "user", "content": "Acorns crashed on a rumor. Buy or sell? One sentence."},
        ],
        [
            {"role": "system", "content": "You are Bramble, an anxious hoarder squirrel."},
            {"role": "user", "content": "Winter is coming. What do you stockpile? One sentence."},
        ],
    ]
    for text in Engine().chat_batch.remote(convos, max_tokens=60, temperature=0.7):
        print("-", text.strip())