"""Self-host Beam-1.0 (NVFP4) on Modal with SGLang to generate distillation traces.

  modal run modal_app/beam_selfhost.py::download            # one-time: weights -> Volume (CPU only)
  modal run modal_app/beam_selfhost.py::bench --ids FILE    # throughput test + parity on 2x B300

Every generation is appended to the `beam-dory-data` Volume under selfhost/, in the
same record format as scripts/generate.py, so build_dataset.py can verify and merge it.
"""

import json
import os
import subprocess
import time
from pathlib import Path

import modal

APP = modal.App("beam-dory-selfhost")
weights = modal.Volume.from_name("beam-weights", create_if_missing=True, version=2)
data = modal.Volume.from_name("beam-dory-data", create_if_missing=True, version=2)

BEAM_REPO = "reflection/Beam-1.0-501B-A23B-NVFP4"
FP8_REPO = "reflection/Beam-1.0-501B-A23B-FP8"
DRAFT_REPO = "reflection/Beam-1.0-501B-A23B-DFlash2"
MODELS = {"nvfp4": ("beam-nvfp4", 2), "fp8": ("beam-fp8", 4)}  # weights subdir, tensor-parallel GPUs
# GPU configs that fit each checkpoint, in order of preference (cost-effectiveness)
GPU_CHOICES = {"nvfp4": ["B300:2"], "fp8": ["B300:4", "H200:8", "B200:8"],
               "auto": ["B300:4", "H200:8", "B200:8", "B200:4", "B300:2"]}  # auto: FP8 if it fits, else NVFP4
GPU_MEM_GB = {"B300": 288, "B200": 180, "H200": 141, "H100": 80}
GPU_RATES = {"B300": 7.10, "B200": 6.25, "H200": 4.54, "H100": 3.95}  # $/GPU-hour on Modal
W = "/weights"
D = "/data"

dl_image = modal.Image.debian_slim(python_version="3.12").pip_install("huggingface_hub[hf_xet]>=0.34")
sglang_image = (
    modal.Image.from_registry("lmsysorg/sglang:beam")
    .pip_install("httpx")
    .env({"HF_HUB_OFFLINE": "1", "SGLANG_DISABLE_REQUEST_LOGGING": "1"})
)


@APP.function(image=dl_image, volumes={W: weights}, cpu=8, memory=16384, timeout=4 * 3600)
def download(which: str = "nvfp4"):
    from huggingface_hub import snapshot_download

    repos = {"nvfp4": [(BEAM_REPO, "beam-nvfp4"), (DRAFT_REPO, "dflash2")], "fp8": [(FP8_REPO, "beam-fp8")]}[which]
    for repo, sub in repos:
        t0 = time.time()
        path = snapshot_download(repo, local_dir=f"{W}/{sub}", max_workers=16)
        size = sum(f.stat().st_size for f in Path(path).rglob("*") if f.is_file())
        print(f"{repo}: {size / 1e9:.1f} GB in {time.time() - t0:.0f}s", flush=True)
        weights.commit()


def start_server(tp, spec, context_len, mem_frac, model_dir="beam-nvfp4", attn="trtllm_mha"):
    cmd = [
        "python", "-m", "sglang.launch_server",
        "--model-path", f"{W}/{model_dir}", "--served-model-name", "beam",
        "--tp", str(tp), "--attention-backend", attn, "--page-size", "64",
        "--context-length", str(context_len), "--mem-fraction-static", str(mem_frac),
        "--reasoning-parser", "beam", "--tool-call-parser", "beam",
        "--host", "127.0.0.1", "--port", "30000",
    ]
    if spec:
        cmd += ["--speculative-algorithm", "DFLASH", "--speculative-draft-model-path", f"{W}/dflash2",
                "--speculative-dflash-block-size", "8"]
    print("launch:", " ".join(cmd), flush=True)
    log = open("/tmp/sglang.log", "w")
    return subprocess.Popen(cmd, stdout=log, stderr=subprocess.STDOUT)


def wait_ready(proc, timeout_s):
    import httpx

    t0 = time.time()
    while time.time() - t0 < timeout_s:
        if proc.poll() is not None:
            raise RuntimeError("sglang exited:\n" + Path("/tmp/sglang.log").read_text()[-6000:])
        try:
            if httpx.get("http://127.0.0.1:30000/health", timeout=5).status_code == 200:
                return time.time() - t0
        except httpx.HTTPError:
            pass
        time.sleep(10)
    raise TimeoutError("sglang not ready:\n" + Path("/tmp/sglang.log").read_text()[-6000:])


