#!/usr/bin/env python3
"""Run a one-shot DeepSeek V4 Flash expert-pack benchmark on one RTX 5090."""

from __future__ import annotations

import argparse
import fcntl
import hashlib
import json
import os
import signal
import socket
import subprocess
import sys
import time
import urllib.request
from pathlib import Path
from typing import Any

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_PROMPT = "Please introduce Shenzhen"
CHAT_PREFIX = "<\uff5cbegin\u2581of\u2581sentence\uff5c>You are a helpful assistant.<\uff5cUser\uff5c>"
CHAT_SUFFIX = "<\uff5cAssistant\uff5c><think>"
DEFAULT_LOCK = Path("/tmp/sglang-deepseek-v4-5090-benchmark.lock")
METADATA_FORMAT_VERSION = 4
ACTIVE_MOE_LAYERS = tuple(range(43))


def cache_root() -> Path:
    return Path(os.environ.get("XDG_CACHE_HOME", Path.home() / ".cache")).expanduser()


def artifact_dir_for_source(path: Path) -> Path:
    stat = path.stat()
    fingerprint = hashlib.sha256(
        f"{path.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:{METADATA_FORMAT_VERSION}".encode()
    ).hexdigest()[:20]
    return cache_root() / "sglang-expert-pack" / "deepseek-v4-flash" / fingerprint


def find_sglang_repo() -> Path:
    configured = os.environ.get("SGLANG_REPO")
    if configured:
        return Path(configured).expanduser().resolve()
    for candidate in (SCRIPT_DIR, *SCRIPT_DIR.parents):
        if (candidate / "python" / "sglang").is_dir() and (
            candidate / "python" / "sglang" / "srt" / "model_loader" / "expert_pack"
        ).is_dir():
            return candidate
    raise RuntimeError("could not locate the SGLang repository")


def format_prompt(prompt: str) -> str:
    return f"{CHAT_PREFIX}{prompt}{CHAT_SUFFIX}"


def server_url(args: argparse.Namespace) -> str:
    return f"http://{args.host}:{args.port}"


def port_in_use(host: str, port: int, timeout: float = 0.5) -> bool:
    try:
        with socket.create_connection((host, port), timeout=timeout):
            return True
    except OSError:
        return False


def detect_rtx_5090() -> str:
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=index,name", "--format=csv,noheader"],
        check=True,
        capture_output=True,
        text=True,
    )
    rows = [line.strip() for line in result.stdout.splitlines() if line.strip()]
    gpu_zero = next(
        (line.split(",", 1)[1].strip() for line in rows if line.startswith("0,")),
        None,
    )
    if gpu_zero is None or "5090" not in gpu_zero:
        raise RuntimeError(
            f"CUDA device 0 must be an RTX 5090; detected: {', '.join(rows) or 'none'}"
        )
    return gpu_zero


def build_server_command(args: argparse.Namespace) -> list[str]:
    extra_config = {
        "cache_vram_mib": args.expert_cache_mib,
        "cache_vram_reserve_mib": args.expert_cache_reserve_mib,
        "stage_slots": args.stage_slots,
        "read_splits": args.read_splits,
        "direct_io": args.direct_io,
        "stats_flush_interval": len(ACTIVE_MOE_LAYERS),
        "stats_path": str(args.stats_path),
    }
    return [
        sys.executable,
        "-m",
        "sglang.launch_server",
        "--model-path",
        str(args.gguf),
        "--trust-remote-code",
        "--load-format",
        "expert_pack",
        "--model-loader-extra-config",
        json.dumps(extra_config, separators=(",", ":")),
        "--attention-backend",
        "dsv4",
        "--tp-size",
        "1",
        "--ep-size",
        "1",
        "--disable-flashinfer-autotune",
        "--skip-server-warmup",
        "--context-length",
        str(args.context_length),
        "--max-total-tokens",
        str(args.max_total_tokens),
        "--max-running-requests",
        "1",
        "--mem-fraction-static",
        str(args.mem_fraction_static),
        "--watchdog-timeout",
        str(args.watchdog_timeout),
        "--host",
        args.host,
        "--port",
        str(args.port),
    ]


