#!/usr/bin/env python3
"""Balanced synchronous serving benchmark for MiniMax-H3 on vLLM-Omni."""

from __future__ import annotations

import argparse
import json
import math
import mimetypes
import statistics
import time
from pathlib import Path
from typing import Any

import requests


DEFAULT_PROMPT = "A cinematic landscape with natural motion and realistic lighting."


def percentile(values: list[float], q: float) -> float:
    if not values:
        return 0.0
    values = sorted(values)
    pos = (len(values) - 1) * q
    lo, hi = math.floor(pos), math.ceil(pos)
    if lo == hi:
        return values[lo]
    return values[lo] * (hi - pos) + values[hi] * (pos - lo)


def load_prompts(path: Path, count: int) -> list[str]:
    prompts: list[str] = []
    if path.is_file():
        prompts = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    if not prompts:
        prompts = [DEFAULT_PROMPT]
    repeats = (count + len(prompts) - 1) // len(prompts)
    return (prompts * repeats)[:count]


def parse_resolution_map(raw: str) -> dict[int, tuple[int, int]]:
    mapping: dict[int, tuple[int, int]] = {}
    for entry in raw.split(","):
        short_edge, shape = entry.split(":", 1)
        width, height = shape.lower().split("x", 1)
        mapping[int(short_edge)] = (int(width), int(height))
    return mapping


def build_plan(args: argparse.Namespace) -> list[dict[str, Any]]:
    resolution_map = parse_resolution_map(args.resolution_map)
    prompts = load_prompts(args.prompt_file, args.requests_per_resolution)
    plan: list[dict[str, Any]] = []
    for prompt_index, prompt in enumerate(prompts):
        for short_edge, (width, height) in resolution_map.items():
            plan.append(
                {
                    "request_id": f"{args.task}-r{short_edge}-p{prompt_index:02d}",
                    "task": args.task,
                    "short_edge": short_edge,
                    "width": width,
                    "height": height,
                    "prompt_index": prompt_index,
                    "prompt": prompt,
                    "seed": args.seed + prompt_index,
                }
            )
    return plan


def make_form(args: argparse.Namespace, item: dict[str, Any], steps: int) -> dict[str, str]:
    extra_params = {
        "task": args.task,
        "duration": args.duration_seconds,
        "audio_flow_shift": args.audio_flow_shift,
    }
    return {
        "prompt": item["prompt"],
        "width": str(item["width"]),
        "height": str(item["height"]),
        "aspect_ratio": args.aspect_ratio,
        "fps": str(args.fps),
        "num_inference_steps": str(steps),
        "flow_shift": str(args.flow_shift),
        "seed": str(item["seed"]),
        "quality": args.quality,
        "extra_params": json.dumps(extra_params, separators=(",", ":")),
    }


def run_one(
    session: requests.Session,
    args: argparse.Namespace,
    item: dict[str, Any],
    steps: int,
    save_video: bool,
) -> dict[str, Any]:
    started_epoch = time.time()
    started = time.monotonic()
    output_path = args.video_dir / f"{item['request_id']}.mp4"
    result: dict[str, Any] = {
        **item,
        "backend": "vllm-omni",
        "replica_index": args.replica_index,
        "port": args.port,
        "gpus_per_instance": args.gpus_per_instance,
        "dit_tensor_parallel_size": args.dit_tp,
        "ulysses_degree": args.usp,
        "num_inference_steps": steps,
        "duration_seconds": args.duration_seconds,
        "aspect_ratio": args.aspect_ratio,
        "started_at_epoch": started_epoch,
        "success": False,
        "error": None,
    }
    try:
        mime_type = mimetypes.guess_type(args.reference_image.name)[0] or "application/octet-stream"
        with args.reference_image.open("rb") as reference_handle:
            files = {
                "input_reference": (
                    args.reference_image.name,
                    reference_handle,
                    mime_type,
                )
            }
            with session.post(
                f"http://{args.host}:{args.port}/v1/videos/sync",
                data=make_form(args, item, steps),
                files=files,
                timeout=(args.connect_timeout, args.request_timeout),
                stream=True,
            ) as response:
                if response.status_code != 200:
                    body = response.content[:2000].decode("utf-8", errors="replace")
                    raise RuntimeError(f"HTTP {response.status_code}: {body}")
                result["content_type"] = response.headers.get("content-type")
                result["stage_durations"] = response.headers.get("x-stage-durations")
                result["peak_memory_mb"] = response.headers.get("x-peak-memory-mb")
                byte_count = 0
                first_bytes = b""
                output_handle = output_path.open("wb") if save_video else None
                try:
                    for chunk in response.iter_content(chunk_size=1024 * 1024):
                        if not chunk:
                            continue
                        if len(first_bytes) < 32:
                            first_bytes += chunk[: 32 - len(first_bytes)]
                        byte_count += len(chunk)
                        if output_handle is not None:
                            output_handle.write(chunk)
                finally:
                    if output_handle is not None:
                        output_handle.close()
                if byte_count < 12 or b"ftyp" not in first_bytes:
                    raise RuntimeError(
                        f"response is not a valid-looking MP4: bytes={byte_count}, "
                        f"content_type={result['content_type']}"
                    )
                result["response_bytes"] = byte_count
                if save_video:
                    result["file_path"] = str(output_path)
                result["success"] = True
    except Exception as exc:
        result["error"] = f"{type(exc).__name__}: {exc}"
        if save_video and output_path.exists():
            output_path.unlink()
    result["latency_s"] = time.monotonic() - started
    result["finished_at_epoch"] = time.time()
    return result