async def run_phase(items, concurrency, seconds, out_f, label, max_tokens):
    """Pull prompts from items with `concurrency` workers; stop issuing after `seconds`."""
    import asyncio

    import httpx

    queue = asyncio.Queue()
    for it in items:
        queue.put_nowait(it)
    stats = {"n": 0, "tokens": 0, "errors": 0}
    t_start = time.time()
    deadline = t_start + seconds if seconds else None

    async def worker(client):
        while not queue.empty():
            if deadline and time.time() > deadline:
                return
            it = queue.get_nowait()
            body = {
                "model": "beam",
                "messages": [{"role": "user", "content": it["prompt"]}],
                "max_tokens": max_tokens,
                "temperature": 0.7,
                "top_p": 0.99,
                "chat_template_kwargs": {"reasoning_effort": "medium"},
                "separate_reasoning": True,
            }
            t0 = time.time()
            try:
                r = await client.post("http://127.0.0.1:30000/v1/chat/completions", json=body)
                r.raise_for_status()
                resp = r.json()
            except Exception as e:  # keep the phase going; record the failure
                stats["errors"] += 1
                out_f.write(json.dumps({"id": it["id"], "phase": label, "error": repr(e)[:300]}) + "\n")
                continue
            choice = resp["choices"][0]
            usage = resp.get("usage") or {}
            stats["n"] += 1
            stats["tokens"] += usage.get("completion_tokens", 0)
            rec = {k: it.get(k) for k in ("id", "domain", "source", "prompt", "difficulty", "topic", "answer_type")}
            rec.update({
                "effort": "medium",
                "gold": it["answer"],
                "teacher": "beam-nvfp4-selfhost-sglang",
                "phase": label,
                "reasoning": choice["message"].get("reasoning_content") or "",
                "content": choice["message"].get("content") or "",
                "finish_reason": choice.get("finish_reason"),
                "usage": usage,
                "latency_s": round(time.time() - t0, 2),
                "time": time.strftime("%Y-%m-%dT%H:%M:%S"),
            })
            out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
            out_f.flush()

    limits = httpx.Limits(max_connections=concurrency + 8)
    async with httpx.AsyncClient(timeout=httpx.Timeout(1800), limits=limits) as client:
        await asyncio.gather(*(worker(client) for _ in range(concurrency)))
    wall = time.time() - t_start
    stats.update({"label": label, "concurrency": concurrency, "wall_s": round(wall, 1),
                  "out_tok_per_s": round(stats["tokens"] / wall, 1)})
    print("PHASE", json.dumps(stats), flush=True)
    return stats


@APP.function(image=sglang_image, gpu="B300:2", volumes={W: weights, D: data}, timeout=60 * 60,
              memory=262144, cpu=16)
def bench(items: list, parity_ids: list, sweep: list, phase_seconds: int = 300, spec: bool = False,
          max_tokens: int = 16384, tag: str = "bench"):
    import asyncio

    t_boot = time.time()
    proc = start_server(tp=2, spec=spec, context_len=32768, mem_frac=0.88)
    ready_s = wait_ready(proc, timeout_s=30 * 60)
    print(f"server ready after {ready_s:.0f}s", flush=True)
    print(subprocess.run(["nvidia-smi", "--query-gpu=name,memory.used,memory.total", "--format=csv"],
                         capture_output=True, text=True).stdout, flush=True)

    out_dir = Path(D) / "selfhost"
    out_dir.mkdir(parents=True, exist_ok=True)
    out_path = out_dir / f"{tag}-{time.strftime('%Y%m%d-%H%M%S')}.jsonl"
    parity = [it for it in items if it["id"] in set(parity_ids)]
    fresh = [it for it in items if it["id"] not in set(parity_ids)]
    results = {"ready_s": round(ready_s, 1), "spec": spec, "phases": [], "out_file": str(out_path)}
    with open(out_path, "a") as f:
        if parity:
            results["phases"].append(asyncio.run(run_phase(parity, len(parity), 0, f, "parity", max_tokens)))
            data.commit()
        pos = 0
        for conc in sweep:
            chunk = fresh[pos:]
            st = asyncio.run(run_phase(chunk, conc, phase_seconds, f, f"c{conc}", max_tokens))
            pos += st["n"] + st["errors"]
            results["phases"].append(st)
            data.commit()
    proc.terminate()
    results["total_s"] = round(time.time() - t_boot, 1)
    print("RESULTS", json.dumps(results), flush=True)
    return results


