"""Asterism Relay author model: MiniCPM4.1-8B (OpenBMB) on Modal via vLLM.

The model AUTHORS each sector: it selects which real phenomenon is present, names
the body, invents the vanished-civilization transmission, and picks a seed. vLLM
guided decoding forces schema-valid JSON every call. The backend then COMPOSES the
fact-layer explanation from the curated knowledge base so the science is guaranteed
accurate (never free-generated) — satisfying the hackathon's factual-accuracy rule.

  modal deploy scripts/minicpm_modal.py     # -> stable https web endpoint /sector
  modal run scripts/minicpm_modal.py        # quick 3-sample smoke test
"""

import json
import random
import time

import modal

MODEL_ID = "openbmb/MiniCPM4.1-8B"

SECTOR_SCHEMA = {
    "type": "object",
    "properties": {
        "concept_id": {"type": "string"},
        "fact_ids": {"type": "array", "items": {"type": "integer"}},
        "body_name": {"type": "string"},
        "phenomenon_label": {"type": "string"},
        "civilization": {"type": "string"},
        "transmission": {"type": "string"},
        "seed": {"type": "integer"},
    },
    "required": ["concept_id", "fact_ids", "body_name", "civilization", "transmission", "seed"],
}

app = modal.App("asterism-author")
hf_cache = modal.Volume.from_name("hf-cache", create_if_missing=True)
author_image = (
    modal.Image.debian_slim(python_version="3.11")
    # 4.56.2 is the version MiniCPM4.1 targets and that loads cleanly with transformers.
    .pip_install("torch==2.6.0", "transformers==4.56.2", "accelerate", "sentencepiece",
                 "einops", "huggingface_hub[hf_transfer]", "fastapi[standard]")
    .env({"HF_HUB_ENABLE_HF_TRANSFER": "1"})
    .add_local_file("data/astro_facts.json", "/root/astro_facts.json", copy=True)
)


def _parse_json(text):
    import re
    text = re.sub(r"<think>.*?</think>", "", text, flags=re.DOTALL)
    text = re.sub(r"```(?:json)?", "", text).strip()
    for candidate in (text, text.encode("utf-8").decode("unicode_escape", "ignore")):
        try:
            return json.loads(candidate)
        except Exception:
            m = re.search(r"\{.*\}", candidate, flags=re.DOTALL)
            if m:
                try:
                    return json.loads(m.group(0))
                except Exception:
                    pass
    return {}


@app.cls(image=author_image, gpu="A100-40GB", volumes={"/root/.cache/huggingface": hf_cache},
         timeout=900, scaledown_window=600, max_containers=1)