def run_command(args: argparse.Namespace) -> int:
    if not args.reference_image.is_file():
        raise SystemExit(f"reference image not found: {args.reference_image}")
    full_plan = build_plan(args)
    # Stratified allocation: every replica gets the same number of prompts
    # from every resolution bucket.
    shard = [
        item
        for item in full_plan
        if item["prompt_index"] % args.num_replicas == args.replica_index
    ]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.video_dir.mkdir(parents=True, exist_ok=True)
    completed_ids: set[str] = set()
    if args.output.is_file():
        for line in args.output.read_text(encoding="utf-8").splitlines():
            try:
                completed_ids.add(json.loads(line)["request_id"])
            except (json.JSONDecodeError, KeyError):
                continue
    shard = [item for item in shard if item["request_id"] not in completed_ids]
    print(
        f"task={args.task} replica={args.replica_index}/{args.num_replicas} "
        f"requests={len(shard)} port={args.port} gpus={args.gpus_per_instance} "
        f"dit_tp={args.dit_tp} usp={args.usp}",
        flush=True,
    )
    failures = 0
    with requests.Session() as session, args.output.open("a", encoding="utf-8", buffering=1) as output:
        for warmup_index in range(args.warmup_requests):
            warmup_item = (shard or full_plan)[warmup_index % len(shard or full_plan)].copy()
            warmup_item["request_id"] = f"warmup-{warmup_index}-{warmup_item['request_id']}"
            warmup = run_one(
                session,
                args,
                warmup_item,
                args.warmup_inference_steps,
                save_video=False,
            )
            print(
                f"warmup {warmup_index + 1}/{args.warmup_requests}: "
                f"success={warmup['success']} latency={warmup['latency_s']:.2f}s "
                f"error={warmup['error']}",
                flush=True,
            )
            if not warmup["success"]:
                raise SystemExit("warmup failed")
        for index, item in enumerate(shard, start=1):
            result = run_one(session, args, item, args.num_inference_steps, save_video=True)
            output.write(json.dumps(result, ensure_ascii=False) + "\n")
            output.flush()
            failures += int(not result["success"])
            print(
                f"request {index}/{len(shard)} id={item['request_id']} "
                f"success={result['success']} latency={result['latency_s']:.2f}s "
                f"error={result['error']}",
                flush=True,
            )
    return int(failures > 0)


