#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
"""Run a one-shot Kimi K3 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

SCRIPT_DIR = Path(__file__).resolve().parent
DEFAULT_PROMPT = "请介绍深圳"
DEFAULT_SERVER_LOG = SCRIPT_DIR / "logs/kimi-k3-5090-benchmark-server.log"
DEFAULT_LOCK = "/tmp/sglang-kimi-k3-5090-benchmark.lock"
METADATA_FORMAT_VERSION = 3
ACTIVE_MOE_LAYERS = tuple(range(1, 93))
IMMUTABLE_TOP_K = 16


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


def artifact_dir_for_source(gguf: Path) -> Path:
    stat = gguf.stat()
    fingerprint = hashlib.sha256(
        f"{gguf.parent.resolve()}:{stat.st_size}:{stat.st_mtime_ns}:"
        f"{METADATA_FORMAT_VERSION}".encode()
    ).hexdigest()[:20]
    return cache_root() / "sglang-expert-pack" / "kimi-k3" / 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 write_json_atomic(path: Path, value: object) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_name(path.name + f".{os.getpid()}.tmp")
    temporary.write_text(
        json.dumps(value, ensure_ascii=False, indent=2, sort_keys=True) + "\n",
        encoding="utf-8",
    )
    temporary.replace(path)


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)}"
        )
    return gpu_zero


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 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=(",", ":")),
        "--tp-size",
        "1",
        "--ep-size",
        "1",
        "--disable-radix-cache",
        "--mamba-radix-cache-strategy",
        "no_buffer",
        "--disable-overlap-schedule",
        "--skip-server-warmup",
        "--context-length",
        str(args.context_length),
        "--max-total-tokens",
        str(args.max_total_tokens),
        "--chunked-prefill-size",
        str(args.chunked_prefill_size),
        "--watchdog-timeout",
        str(args.watchdog_timeout),
        "--max-running-requests",
        "1",
        "--mem-fraction-static",
        str(args.mem_fraction_static),
        "--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)
    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,
        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, server_url(args))
        raise


def stop_server(process: subprocess.Popen | None, url: str) -> 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(*server_address(url), timeout=0.2)
            and time.monotonic() < deadline
        ):
            time.sleep(0.2)
        print(f"SERVICE_STOPPED pid={process.pid} url={url}", flush=True)
    finally:
        if log is not None:
            log.close()


def server_address(url: str) -> tuple[str, int]:
    without_scheme = url.removeprefix("http://")
    host, port = without_scheme.rsplit(":", 1)
    return host, int(port)


def make_prompt(model_dir: Path, prompt: str) -> tuple[list[int], str]:
    from transformers import AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(
        model_dir, trust_remote_code=True, local_files_only=True
    )
    input_ids = tokenizer.apply_chat_template(
        [{"role": "user", "content": prompt}],
        tokenize=True,
        add_generation_prompt=True,
    )
    prompt_text = tokenizer.decode(input_ids, skip_special_tokens=False)
    if not input_ids:
        raise ValueError("Kimi chat template produced an empty prompt")
    return [int(value) for value in input_ids], prompt_text


def generate(
    url: str,
    input_ids: list[int],
    prompt: str,
    max_new_tokens: int,
    temperature: float,
    top_p: float,
    seed: int,
    *,
    stream_output: bool,
) -> dict:
    payload = {
        "input_ids": input_ids,
        "sampling_params": {
            "temperature": temperature,
            "top_p": top_p,
            "sampling_seed": seed,
            "max_new_tokens": max_new_tokens,
            "ignore_eos": True,
        },
        "return_logprob": True,
        "stream": True,
    }
    request = urllib.request.Request(
        url.rstrip("/") + "/generate",
        data=json.dumps(payload).encode("utf-8"),
        headers={"Content-Type": "application/json"},
        method="POST",
    )
    started = time.perf_counter_ns()
    first_token = None
    last_token = None
    completion_tokens = 0
    prompt_tokens = None
    output = ""
    finish_reason = None
    output_token_ids: dict[int, int] = {}

    if stream_output:
        print(f"prompt: {prompt}", flush=True)
        print("output: ", end="", flush=True)
    with urllib.request.urlopen(request, timeout=3600) 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"])
            logprobs = meta.get("output_token_logprobs") or []
            logprob_length = int(
                meta.get("output_token_logprobs_length", current_tokens)
            )
            offset = logprob_length - len(logprobs)
            for index, item in enumerate(logprobs):
                output_token_ids[offset + index] = int(item[1])
            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 omitted token timing metadata")
    if prompt_tokens != len(input_ids):
        raise RuntimeError(
            f"server prompt token count {prompt_tokens} != tokenizer count {len(input_ids)}"
        )
    ordered_token_ids = [output_token_ids[index] for index in sorted(output_token_ids)]
    if len(ordered_token_ids) != completion_tokens:
        raise RuntimeError(
            "SGLang response omitted output token IDs: "
            f"{len(ordered_token_ids)} != {completion_tokens}"
        )
    ttft_s = (first_token - started) / 1e9
    decode_span_s = (last_token - first_token) / 1e9
    total_s = (time.perf_counter_ns() - started) / 1e9
    return {
        "output": output,
        "output_token_ids": ordered_token_ids,
        "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": (
            (completion_tokens - 1) / decode_span_s
            if completion_tokens > 1 and decode_span_s > 0
            else None
        ),
        "tpot_ms": (
            decode_span_s * 1000 / (completion_tokens - 1)
            if completion_tokens > 1
            else None
        ),
        "total_elapsed_s": total_s,
        "end_to_end_token_rate": completion_tokens / total_s if total_s > 0 else None,
    }


def read_stats(path: Path) -> dict:
    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, expected_tokens: int) -> None:
    token_counts = stats["route_tokens_by_layer"]
    call_counts = stats["route_calls_by_layer"]
    for layer in ACTIVE_MOE_LAYERS:
        if (
            token_counts[layer] != expected_tokens
            or call_counts[layer] != expected_tokens
        ):
            raise RuntimeError(
                f"layer {layer} routed {token_counts[layer]} tokens in "
                f"{call_counts[layer]} calls; expected {expected_tokens} exact Top-16 calls"
            )
    if any(
        token_counts[layer]
        for layer in set(range(len(token_counts))) - set(ACTIVE_MOE_LAYERS)
    ):
        raise RuntimeError("routed experts outside model layers 1..92")
    if int(stats.get("fallback_count", 0)) != 0:
        raise RuntimeError("the request used an expert fallback")


def git_commit(repo: Path) -> str:
    return subprocess.run(
        ["git", "rev-parse", "HEAD"],
        cwd=repo,
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--gguf",
        type=Path,
        required=True,
        help="one Kimi-K3 GGUF shard; sibling shards and serving artifacts are derived",
    )
    parser.add_argument("--prompt", default=DEFAULT_PROMPT)
    parser.add_argument("--max-new-tokens", type=int, default=200)
    parser.set_defaults(
        host="127.0.0.1",
        port=30001,
        temperature=0.0,
        top_p=0.95,
        seed=20260813,
        startup_timeout=1200,
        context_length=384,
        max_total_tokens=512,
        chunked_prefill_size=64,
        watchdog_timeout=1800,
        mem_fraction_static=0.98,
        expert_cache_mib=5120,
        expert_cache_reserve_mib=1536,
        stage_slots=16,
        read_splits=1,
        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 / DEFAULT_SERVER_LOG.name
    args.stats_path = args.artifact_dir / "kimi-k3-expert-pack.stats.json"
    args.report_path = args.artifact_dir / "kimi-k3-5090-benchmark.json"
    return args


def print_result(result: dict, stats: dict) -> None:
    def value(name: str) -> str:
        item = result[name]
        return "n/a" if item is None else f"{item:.3f}"

    print(f"prompt_tokens: {result['prompt_tokens']}")
    print(f"completion_tokens: {result['completion_tokens']}")
    print(f"ttft_ms: {value('ttft_ms')}")
    print(f"prefill_token_rate: {value('prefill_token_rate')} tok/s")
    print(f"decode_token_rate: {value('decode_token_rate')} tok/s")
    print(f"tpot_ms: {value('tpot_ms')} ms/token")
    print(f"end_to_end_token_rate: {value('end_to_end_token_rate')} tok/s")
    print(
        f"expert_cache: hits={stats['cache_hits']} misses={stats['cache_misses']} "
        f"evictions={stats['cache_evictions']} reads={stats['pack_reads']} "
        f"read_bytes={stats['pack_read_bytes']} h2d_bytes={stats['h2d_bytes']}"
    )


def main() -> int:
    args = parse_args()
    signal.signal(signal.SIGTERM, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt()))
    signal.signal(signal.SIGHUP, lambda *_: (_ for _ in ()).throw(KeyboardInterrupt()))
    args.artifact_dir.mkdir(parents=True, exist_ok=True)
    lock = Path(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()

    process = None
    try:
        gpu = detect_rtx_5090()
        prepared = {
            "gpu": gpu,
            "source_gguf": str(args.gguf),
            "top_k": IMMUTABLE_TOP_K,
        }
        print(
            f"MODEL_INPUT_READY gpu={gpu} top_k={IMMUTABLE_TOP_K} gguf={args.gguf}",
            flush=True,
        )

        if args.stats_path.exists():
            args.stats_path.unlink()
        process = start_server(args)
        prompt_ids, prompt_text = make_prompt(
            args.artifact_dir / "model-meta", args.prompt
        )
        result = generate(
            server_url(args),
            prompt_ids,
            args.prompt,
            args.max_new_tokens,
            args.temperature,
            args.top_p,
            args.seed,
            stream_output=True,
        )
        model_tokens = result["prompt_tokens"] + result["completion_tokens"] - 1
        stop_server(process, server_url(args))
        process = None
        stats = read_stats(args.stats_path)
        audit_routes(stats, model_tokens)

        report = {
            **prepared,
            "status": "passed",
            "sglang_commit": git_commit(args.sglang_repo),
            "python": sys.version,
            "command": sys.argv,
            "server_url": server_url(args),
            "server_log": str(args.server_log),
            "stats_path": str(args.stats_path),
            "prompt": args.prompt,
            "formatted_prompt": prompt_text,
            "result": result,
            "expert_pack_stats": stats,
            "route_audit": {
                "active_moe_layers": list(ACTIVE_MOE_LAYERS),
                "immutable_top_k": IMMUTABLE_TOP_K,
                "model_tokens_per_layer": model_tokens,
                "fallback_count": int(stats.get("fallback_count", 0)),
            },
        }
        write_json_atomic(args.report_path, report)
        print_result(result, stats)
        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:
        stop_server(process, server_url(args))


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