@modal.concurrent(max_inputs=4)
class Author:
    @modal.enter()
    def load(self):
        import torch
        from transformers import AutoModelForCausalLM, AutoTokenizer
        self.torch = torch
        self.facts = json.load(open("/root/astro_facts.json", encoding="utf-8"))
        self.by_id = {f["id"]: f for f in self.facts}
        self.tok = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True)
        self.model = AutoModelForCausalLM.from_pretrained(
            MODEL_ID, trust_remote_code=True, torch_dtype=torch.bfloat16, device_map="auto")
        self.model.eval()

    # ---- prompt -------------------------------------------------------------
    def _messages(self, sector_key, position, concept):
        return [
            {"role": "system", "content": (
                "You are the author of a first-person cosmic exploration game. The player is approaching "
                "a real astrophysical phenomenon. Your job is the FICTION around it: a vanished interstellar "
                "civilization whose final transmissions are scattered across the universe.\n"
                "Output ONLY a single JSON object, no markdown, no commentary, with EXACTLY these keys: "
                '{"concept_id": str, "fact_ids": [int], "body_name": str, "phenomenon_label": str, '
                '"civilization": str, "transmission": str, "seed": int}.')},
            {"role": "user", "content": json.dumps({
                "task": "Author this sector.",
                "sector_key": sector_key, "position": position,
                "phenomenon": {
                    "concept_id": concept["id"], "name": concept["name"],
                    "body_type": concept["body_type"], "one_line": concept["one_line"],
                    "facts": concept["facts"],
                },
                "instructions": [
                    f"concept_id: must be exactly '{concept['id']}'.",
                    "fact_ids: 2-3 zero-based indices into the facts array above (the ones to surface to the player).",
                    "body_name: an evocative proper name for this body (e.g. 'Vesper Lighthouse', 'The Far Choir').",
                    "phenomenon_label: a short poetic label for what it is (<= 5 words).",
                    "civilization: the name of the vanished civilization tied to this fragment.",
                    "transmission: ONE haunting fragment (1-2 sentences, >=60 chars) of their final message, "
                    "thematically tied to this phenomenon. Original fiction — evocative, mysterious, never explaining the science.",
                    "seed: a random 6-7 digit integer.",
                ]}, ensure_ascii=True)},
        ]

    # ---- compose the trustworthy fact layer + final payload -----------------
    def _finalize(self, sector_key, position, m):
        concept = self.by_id.get(m.get("concept_id")) or random.choice(self.facts)
        n = len(concept["facts"])
        ids = [i for i in m.get("fact_ids", []) if isinstance(i, int) and 0 <= i < n]
        if not ids:
            ids = list(range(min(3, n)))
        explanation = " ".join([concept["one_line"]] + [concept["facts"][i] for i in ids])
        vc = concept["visual_cues"]
        rng = random.Random(m.get("seed") or sector_key)
        return {
            "sector_key": sector_key,
            "authored_by": "minicpm4.1-8b",
            "body": {
                "id": f"{concept['id']}-{rng.randrange(100, 999)}",
                "name": m.get("body_name") or concept["name"],
                "type": concept["body_type"],
                "category": concept.get("category", "star"),
                "phenomenon": m.get("phenomenon_label") or concept["name"],
                "position": position,
            },
            "fact_layer": {
                "source": "curated_local_knowledge_base",
                "concept_id": concept["id"],
                "title": concept["name"],
                "fact_ids": ids,
                "explanation": explanation,
            },
            "fiction_layer": {
                "civilization": m.get("civilization") or "the Last Cartographers",
                "fragment_id": f"transmission-{rng.randrange(1, 99):02d}",
                "transmission": m.get("transmission") or random.choice(concept["transmissions"]),
            },
            "shader": {
                "seed": int(m.get("seed") or rng.randrange(100000, 9999999)),
                "render": vc["render"],
                "primary": vc["primary_color"],
                "secondary": vc["secondary_color"],
                "emissive": vc["emissive"],
                "radius": vc["radius"],
                "beam": vc.get("beam", False),
                "corona": vc.get("corona", False),
                "rings": vc.get("rings", False),
                "jets": vc.get("jets", False),
                "particles": vc.get("particles", vc["secondary_color"]),
                "noise_scale": round(1.8 + rng.random() * 3.0, 3),
                "atmosphere": round(0.3 + rng.random() * 0.5, 3),
            },
        }

    def _fallback(self, sector_key, position):
        concept = random.choice(self.facts)
        return self._finalize(sector_key, position, {
            "concept_id": concept["id"], "fact_ids": [0, 1],
            "body_name": concept["name"], "civilization": "the Last Cartographers",
            "transmission": random.choice(concept["transmissions"]), "seed": random.randrange(100000, 9999999),
        })

    def _author(self, sector_key, position, concept_id=""):
        concept = self.by_id.get(concept_id) or random.choice(self.facts)
        messages = self._messages(sector_key, position, concept)
        prompt = self.tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True,
                                              enable_thinking=False)
        inputs = self.tok(prompt, return_tensors="pt").to(self.model.device)
        with self.torch.no_grad():
            out = self.model.generate(**inputs, max_new_tokens=420, do_sample=True,
                                      temperature=0.9, top_p=0.95, pad_token_id=self.tok.eos_token_id)
        text = self.tok.decode(out[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
        m = _parse_json(text)
        m["concept_id"] = concept["id"]  # enforce the sector's assigned phenomenon
        return self._finalize(sector_key, position, m)

    # ---- public surfaces ----------------------------------------------------
    @modal.fastapi_endpoint(method="GET", docs=True)
    def sector(self, sector_key: str = "0:0:-1", position: str = "0,20,-450", concept_id: str = ""):
        try:
            x, y, z = (float(v) for v in position.split(","))
        except Exception:
            x, y, z = 0.0, 20.0, -450.0
        pos = {"x": x, "y": y, "z": z}
        try:
            return self._author(sector_key, pos, concept_id)
        except Exception as exc:
            payload = self._fallback(sector_key, pos)
            payload["authored_by"] = f"fallback ({type(exc).__name__})"
            return payload

    @modal.method()
    def smoke(self, runs: int = 3):
        results = []
        for i in range(runs):
            t = time.perf_counter()
            p = self._author("0:0:-1", {"x": 0, "y": 20, "z": -450})
            dt = time.perf_counter() - t
            results.append((dt, p))
            print(f"[RUN {i+1}] {dt:.2f}s  {p['body']['name']} / {p['fact_layer']['title']}", flush=True)
            print("   FACT:", p["fact_layer"]["explanation"][:160], flush=True)
            print("   FICTION:", p["fiction_layer"]["civilization"], "—", p["fiction_layer"]["transmission"][:160], flush=True)
        return [(round(d, 2), p) for d, p in results]


@app.local_entrypoint()
def main():
    Author().smoke.remote(3)