#!/usr/bin/env python3
"""Run the community-vLLM side of the frozen Qwen235B prefill grid."""

from __future__ import annotations

import argparse
import hashlib
import json
import os
import random
import subprocess
import time
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any


SCHEMA = "community-vllm-qwen235b-prefill-grid-v2"
MODEL = "/home/admin/cpfs/wjh/models/Qwen/Qwen3-235B-A22B-FP8"
SERVED_MODEL = "qwen3-235b-community-prefill"
TRACE = (
    "/home/admin/cpfs/wjh/aituner/aituner/trace_windows/traces/"
    "thinking_w20260327_1000.jsonl"
)
WINDOWS = "/home/admin/cpfs/wjh/aituner/aituner/trace_windows/windows.json"
TRACE_SHA256 = "f878e9af18f94dcfaced94a8e1e6b20a2f7d97d64aa862448025660dbbd965b2"
ORDER_SEED = 20260715
CLIENT_MAX_CONCURRENCY = 256


@dataclass(frozen=True)
class GridConfig:
    tp: int
    mns: int
    mbt: int
    expert_parallel: bool
    num_gpu_blocks: int

    @property
    def name(self) -> str:
        return f"tp{self.tp}_mns{self.mns}_mbt{self.mbt}"


GRID = tuple(
    GridConfig(
        tp=tp,
        mns=mns,
        mbt=mbt,
        expert_parallel=tp == 8,
        num_gpu_blocks=26101 if tp == 4 else 62351,
    )
    for tp in (4, 8)
    for mns in (64, 128)
    for mbt in (8192, 16384)
)


def sha256(path: Path) -> str:
    digest = hashlib.sha256()
    with path.open("rb") as source:
        for chunk in iter(lambda: source.read(1 << 20), b""):
            digest.update(chunk)
    return digest.hexdigest()


def write_json(path: Path, payload: Any) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    temporary = path.with_suffix(path.suffix + ".tmp")
    temporary.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
    os.replace(temporary, path)


def search_spec() -> dict[str, Any]:
    return {
        "low": 0.0,
        "high": 0.125,
        "tolerance": 0.001,
        "max_probes": 6,
        "sample_seed": 20260325,
        "inherit_incumbent_floor": False,
        "auto_high": {
            "enabled": False,
            "max_sampling_u": 1.0,
            "require_human_confirmation_beyond_trace": True,
        },
    }


def study_payload(
    *,
    tp: int,
    repo: Path,
    python: Path,
    vllm: Path,
    model: Path,
    trace: Path,
    windows: Path,
    port: int,
) -> dict[str, Any]:
    base_flags: dict[str, Any] = {
        "host": "127.0.0.1",
        "port": port,
        "served-model-name": SERVED_MODEL,
        "tensor-parallel-size": tp,
        "disable-custom-all-reduce": True,
        "quantization": "fp8",
        "gpu-memory-utilization": 0.80,
        "num-gpu-blocks-override": 26101 if tp == 4 else 62351,
        "kv-cache-dtype": "auto",
        "max-model-len": 40960,
        "max-num-batched-tokens": 8192,
        "max-num-seqs": 64,
        "enable-prefix-caching": False,
        "enable-chunked-prefill": True,
        "enforce-eager": True,
        "disable-log-requests": True,
    }
    if tp == 8:
        base_flags["enable-expert-parallel"] = True
    return {
        "study_id": f"community-qwen235b-prefill-tp{tp}-v2",
        "hardware": {
            "gpu_count": tp,
            "gpu_model": "NVIDIA H20",
            "host_candidates": ["dash0"],
        },
        "model": {"model_id": str(model), "served_model_name": SERVED_MODEL},
        "engine": {
            "engine_name": "vllm",
            "engine_version": "0.10.2-community",
            "exec_path": str(vllm),
            "cwd": str(repo),
            "host": "127.0.0.1",
            "port": port,
            "ready_timeout_s": 1800,
            "request_timeout_s": 1800,
            "healthcheck_path": "/v1/models",
            "launch_args": ["serve", str(model)],
            "base_envs": {
                "CUDA_VISIBLE_DEVICES": ",".join(str(index) for index in range(tp)),
            },
            "base_flags": base_flags,
            "tunable_envs": [],
            "tunable_flags": ["max-num-seqs", "max-num-batched-tokens"],
            "python_executable": str(python),
        },
        "trace": {
            "windows_path": str(windows),
            "window_id": "thinking_w20260327_1000",
            "trace_file_override": str(trace),
            "request_mode": "raw_completion",
            "completion_tokens_override": 1,
            "u_field": "sampling_u",
            "timestamp_field": "timestamp",
            # Keep the request generator above the largest server-side MNS so
            # MNS=128 is exercised by vLLM rather than clipped by the client.
            # The extra headroom also leaves an explicit engine waiting queue
            # under overload instead of moving that queue into the replay loop.
            "max_concurrency": CLIENT_MAX_CONCURRENCY,
            "input_length_filter": {
                "min_input_tokens": 0,
                "max_input_tokens": 32768,
            },
            "replay_time_scale": 1.0,
            "early_stop_max_lag_s": 180.0,
            "early_stop_max_elapsed_s": 1200.0,
            "restart_engine_after_early_stop": False,
            "adaptive_stop": {"enabled": False},
        },
        "slo": {
            "target_pass_rate": 0.95,
            "ttft_rule": {
                "kind": "step_ms",
                "buckets": [
                    {"max_input_tokens": 8191, "threshold_ms": 1000},
                    {"max_input_tokens": 32767, "threshold_ms": 2000},
                    {"threshold_ms": 2000},
                ],
            },
        },
        "search": search_spec(),
        "llm": {"use_harness": False},
    }


