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

B300 variant: adds per-instance in-flight concurrency (--in-flight) so the
server-side --batching-max-size can be exercised, and extends the summary
with topo/precision/batching columns. In-flight=1 reproduces the 6000D
serial-per-instance behaviour exactly.
"""

from __future__ import annotations

import argparse
import json
import math
import queue
import statistics
import threading
import time
from concurrent.futures import ThreadPoolExecutor
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} in_flight={args.in_flight}",
        flush=True,
    )
    failures = 0
    write_lock = threading.Lock()


    with requests.Session() as session:
        # Warmup stays serial so a slow first request cannot stall concurrency probes.
        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")
            with write_lock, args.output.open("a", encoding="utf-8") as _out:
                _out.write(json.dumps(warmup, ensure_ascii=False) + "\n")
        if not shard:
            return int(failures > 0)
        if args.in_flight <= 1:
            with args.output.open("a", encoding="utf-8", buffering=1) as output:
                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")
                    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)

        # Concurrent: fixed in-flight window over a worker pool.
        task_queue: queue.Queue[dict[str, Any] | None] = queue.Queue()
        for item in shard:
            task_queue.put(item)
        for _ in range(args.in_flight):
            task_queue.put(None)  # sentinel

        completed = 0

        def worker() -> None:
            nonlocal completed
            with args.output.open("a", encoding="utf-8", buffering=1) as output:
                while True:
                    item = task_queue.get()
                    if item is None:
                        task_queue.task_done()
                        return
                    result = run_one(session, args, item, args.num_inference_steps)
                    with write_lock:
                        output.write(json.dumps(result, ensure_ascii=False) + "\n")
                        completed += 1
                        failures += int(not result["success"])
                        print(
                            f"request {completed}/{len(shard)} id={item['request_id']} "
                            f"success={result['success']} latency={result['latency_s']:.2f}s "
                            f"error={result['error']}",
                            flush=True,
                        )
                    task_queue.task_done()

        with ThreadPoolExecutor(max_workers=args.in_flight) as pool:
            futures = [pool.submit(worker) for _ in range(args.in_flight)]
            for future in futures:
                future.result()
    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 = {
        "topo": args.topo,
        "prec": args.prec,
        "replicas": args.replicas,
        "tp": args.tp,
        "ulysses": args.ulysses,
        "batching": args.batching,
        "in_flight": args.in_flight,
        "task": args.task,
        "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.topo),
                str(args.prec),
                str(args.replicas),
                str(args.tp),
                str(args.ulysses),
                str(args.batching),
                str(args.in_flight),
                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("--in-flight", type=int, default=1)
    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)
    run.set_defaults(func=run_command)

    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("--topo", required=True)
    summarize.add_argument("--prec", default="bf16")
    summarize.add_argument("--tp", type=int, required=True)
    summarize.add_argument("--ulysses", type=int, required=True)
    summarize.add_argument("--replicas", type=int, required=True)
    summarize.add_argument("--batching", type=int, required=True)
    summarize.add_argument("--in-flight", type=int, required=True)
    summarize.add_argument("--expected-requests", type=int, required=True)
    summarize.set_defaults(func=summarize_command)
    return parser.parse_args()


def main() -> int:
    args = parse_args()
    return args.func(args)


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