def start_server(args: argparse.Namespace) -> subprocess.Popen:
    if port_in_use(args.host, args.port):
        raise RuntimeError(f"server address is already in use: {server_url(args)}")
    args.server_log.parent.mkdir(parents=True, exist_ok=True)
    log = args.server_log.open("wb", buffering=0)
    env = os.environ.copy()
    python_path = [str(args.sglang_repo), str(args.sglang_repo / "python")]
    if env.get("PYTHONPATH"):
        python_path.append(env["PYTHONPATH"])
    env["CUDA_VISIBLE_DEVICES"] = "0"
    env["PYTHONPATH"] = os.pathsep.join(python_path)
    env.setdefault("SGLANG_OPT_USE_TILELANG_INDEXER", "1")
    conda_lib = str(Path(sys.prefix) / "lib")
    cuda_root = Path("/usr/local/cuda")
    if (cuda_root / "bin" / "nvcc").is_file():
        env["CUDA_HOME"] = str(cuda_root)
        env["CUDA_PATH"] = str(cuda_root)
        env["PATH"] = os.pathsep.join((str(cuda_root / "bin"), env.get("PATH", "")))
    env["LD_LIBRARY_PATH"] = os.pathsep.join(
        value
        for value in (
            conda_lib,
            str(cuda_root / "lib64") if (cuda_root / "lib64").is_dir() else None,
            env.get("LD_LIBRARY_PATH"),
        )
        if value
    )
    command = build_server_command(args)
    print(
        f"SERVICE_STARTING url={server_url(args)} timeout={args.startup_timeout:.0f}s "
        f"log={args.server_log}",
        flush=True,
    )
    process = subprocess.Popen(
        command,
        cwd=args.sglang_repo,
        stdin=subprocess.DEVNULL,
        stdout=log,
        stderr=subprocess.STDOUT,
        env=env,
        start_new_session=True,
    )
    process._benchmark_log = log  # type: ignore[attr-defined]
    try:
        deadline = time.monotonic() + args.startup_timeout
        while time.monotonic() < deadline:
            if process.poll() is not None:
                raise RuntimeError(
                    f"SGLang exited during startup with code {process.returncode}; "
                    f"see {args.server_log}"
                )
            if port_in_use(args.host, args.port):
                print(
                    f"SERVICE_READY pid={process.pid} url={server_url(args)}",
                    flush=True,
                )
                return process
            time.sleep(2)
        raise TimeoutError(
            f"SGLang did not become ready within {args.startup_timeout:.0f}s; "
            f"see {args.server_log}"
        )
    except BaseException:
        stop_server(process, args)
        raise


def stop_server(process: subprocess.Popen | None, args: argparse.Namespace) -> None:
    if process is None:
        return
    log = getattr(process, "_benchmark_log", None)
    try:
        if process.poll() is None:
            try:
                os.killpg(process.pid, signal.SIGTERM)
            except ProcessLookupError:
                pass
            try:
                process.wait(timeout=45)
            except subprocess.TimeoutExpired:
                try:
                    os.killpg(process.pid, signal.SIGKILL)
                except ProcessLookupError:
                    pass
                process.wait(timeout=15)
        deadline = time.monotonic() + 10
        while (
            port_in_use(args.host, args.port, timeout=0.2)
            and time.monotonic() < deadline
        ):
            time.sleep(0.2)
        print(f"SERVICE_STOPPED pid={process.pid} url={server_url(args)}", flush=True)
    finally:
        if log is not None:
            log.close()