@APP.local_entrypoint()
def main(ids: str = "data/selfhost_test_ids.txt", parity: str = "data/pilot_ids.txt",
         sweep: str = "64,128,256", phase_seconds: int = 300, spec: bool = False):
    root = Path(__file__).resolve().parent.parent
    pool = {}
    for name in ("pool.jsonl", "pool_ext.jsonl"):
        for line in open(root / "data" / name):
            r = json.loads(line)
            pool[r["id"]] = r
    wanted = [l.strip() for l in open(root / ids) if l.strip()]
    parity_ids = [l.strip() for l in open(root / parity) if l.strip()]
    items = [pool[i] for i in parity_ids + wanted]
    res = bench.remote(items, parity_ids, [int(x) for x in sweep.split(",")], phase_seconds, spec)
    print(json.dumps(res, indent=1))


TEACHER = "beam-selfhost-sglang"


async def produce(items, n, effort, concurrency, deadline, out_path, max_tokens, log_every=120):
    """Long-running generation: n samples per prompt, stop issuing at `deadline`, commit periodically."""
    import asyncio

    import httpx

    queue = asyncio.Queue()
    for it in items:
        queue.put_nowait(it)
    st = {"prompts": 0, "samples": 0, "tokens": 0, "errors": 0}
    t0 = time.time()
    out_f = open(out_path, "a")

    async def worker(client):
        while not queue.empty() and time.time() < deadline:
            it = queue.get_nowait()
            eff = it.get("_effort", effort)  # multi-effort runs tag each queue item with its effort
            body = {"model": "beam", "messages": [{"role": "user", "content": it["prompt"]}],
                    "max_tokens": max_tokens, "temperature": 0.7, "top_p": 0.99, "n": n,
                    "chat_template_kwargs": {"reasoning_effort": eff}, "separate_reasoning": True}
            t1 = time.time()
            try:
                r = await client.post("http://127.0.0.1:30000/v1/chat/completions", json=body)
                r.raise_for_status()
                resp = r.json()
            except Exception as e:
                st["errors"] += 1
                out_f.write(json.dumps({"id": it["id"], "effort": eff, "error": repr(e)[:300]}) + "\n")
                continue
            usage = resp.get("usage") or {}
            per = usage.get("completion_tokens", 0) / max(1, len(resp["choices"]))
            for k, choice in enumerate(resp["choices"]):
                rec = {k2: it.get(k2) for k2 in ("id", "domain", "source", "prompt", "difficulty", "topic", "answer_type")}
                rec.update({"effort": eff, "gold": it["answer"], "teacher": TEACHER,
                            "sample_idx": k, "reasoning": choice["message"].get("reasoning_content") or "",
                            "content": choice["message"].get("content") or "",
                            "finish_reason": choice.get("finish_reason"),
                            # SGLang reports usage for the whole request; split evenly as an estimate
                            "usage": {"completion_tokens": round(per), "request_completion_tokens": usage.get("completion_tokens")},
                            "latency_s": round(time.time() - t1, 2), "time": time.strftime("%Y-%m-%dT%H:%M:%S")})
                out_f.write(json.dumps(rec, ensure_ascii=False) + "\n")
                st["samples"] += 1
            st["prompts"] += 1
            st["tokens"] += usage.get("completion_tokens", 0)

    async def reporter():
        while True:
            await asyncio.sleep(log_every)
            out_f.flush()
            await asyncio.to_thread(data.commit)
            el = time.time() - t0
            print(f"PROGRESS {time.strftime('%H:%M:%S')} prompts={st['prompts']} samples={st['samples']} "
                  f"err={st['errors']} tok/s={st['tokens'] / el:.0f} left={queue.qsize()} "
                  f"stop_in={(deadline - time.time()) / 60:.0f}min", flush=True)

    limits = httpx.Limits(max_connections=concurrency + 8)
    async with httpx.AsyncClient(timeout=httpx.Timeout(3600), limits=limits) as client:
        rep = asyncio.create_task(reporter())
        await asyncio.gather(*(worker(client) for _ in range(concurrency)))
        rep.cancel()
    out_f.close()
    data.commit()
    st["wall_s"] = round(time.time() - t0)
    return st