def git_fingerprint(repo: Path) -> dict[str, Any]:
    def run(*args: str) -> str:
        return subprocess.run(
            ["git", *args],
            cwd=repo,
            check=True,
            stdout=subprocess.PIPE,
            text=True,
        ).stdout.strip()

    return {
        "commit": run("rev-parse", "HEAD"),
        "status_porcelain": run("status", "--porcelain=v1").splitlines(),
    }


def prepare(args: argparse.Namespace) -> None:
    repo = args.repo.resolve()
    python = args.python.resolve()
    vllm = args.vllm.resolve()
    model = args.model.resolve()
    trace = args.trace.resolve()
    windows = args.windows.resolve()
    for path in (repo, python, vllm, model, trace, windows):
        if not path.exists():
            raise FileNotFoundError(path)
    actual_trace_sha = sha256(trace)
    if actual_trace_sha != args.expected_trace_sha256:
        raise ValueError(
            f"trace SHA256 mismatch: expected={args.expected_trace_sha256}, "
            f"actual={actual_trace_sha}"
        )
    output = args.output_root.resolve()
    output.mkdir(parents=True, exist_ok=True)
    studies: dict[int, dict[str, str]] = {}
    for tp, port in ((4, args.port), (8, args.port)):
        study_dir = output / "studies" / f"tp{tp}"
        study_path = study_dir / "study.json"
        pointer_path = study_dir / "study_spec.source"
        write_json(
            study_path,
            study_payload(
                tp=tp,
                repo=repo,
                python=python,
                vllm=vllm,
                model=model,
                trace=trace,
                windows=windows,
                port=port,
            ),
        )
        pointer_path.write_text(str(study_path) + "\n")
        studies[tp] = {
            "study_path": str(study_path),
            "study_sha256": sha256(study_path),
            "pointer_path": str(pointer_path),
        }

    order = list(GRID)
    random.Random(args.order_seed).shuffle(order)
    trials = []
    for index, config in enumerate(order, start=1):
        trial_dir = output / "trials" / config.name
        trial_path = trial_dir / "trial_spec.json"
        trial = {
            "study_id": f"community-qwen235b-prefill-tp{config.tp}-v2",
            "trial_id": config.name,
            "config_patch": {
                "env_patch": {},
                "flag_patch": {
                    "max-num-seqs": config.mns,
                    "max-num-batched-tokens": config.mbt,
                },
            },
            "search": search_spec(),
            "study_spec_path": studies[config.tp]["pointer_path"],
            "artifact_dir": str(trial_dir),
            "probe_log_path": str(trial_dir / "probe_history.json"),
            "engine_log_path": str(trial_dir / "engine.log"),
            "result_path": str(trial_dir / "result.json"),
            "search_evidence": {
                "enabled": False,
                "original_high": 0.125,
                "effective_high": 0.125,
                "trace_max_sampling_u": None,
                "max_sampling_u": 1.0,
                "require_human_confirmation_beyond_trace": True,
                "reason": "auto_high_disabled",
            },
        }
        write_json(trial_path, trial)
        trials.append(
            {
                "execution_index": index,
                "config": asdict(config) | {"name": config.name},
                "trial_spec_path": str(trial_path),
                "trial_spec_sha256": sha256(trial_path),
            }
        )
    manifest = {
        "schema": SCHEMA,
        "created_unix_s": time.time(),
        "repository": git_fingerprint(repo),
        "runtime": {"python": str(python), "vllm": str(vllm)},
        "model": {
            "path": str(model),
            "config_sha256": sha256(model / "config.json"),
        },
        "trace": {"path": str(trace), "sha256": actual_trace_sha},
        "contract": {
            "metric": "maximum SLO-feasible request_rate_per_gpu",
            "request_mode": "raw_completion",
            "completion_tokens_override": 1,
            "target_pass_rate": 0.95,
            "search": search_spec(),
            "prefix_caching": False,
            "chunked_prefill": True,
            "cuda_graphs": False,
            "custom_all_reduce": False,
            "client_max_concurrency": CLIENT_MAX_CONCURRENCY,
            "max_server_mns": max(config.mns for config in GRID),
            "kv_blocks_source": "community_vllm_measured_and_fixed_per_topology",
        },
        "order": {"method": "python_random_shuffle", "seed": args.order_seed},
        "studies": studies,
        "trials": trials,
    }
    write_json(output / "run_manifest.json", manifest)
    print(output / "run_manifest.json")