def generate(
    args: argparse.Namespace,
    prompt: str,
    max_new_tokens: int,
    *,
    stream_output: bool,
) -> dict[str, Any]:
    payload = {
        "text": format_prompt(prompt),
        "sampling_params": {
            "temperature": args.temperature,
            "top_p": args.top_p,
            "sampling_seed": args.seed,
            "max_new_tokens": max_new_tokens,
            "ignore_eos": True,
        },
        "stream": True,
    }
    request = urllib.request.Request(
        server_url(args) + "/generate",
        data=json.dumps(payload, ensure_ascii=False).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter_ns()
    first_token = last_token = None
    completion_tokens = 0
    prompt_tokens = None
    output = ""
    finish_reason = None
    if stream_output:
        print(f"prompt: {prompt}", flush=True)
        print("output: ", end="", flush=True)

    with urllib.request.urlopen(request, timeout=args.request_timeout) as response:
        for raw_line in response:
            now = time.perf_counter_ns()
            line = raw_line.decode("utf-8").strip()
            if not line:
                continue
            if line.startswith("data: "):
                line = line[6:]
            if line == "[DONE]":
                continue
            event = json.loads(line)
            meta = event.get("meta_info") or {}
            current_tokens = int(meta.get("completion_tokens", 0))
            if current_tokens > completion_tokens:
                first_token = first_token or now
                last_token = now
                completion_tokens = current_tokens
            if meta.get("prompt_tokens") is not None:
                prompt_tokens = int(meta["prompt_tokens"])
            event_output = event.get("text")
            if event_output is not None:
                if stream_output and event_output != output:
                    if event_output.startswith(output):
                        print(event_output[len(output) :], end="", flush=True)
                    else:
                        print(f"\n[output revised]\n{event_output}", end="", flush=True)
                output = event_output
            finish_reason = meta.get("finish_reason", finish_reason)
    if stream_output:
        print(flush=True)
    if first_token is None or last_token is None or prompt_tokens is None:
        raise RuntimeError(
            "SGLang response did not contain complete token timing metadata"
        )
    ttft_s = (first_token - started) / 1e9
    decode_span_s = (last_token - first_token) / 1e9
    total_s = (time.perf_counter_ns() - started) / 1e9
    decode_intervals = max(0, completion_tokens - 1)
    return {
        "prompt": prompt,
        "output": output,
        "prompt_tokens": prompt_tokens,
        "completion_tokens": completion_tokens,
        "finish_reason": finish_reason,
        "ttft_ms": ttft_s * 1000,
        "prefill_token_rate": prompt_tokens / ttft_s if ttft_s > 0 else None,
        "decode_token_rate": (
            decode_intervals / decode_span_s if decode_span_s > 0 else None
        ),
        "tpot_ms": (
            decode_span_s * 1000 / decode_intervals if decode_intervals else None
        ),
        "total_elapsed_s": total_s,
        "end_to_end_token_rate": completion_tokens / total_s if total_s > 0 else None,
    }


def run_benchmark(args: argparse.Namespace) -> dict[str, Any]:
    process = None
    try:
        if args.stats_path.exists():
            args.stats_path.unlink()
        process = start_server(args)
        return generate(args, args.prompt, args.max_new_tokens, stream_output=True)
    finally:
        stop_server(process, args)


def read_stats(path: Path) -> dict[str, Any]:
    if not path.is_file():
        raise FileNotFoundError(f"expert-pack stats were not written: {path}")
    return json.loads(path.read_text(encoding="utf-8"))


def audit_routes(stats: dict[str, Any], expected_tokens: int) -> None:
    token_counts = stats.get("route_tokens_by_layer") or []
    call_counts = stats.get("route_calls_by_layer") or []
    if len(token_counts) != len(ACTIVE_MOE_LAYERS) or len(call_counts) != len(
        ACTIVE_MOE_LAYERS
    ):
        raise RuntimeError("DeepSeek Expert Pack stats have an unexpected layer count")
    for layer in ACTIVE_MOE_LAYERS:
        if call_counts[layer] <= 0 or token_counts[layer] != expected_tokens:
            raise RuntimeError(
                f"layer {layer} routed {token_counts[layer]} tokens in "
                f"{call_counts[layer]} calls; expected {expected_tokens} tokens"
            )
    if int(stats.get("fallback_count", 0)) != 0:
        raise RuntimeError("the request used an expert fallback")
    if int(stats.get("io_errors", 0)) != 0:
        raise RuntimeError("the request encountered Expert Pack I/O errors")


def _git_sha(repo: Path) -> str | None:
    result = subprocess.run(
        ["git", "-C", str(repo), "rev-parse", "HEAD"],
        capture_output=True,
        text=True,
    )
    return result.stdout.strip() if result.returncode == 0 else None


def write_report(args: argparse.Namespace, gpu: str, result: dict[str, Any]) -> None:
    report = {
        "format": "SGLANG-DEEPSEEK-V4-FLASH-EXPERT-PACK-BENCHMARK-v1",
        "git_sha": _git_sha(args.sglang_repo),
        "gpu": gpu,
        "source_path": str(args.gguf),
        "result": result,
        "expert_pack_stats": read_stats(args.stats_path),
        "server_log": str(args.server_log),
    }
    temporary = args.report_path.with_suffix(args.report_path.suffix + ".tmp")
    temporary.write_text(
        json.dumps(report, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    temporary.replace(args.report_path)


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--gguf",
        type=Path,
        required=True,
        help="DeepSeek-V4-Flash source GGUF; server startup derives all artifacts",
    )
    parser.add_argument("--max-new-tokens", type=int, default=200)
    parser.set_defaults(
        prompt=DEFAULT_PROMPT,
        temperature=0.0,
        top_p=0.95,
        seed=20260810,
        host="127.0.0.1",
        port=30000,
        startup_timeout=1200,
        request_timeout=3600,
        watchdog_timeout=1800,
        context_length=32768,
        max_total_tokens=32768,
        mem_fraction_static=0.96,
        expert_cache_mib=21 * 1024,
        expert_cache_reserve_mib=2 * 1024,
        stage_slots=12,
        read_splits=4,
        direct_io=True,
    )
    args = parser.parse_args(argv)
    if args.max_new_tokens < 1:
        parser.error("--max-new-tokens must be positive")
    if not 1 <= args.port <= 65535:
        parser.error("--port must be between 1 and 65535")
    for name in (
        "expert_cache_mib",
        "expert_cache_reserve_mib",
        "stage_slots",
        "read_splits",
    ):
        if getattr(args, name) < 1:
            parser.error(f"--{name.replace('_', '-')} must be positive")
    args.gguf = args.gguf.expanduser().resolve(strict=True)
    args.sglang_repo = find_sglang_repo()
    args.artifact_dir = artifact_dir_for_source(args.gguf).resolve()
    args.server_log = args.artifact_dir / "deepseek-v4-5090-server.log"
    args.stats_path = args.artifact_dir / "deepseek-v4-expert-pack.stats.json"
    args.report_path = args.artifact_dir / "deepseek-v4-5090-benchmark.json"
    return args


def handle_termination(signum: int, _frame: object) -> None:
    raise KeyboardInterrupt(f"received signal {signum}")


def print_result(result: dict[str, Any], gpu: str) -> None:
    print(f"gpu: {gpu}")
    print(f"prompt_tokens: {result['prompt_tokens']}")
    print(f"completion_tokens: {result['completion_tokens']}")
    for name, suffix in (
        ("ttft_ms", ""),
        ("prefill_token_rate", " tok/s"),
        ("decode_token_rate", " tok/s"),
        ("tpot_ms", " ms/token"),
        ("end_to_end_token_rate", " tok/s"),
    ):
        value = result[name]
        print(f"{name}: {'n/a' if value is None else f'{value:.3f}'}{suffix}")


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    signal.signal(signal.SIGTERM, handle_termination)
    signal.signal(signal.SIGHUP, handle_termination)
    args.artifact_dir.mkdir(parents=True, exist_ok=True)
    lock = DEFAULT_LOCK.open("w")
    try:
        fcntl.flock(lock.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
    except BlockingIOError:
        print(
            f"error: another benchmark is running (lock: {DEFAULT_LOCK})",
            file=sys.stderr,
        )
        return 2
    lock.write(f"{os.getpid()}\n")
    lock.flush()
    try:
        gpu = detect_rtx_5090()
        print(f"MODEL_INPUT_READY gpu={gpu} gguf={args.gguf}", flush=True)
        result = run_benchmark(args)
        stats = read_stats(args.stats_path)
        audit_routes(stats, result["prompt_tokens"] + result["completion_tokens"])
        write_report(args, gpu, result)
        print_result(result, gpu)
        print(f"report: {args.report_path}")
        print(f"server_log: {args.server_log}")
        return 0
    except KeyboardInterrupt:
        print("error: benchmark interrupted", file=sys.stderr)
        return 130
    except Exception as exc:
        print(f"error: {exc}", file=sys.stderr)
        return 1
    finally:
        lock.close()


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