def _run_impl(prompts_file: str, tag: str, n: int = 2, effort: str = "medium", concurrency: int = 320,
              max_hours: float = 9.0, max_tokens: int = 16384, model: str = "nvfp4", max_dollars: float = 0):
    global TEACHER
    """Production run over a prompt file on the data Volume. Resumable: (prompt, effort) pairs already in
    selfhost/run-<tag>-*.jsonl are skipped (survives preemption). `effort` may list several levels
    ("low,high"): every prompt is queued once per level, interleaved so all levels progress together."""
    import asyncio

    t_start = time.time()
    deadline = t_start + max_hours * 3600
    items = [json.loads(l) for l in open(f"{D}/{prompts_file}")]
    out_dir = Path(D) / "selfhost"
    out_dir.mkdir(parents=True, exist_ok=True)
    done = set()
    for f in out_dir.glob(f"run-{tag}-*.jsonl"):
        for line in open(f):
            try:
                r = json.loads(line)
            except json.JSONDecodeError:  # torn last line after preemption
                continue
            if "error" not in r:
                done.add((r["id"], r.get("effort", "medium")))
    efforts = effort.split(",")
    todo = [dict(it, _effort=e) for it in items for e in efforts if (it["id"], e) not in done]
    TEACHER = f"beam-{model}-selfhost-sglang"
    print(f"{len(items)} prompts, {len(done)} already done, {len(todo)} to go; n={n} effort={effort}", flush=True)

    gpus = subprocess.run(["nvidia-smi", "--query-gpu=name", "--format=csv,noheader"],
                          capture_output=True, text=True).stdout.split("\n")
    gpus = [g for g in gpus if g.strip()]
    kind = next(k for k in GPU_RATES if k in gpus[0])
    tp = len(gpus)
    if model == "auto":  # FP8 (503 GB) needs ~1 TB of GPU memory; otherwise use NVFP4 (294 GB)
        model = "fp8" if tp * GPU_MEM_GB[kind] >= 1000 else "nvfp4"
        TEACHER = f"beam-{model}-selfhost-sglang"
    model_dir, _ = MODELS[model]
    rate = tp * GPU_RATES[kind] + 2.8  # $/h incl. CPU and memory
    if max_dollars:
        deadline = min(deadline, t_start + max_dollars / rate * 3600)
    attn = "fa3" if kind in ("H200", "H100") else "trtllm_mha"
    print(f"got {tp}x {kind} (~${rate:.1f}/h); model={model}; attention={attn}; stop issuing in "
          f"{(deadline - time.time()) / 3600:.2f}h", flush=True)
    proc = start_server(tp=tp, spec=False, context_len=32768, mem_frac=0.88, model_dir=model_dir, attn=attn)
    print(f"server ready after {wait_ready(proc, timeout_s=40 * 60):.0f}s", flush=True)
    out_path = out_dir / f"run-{tag}-{time.strftime('%Y%m%d-%H%M%S')}.jsonl"
    st = asyncio.run(produce(todo, n, effort, concurrency, deadline - 300, out_path, max_tokens))
    proc.terminate()
    st["total_h"] = round((time.time() - t_start) / 3600, 2)
    print("RUN DONE", json.dumps(st), flush=True)
    return st


RUN_KW = dict(image=sglang_image, volumes={W: weights, D: data}, timeout=12 * 3600, memory=262144, cpu=16)


@APP.function(gpu=GPU_CHOICES["nvfp4"], **RUN_KW)
def run(*args, **kwargs):
    return _run_impl(*args, **kwargs)


@APP.function(gpu=GPU_CHOICES["fp8"], **RUN_KW)
def run_fp8(*args, **kwargs):
    return _run_impl(*args, **kwargs)


@APP.function(gpu=GPU_CHOICES["auto"], **RUN_KW)
def run_auto(*args, **kwargs):
    return _run_impl(*args, **kwargs)


@APP.local_entrypoint()
def launch(prompts: str, tag: str, model: str = "fp8", n: int = 1, effort: str = "medium",
           concurrency: int = 384, max_hours: float = 8.0, max_dollars: float = 0, gpu: str = ""):
    """prompts: path inside the beam-dory-data Volume (upload first with `modal volume put`).
    Takes the first available GPU config from GPU_CHOICES[model]; max_dollars caps the run."""
    fn = {"fp8": run_fp8, "auto": run_auto}.get(model, run)
    if gpu:  # pin one exact GPU config instead of the fallback list
        fn = fn.with_options(gpu=gpu)
    res = fn.remote(prompts, tag, n, effort, concurrency, max_hours, 16384, model, max_dollars)
    print(json.dumps(res, indent=1))