#!/usr/bin/env python3
"""Run or summarize a stratified MiniMax-H3 FL2VA/Ref2VA serving workload."""

from __future__ import annotations

import argparse
import json
import math
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 build_plan(args: argparse.Namespace) -> list[dict[str, Any]]:
    resolutions = [int(item) for item in args.resolutions.split(",") if item.strip()]
    prompts = load_prompts(args.prompt_file, args.requests_per_resolution)
    plan: list[dict[str, Any]] = []
    # Interleave resolutions so any slow drift affects every bucket similarly.
    for prompt_index, prompt in enumerate(prompts):
        for short_edge in resolutions:
            plan.append(
                {
                    "request_id": f"{args.task}-r{short_edge}-p{prompt_index:02d}",
                    "task": args.task,
                    "short_edge": short_edge,
                    "prompt_index": prompt_index,
                    "prompt": prompt,
                    "seed": args.seed + prompt_index,
                }
            )
    return plan


def make_payload(args: argparse.Namespace, item: dict[str, Any], steps: int) -> dict[str, Any]:
    condition: dict[str, Any] = {
        "type": "image",
        "uri": str(args.reference_image),
        "role": "keyframe" if args.task == "fl2va" else "reference",
    }
    if args.task == "fl2va":
        condition["frame_index"] = 0
    return {
        "model": args.model,
        "prompt": item["prompt"],
        "num_outputs_per_prompt": 1,
        "num_inference_steps": steps,
        "flow_shift": args.flow_shift,
        "audio_flow_shift": args.audio_flow_shift,
        "seed": item["seed"],
        "task": args.task,
        "conditions": [condition],
        "target": {
            "short_edge": item["short_edge"],
            "aspect_ratio": args.aspect_ratio,
            "duration_seconds": args.duration_seconds,
        },
    }


def run_one(
    session: requests.Session,
    args: argparse.Namespace,
    item: dict[str, Any],
    steps: int,
) -> dict[str, Any]:
    started_epoch = time.time()
    started = time.monotonic()
    result: dict[str, Any] = {
        **item,
        "replica_index": args.replica_index,
        "port": args.port,
        "num_inference_steps": steps,
        "duration_seconds": args.duration_seconds,
        "aspect_ratio": args.aspect_ratio,
        "started_at_epoch": started_epoch,
        "success": False,
        "error": None,
    }
    try:
        response = session.post(
            f"http://{args.host}:{args.port}/v1/videos",
            json=make_payload(args, item, steps),
            timeout=args.submit_timeout,
        )
        if response.status_code != 200:
            raise RuntimeError(f"submit HTTP {response.status_code}: {response.text[:1000]}")
        status = response.json()
        video_id = status.get("id")
        if not video_id:
            raise RuntimeError(f"submit response has no id: {status}")
        result["video_id"] = video_id
        deadline = time.monotonic() + args.request_timeout
        while status.get("status") not in {"completed", "failed"}:
            if time.monotonic() >= deadline:
                raise TimeoutError(f"video job {video_id} exceeded {args.request_timeout}s")
            time.sleep(args.poll_interval)
            poll = session.get(
                f"http://{args.host}:{args.port}/v1/videos/{video_id}",
                timeout=args.poll_timeout,
            )
            if poll.status_code != 200:
                raise RuntimeError(f"poll HTTP {poll.status_code}: {poll.text[:1000]}")
            status = poll.json()
        if status.get("status") != "completed":
            raise RuntimeError(f"job failed: {status.get('error') or status}")
        result["success"] = True
        result["inference_time_s"] = status.get("inference_time_s")
        result["peak_memory_mb"] = status.get("peak_memory_mb")
        result["file_path"] = status.get("file_path")
    except Exception as exc:  # Keep the rest of the matrix running and record the cell failure.
        result["error"] = f"{type(exc).__name__}: {exc}"
    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)
    # Stratify by prompt index so every replica receives the same number of
    # samples from every resolution. This avoids assigning an entire slow
    # resolution bucket (for example 1080p) to only one replica.
    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)
    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}",
        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)
            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)
            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 = {
        "task": args.task,
        "tp": args.tp,
        "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.tp),
                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("--model", default="/data/hf_models/MiniMax-H3")
    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, default=Path.home() / ".cache/sglang/vbench_subject_consistency.txt")
    run.add_argument("--resolutions", default="480,720,768,1080")
    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("--num-inference-steps", type=int, default=20)
    run.add_argument("--warmup-requests", type=int, default=1)
    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("--flow-shift", type=float, default=12.0)
    run.add_argument("--audio-flow-shift", type=float, default=3.0)
    run.add_argument("--seed", type=int, default=1101)
    run.add_argument("--submit-timeout", type=float, default=120.0)
    run.add_argument("--poll-timeout", type=float, default=30.0)
    run.add_argument("--poll-interval", type=float, default=1.0)
    run.add_argument("--request-timeout", type=float, default=3600.0)
    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("--tp", 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())