#!/usr/bin/env python
"""Start the ACDiR-LLaDA OpenAI-compatible local server."""

from __future__ import annotations

import argparse
import json
import os
import subprocess
import sys
from pathlib import Path

from eval_math500 import resolve_base_model


ROOT = Path(__file__).resolve().parent
DEFAULT_CONFIG = ROOT / "configs" / "math500_44.json"


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=str(DEFAULT_CONFIG))
    parser.add_argument("--base_model", default="")
    parser.add_argument("--critic_ckpt", default="")
    parser.add_argument("--host", default="127.0.0.1")
    parser.add_argument("--port", type=int, default=23333)
    parser.add_argument("--model_id", default="acdir-llada-math500")
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    config_path = Path(args.config)
    if not config_path.is_absolute():
        config_path = ROOT / config_path
    with config_path.open("r", encoding="utf-8") as f:
        cfg = json.load(f)

    decode_cfg = cfg["decode"]
    runtime_cfg = cfg["runtime"]
    weights = cfg["released_weights"]
    base_model_source = args.base_model or cfg["base_model"]
    base_model_revision = "" if args.base_model else str(cfg.get("base_model_revision", ""))
    base_model = resolve_base_model(
        str(base_model_source),
        revision=base_model_revision,
        cache_dir=ROOT / ".cache" / "huggingface" / "hub",
    )
    critic_ckpt = Path(args.critic_ckpt or weights["critic"])
    if not critic_ckpt.is_absolute():
        critic_ckpt = ROOT / critic_ckpt

    env = os.environ.copy()
    env.setdefault("LLADA_EXACT_BACKEND", runtime_cfg["llada_exact_backend"])
    env.setdefault("LLADA_LMDEPLOY_FAST_MODE", runtime_cfg["llada_fast_mode"])
    env.setdefault("LLADA_LMDEPLOY_CUDAGRAPH", "0")
    env.setdefault("HF_HOME", str(ROOT / ".cache" / "huggingface"))

    cmd = [
        sys.executable,
        str(ROOT / "tools" / "acdir_openai_server.py"),
        "--host",
        args.host,
        "--port",
        str(args.port),
        "--model_id",
        args.model_id,
        "--ckpt_path",
        str(base_model),
        "--critic_ckpt_path",
        str(critic_ckpt),
        "--remask_method",
        decode_cfg["remask_method"],
        "--reforward_after_remask",
        "--no_deterministic_joint_argmax",
        "--force_remask_window",
        str(decode_cfg["force_remask_window"]),
        "--actor_forward_backend",
        runtime_cfg["actor_forward_backend"],
        "--actor_forward_dtype",
        runtime_cfg["actor_forward_dtype"],
        "--llada_fast_mode",
        runtime_cfg["llada_fast_mode"],
    ]
    print("[acdir] starting local server:", " ".join(cmd), flush=True)
    return subprocess.call(cmd, cwd=str(ROOT), env=env)


if __name__ == "__main__":
    raise SystemExit(main())