def summarize_command(args: argparse.Namespace) -> int:
    rows: list[dict[str, Any]] = []
    for path in sorted(args.input_dir.glob("client_*/results.jsonl")):
        for line in path.read_text(encoding="utf-8").splitlines():
            if line.strip():
                rows.append(json.loads(line))
    successful = [row for row in rows if row.get("success")]
    latencies = [float(row["latency_s"]) for row in successful]
    started = [float(row["started_at_epoch"]) for row in rows]
    finished = [float(row["finished_at_epoch"]) for row in rows]
    wall_s = max(finished) - min(started) if started and finished else 0.0
    buckets: dict[str, dict[str, Any]] = {}
    for short_edge in sorted({int(row["short_edge"]) for row in rows}):
        bucket_rows = [row for row in rows if int(row["short_edge"]) == short_edge]
        bucket_success = [row for row in bucket_rows if row.get("success")]
        bucket_latencies = [float(row["latency_s"]) for row in bucket_success]
        buckets[str(short_edge)] = {
            "requests": len(bucket_rows),
            "completed": len(bucket_success),
            "failed": len(bucket_rows) - len(bucket_success),
            "latency_mean_s": statistics.fmean(bucket_latencies) if bucket_latencies else 0.0,
            "latency_p95_s": percentile(bucket_latencies, 0.95),
        }
    summary = {
        "backend": "vllm-omni",
        "task": args.task,
        "gpus_per_instance": args.gpus_per_instance,
        "dit_tensor_parallel_size": args.dit_tp,
        "ulysses_degree": args.usp,
        "replicas": args.replicas,
        "expected_requests": args.expected_requests,
        "requests_recorded": len(rows),
        "completed": len(successful),
        "failed": len(rows) - len(successful),
        "machine_wall_s": wall_s,
        "machine_qps": len(successful) / wall_s if wall_s else 0.0,
        "latency_mean_s": statistics.fmean(latencies) if latencies else 0.0,
        "latency_p50_s": percentile(latencies, 0.50),
        "latency_p95_s": percentile(latencies, 0.95),
        "by_short_edge": buckets,
    }
    args.output.parent.mkdir(parents=True, exist_ok=True)
    args.output.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(
        "\t".join(
            [
                str(args.gpus_per_instance),
                str(args.dit_tp),
                str(args.usp),
                str(args.replicas),
                args.task,
                str(args.expected_requests),
                str(len(rows)),
                str(len(successful)),
                str(len(rows) - len(successful)),
                f"{summary['machine_qps']:.8f}",
                f"{summary['latency_mean_s']:.6f}",
                f"{summary['latency_p95_s']:.6f}",
                f"{wall_s:.3f}",
            ]
        )
    )
    return int(len(rows) != args.expected_requests or len(successful) != len(rows))


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser()
    subparsers = parser.add_subparsers(dest="command", required=True)
    run = subparsers.add_parser("run")
    run.add_argument("--host", default="127.0.0.1")
    run.add_argument("--port", type=int, required=True)
    run.add_argument("--task", choices=["fl2va", "ref2va"], required=True)
    run.add_argument("--reference-image", type=Path, required=True)
    run.add_argument("--prompt-file", type=Path, required=True)
    run.add_argument("--resolution-map", default="480:864x480,720:1280x736,768:1344x768,1080:1920x1088")
    run.add_argument("--requests-per-resolution", type=int, default=8)
    run.add_argument("--replica-index", type=int, required=True)
    run.add_argument("--num-replicas", type=int, required=True)
    run.add_argument("--gpus-per-instance", type=int, required=True)
    run.add_argument("--dit-tp", type=int, required=True)
    run.add_argument("--usp", type=int, required=True)
    run.add_argument("--num-inference-steps", type=int, default=20)
    run.add_argument("--warmup-requests", type=int, default=2)
    run.add_argument("--warmup-inference-steps", type=int, default=5)
    run.add_argument("--duration-seconds", type=float, default=5.0)
    run.add_argument("--aspect-ratio", default="16:9")
    run.add_argument("--fps", type=int, default=24)
    run.add_argument("--flow-shift", type=float, default=12.0)
    run.add_argument("--audio-flow-shift", type=float, default=3.0)
    run.add_argument("--quality", default="lossless")
    run.add_argument("--seed", type=int, default=1101)
    run.add_argument("--connect-timeout", type=float, default=30.0)
    run.add_argument("--request-timeout", type=float, default=14400.0)
    run.add_argument("--video-dir", type=Path, required=True)
    run.add_argument("--output", type=Path, required=True)

    summarize = subparsers.add_parser("summarize")
    summarize.add_argument("--input-dir", type=Path, required=True)
    summarize.add_argument("--output", type=Path, required=True)
    summarize.add_argument("--task", required=True)
    summarize.add_argument("--gpus-per-instance", type=int, required=True)
    summarize.add_argument("--dit-tp", type=int, required=True)
    summarize.add_argument("--usp", type=int, required=True)
    summarize.add_argument("--replicas", type=int, required=True)
    summarize.add_argument("--expected-requests", type=int, required=True)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    return run_command(args) if args.command == "run" else summarize_command(args)


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