def run(args: argparse.Namespace) -> None:
    manifest = json.loads(args.manifest.read_text())
    if manifest.get("schema") != SCHEMA:
        raise ValueError(f"unexpected manifest schema: {manifest.get('schema')}")
    repo = args.repo.resolve()
    env = os.environ.copy()
    env["PYTHONPATH"] = str(repo / "src")
    for trial in manifest["trials"]:
        trial_path = Path(trial["trial_spec_path"])
        trial_dir = trial_path.parent
        result_path = trial_dir / "result.json"
        if result_path.is_file():
            result = json.loads(result_path.read_text())
            if result.get("status") == "completed":
                print(json.dumps({"config": trial["config"]["name"], "skipped": True}))
                continue
        command = [
            str(args.python.resolve()),
            "-m",
            "aituner.cli",
            "worker",
            "run-trial",
            "--trial-spec",
            str(trial_path),
        ]
        write_json(trial_dir / "worker_command.json", command)
        started = time.time()
        with (trial_dir / "worker.log").open("w", encoding="utf-8") as output:
            completed = subprocess.run(
                command,
                cwd=repo,
                env=env,
                stdout=output,
                stderr=subprocess.STDOUT,
                check=False,
            )
        record = {
            "config": trial["config"]["name"],
            "returncode": completed.returncode,
            "elapsed_seconds": time.time() - started,
        }
        print(json.dumps(record), flush=True)
        if completed.returncode != 0:
            raise RuntimeError(f"community trial failed: {record}")
    assemble(args.manifest)


def capacity_interval(probes: list[dict[str, Any]]) -> list[float]:
    low, high = 0.0, 0.125
    for probe in probes:
        midpoint = float(probe["threshold"])
        if probe["feasible"]:
            low = midpoint
        else:
            high = midpoint
    return [low, high]


def assemble(manifest_path: Path) -> Path:
    manifest = json.loads(manifest_path.read_text())
    records = []
    for trial in manifest["trials"]:
        result_path = Path(trial["trial_spec_path"]).parent / "result.json"
        if not result_path.is_file():
            raise FileNotFoundError(result_path)
        result = json.loads(result_path.read_text())
        if result.get("status") != "completed":
            raise ValueError(f"trial did not complete: {result_path}")
        tp = int(trial["config"]["tp"])
        request_rate = result.get("best_request_rate")
        records.append(
            {
                "config": trial["config"],
                "result_path": str(result_path),
                "result_sha256": sha256(result_path),
                "capacity_interval_sampling_u": capacity_interval(result["probes"]),
                "best_sampling_u": result.get("best_sampling_u"),
                "best_request_rate": request_rate,
                "best_request_rate_per_gpu": (
                    float(request_rate) / tp if request_rate is not None else None
                ),
                "best_pass_rate": result.get("best_pass_rate"),
            }
        )
    records.sort(
        key=lambda item: (
            -(item["best_request_rate_per_gpu"] or -1.0),
            item["config"]["name"],
        )
    )
    freeze = {
        "schema": SCHEMA,
        "run_manifest_sha256": sha256(manifest_path),
        "ranking": [dict(record, rank=index + 1) for index, record in enumerate(records)],
    }
    path = manifest_path.parent / "community_ranking_frozen.json"
    write_json(path, freeze)
    print(path)
    return path


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    prepare_parser = subparsers.add_parser("prepare")
    prepare_parser.add_argument("--repo", type=Path, required=True)
    prepare_parser.add_argument("--python", type=Path, required=True)
    prepare_parser.add_argument("--vllm", type=Path, required=True)
    prepare_parser.add_argument("--model", type=Path, default=Path(MODEL))
    prepare_parser.add_argument("--trace", type=Path, default=Path(TRACE))
    prepare_parser.add_argument("--windows", type=Path, default=Path(WINDOWS))
    prepare_parser.add_argument("--expected-trace-sha256", default=TRACE_SHA256)
    prepare_parser.add_argument("--output-root", type=Path, required=True)
    prepare_parser.add_argument("--order-seed", type=int, default=ORDER_SEED)
    prepare_parser.add_argument("--port", type=int, default=18918)
    run_parser = subparsers.add_parser("run")
    run_parser.add_argument("--repo", type=Path, required=True)
    run_parser.add_argument("--python", type=Path, required=True)
    run_parser.add_argument("--manifest", type=Path, required=True)
    assemble_parser = subparsers.add_parser("assemble")
    assemble_parser.add_argument("--manifest", type=Path, required=True)
    return parser.parse_args()


def main() -> None:
    args = parse_args()
    if args.command == "prepare":
        prepare(args)
    elif args.command == "run":
        run(args)
    elif args.command == "assemble":
        assemble(args.manifest)
    else:
        raise AssertionError(args.command)


if __name__ == "__main__":
    main()