#!/usr/bin/env python3
"""Run the eight Feishu MiniMax-H3 Ref2VA cases against SGLang servers."""

from __future__ import annotations

import argparse
import hashlib
import json
import re
import statistics
import time
from pathlib import Path
from typing import Any

import requests


REFERENCE_RE = re.compile(
    r"(图|音色)([一二三四五六七八九十\d]+)\s*=\s*([^（(，,\n]+)"
)


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


def load_cases(records_path: Path, assets_root: Path) -> list[dict[str, Any]]:
    envelope = json.loads(records_path.read_text(encoding="utf-8"))
    payload = envelope["data"]
    rows = payload["data"]
    record_ids = payload["record_id_list"]
    fields = payload["fields"]
    if len(rows) != 8 or len(record_ids) != 8:
        raise ValueError(f"expected exactly 8 records, got {len(rows)} rows")
    field_index = {name: index for index, name in enumerate(fields)}
    required_fields = {"提示词", "参考音频", "参考图像"}
    if not required_fields.issubset(field_index):
        raise ValueError(f"missing fields: {sorted(required_fields - set(field_index))}")

    cases: list[dict[str, Any]] = []
    for zero_index, (record_id, row) in enumerate(zip(record_ids, rows, strict=True)):
        case_number = zero_index + 1
        case_dir = assets_root / f"case{case_number:02d}_{record_id}"
        prompt = row[field_index["提示词"]]
        audio_attachments = row[field_index["参考音频"]] or []
        image_attachments = row[field_index["参考图像"]] or []
        attachment_by_name: dict[str, tuple[str, dict[str, Any]]] = {}
        for kind, attachments in (("audio", audio_attachments), ("image", image_attachments)):
            for attachment in attachments:
                name = str(attachment["name"]).strip()
                if name in attachment_by_name:
                    raise ValueError(f"case {case_number}: duplicate attachment name {name!r}")
                attachment_by_name[name] = (kind, attachment)

        references = []
        seen_names: set[str] = set()
        for match in REFERENCE_RE.finditer(prompt.splitlines()[0]):
            label_kind, ordinal, name = match.groups()
            name = name.strip()
            expected_kind = "image" if label_kind == "图" else "audio"
            if name not in attachment_by_name:
                raise ValueError(
                    f"case {case_number}: mapped {label_kind}{ordinal} file {name!r} "
                    f"not present in attachments {sorted(attachment_by_name)}"
                )
            actual_kind, attachment = attachment_by_name[name]
            if actual_kind != expected_kind:
                raise ValueError(
                    f"case {case_number}: {name!r} is {actual_kind}, expected {expected_kind}"
                )
            path = case_dir / name
            if not path.is_file():
                raise ValueError(f"case {case_number}: missing downloaded file {path}")
            references.append(
                {
                    "label": f"{label_kind}{ordinal}",
                    "type": expected_kind,
                    "name": name,
                    "path": str(path),
                    "file_token": attachment["file_token"],
                    "size": path.stat().st_size,
                    "sha256": sha256(path),
                }
            )
            seen_names.add(name)
        if not references:
            raise ValueError(f"case {case_number}: no reference mapping found in first line")
        if seen_names != set(attachment_by_name):
            raise ValueError(
                f"case {case_number}: unmapped attachments "
                f"{sorted(set(attachment_by_name) - seen_names)}"
            )
        cases.append(
            {
                "case_number": case_number,
                "case_id": f"case{case_number:02d}_{record_id}",
                "record_id": record_id,
                "prompt": prompt,
                "seed": 1100 + case_number,
                "references": references,
            }
        )
    return cases


def make_payload(args: argparse.Namespace, case: dict[str, Any], steps: int) -> dict[str, Any]:
    conditions = [
        {"type": ref["type"], "uri": ref["path"], "role": "reference"}
        for ref in case["references"]
    ]
    return {
        "model": args.model,
        "prompt": case["prompt"],
        "num_outputs_per_prompt": 1,
        "num_inference_steps": steps,
        "flow_shift": args.flow_shift,
        "audio_flow_shift": args.audio_flow_shift,
        "seed": case["seed"],
        "task": "ref2va",
        "conditions": conditions,
        "target": {
            "short_edge": args.short_edge,
            "aspect_ratio": args.aspect_ratio,
            "duration_seconds": args.duration_seconds,
        },
    }


def request_video(
    session: requests.Session,
    args: argparse.Namespace,
    case: dict[str, Any],
    steps: int,
) -> dict[str, Any]:
    started_epoch = time.time()
    started = time.monotonic()
    result: dict[str, Any] = {
        "case_number": case["case_number"],
        "case_id": case["case_id"],
        "record_id": case["record_id"],
        "seed": case["seed"],
        "port": args.port,
        "replica_index": args.replica_index,
        "num_inference_steps": steps,
        "short_edge": args.short_edge,
        "aspect_ratio": args.aspect_ratio,
        "duration_seconds": args.duration_seconds,
        "references": case["references"],
        "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, case, steps),
            timeout=args.submit_timeout,
        )
        if response.status_code != 200:
            raise RuntimeError(f"submit HTTP {response.status_code}: {response.text[:2000]}")
        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[:2000]}")
            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")
        result["status"] = status
    except Exception as exc:
        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:
    cases = load_cases(args.records, args.assets_root)
    case = cases[args.case_number - 1]
    args.output.parent.mkdir(parents=True, exist_ok=True)
    with requests.Session() as session:
        if args.warmup_steps > 0:
            warmup = request_video(session, args, case, args.warmup_steps)
            print(
                f"warmup case={case['case_id']} success={warmup['success']} "
                f"latency={warmup['latency_s']:.2f}s error={warmup['error']}",
                flush=True,
            )
            if not warmup["success"]:
                args.output.write_text(
                    json.dumps({"warmup": warmup}, ensure_ascii=False, indent=2),
                    encoding="utf-8",
                )
                return 1
        result = request_video(session, args, case, args.num_inference_steps)
    args.output.write_text(json.dumps(result, ensure_ascii=False, indent=2), encoding="utf-8")
    print(
        f"request case={case['case_id']} success={result['success']} "
        f"latency={result['latency_s']:.2f}s file={result.get('file_path')} "
        f"error={result['error']}",
        flush=True,
    )
    return int(not result["success"])


def summarize_command(args: argparse.Namespace) -> int:
    rows = []
    for path in sorted(args.result_root.glob("batch_*/client_*/result.json")):
        row = json.loads(path.read_text(encoding="utf-8"))
        if "warmup" not in row:
            rows.append(row)
    batches: dict[str, Any] = {}
    for batch_dir in sorted(args.result_root.glob("batch_*")):
        batch_rows = []
        for path in sorted(batch_dir.glob("client_*/result.json")):
            row = json.loads(path.read_text(encoding="utf-8"))
            if "warmup" not in row:
                batch_rows.append(row)
        successful = [row for row in batch_rows if row.get("success")]
        starts = [float(row["started_at_epoch"]) for row in batch_rows]
        finishes = [float(row["finished_at_epoch"]) for row in batch_rows]
        wall_s = max(finishes) - min(starts) if starts and finishes else 0.0
        batches[batch_dir.name] = {
            "requests": len(batch_rows),
            "completed": len(successful),
            "failed": len(batch_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(
                float(row["latency_s"]) for row in successful
            ) if successful else 0.0,
        }
    successful = [row for row in rows if row.get("success")]
    measured_wall_s = sum(batch["machine_wall_s"] for batch in batches.values())
    summary = {
        "method": "base",
        "task": "ref2va",
        "tp": 2,
        "replicas": 4,
        "expected_requests": 8,
        "requests_recorded": len(rows),
        "completed": len(successful),
        "failed": len(rows) - len(successful),
        "measured_batch_wall_s": measured_wall_s,
        "machine_qps": len(successful) / measured_wall_s if measured_wall_s else 0.0,
        "latency_mean_s": statistics.fmean(
            float(row["latency_s"]) for row in successful
        ) if successful else 0.0,
        "batches": batches,
        "cases": rows,
    }
    args.output.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    per_case_path = args.output.with_name("per_case.tsv")
    with per_case_path.open("w", encoding="utf-8") as output:
        output.write(
            "case_number\tcase_id\trecord_id\tseed\tsuccess\t"
            "end_to_end_latency_s\tserver_inference_time_s\tfile_path\terror\n"
        )
        for row in sorted(rows, key=lambda item: int(item["case_number"])):
            fields = [
                row.get("case_number"),
                row.get("case_id"),
                row.get("record_id"),
                row.get("seed"),
                row.get("success"),
                f"{float(row.get('latency_s') or 0.0):.6f}",
                f"{float(row.get('inference_time_s') or 0.0):.6f}",
                row.get("file_path") or "",
                str(row.get("error") or "").replace("\t", " ").replace("\n", " "),
            ]
            output.write("\t".join(str(field) for field in fields) + "\n")
    print(json.dumps({key: value for key, value in summary.items() if key != "cases"}, ensure_ascii=False, indent=2))
    return int(len(rows) != 8 or len(successful) != 8)


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("--replica-index", type=int, required=True)
    run.add_argument("--case-number", type=int, choices=range(1, 9), required=True)
    run.add_argument("--records", type=Path, required=True)
    run.add_argument("--assets-root", type=Path, required=True)
    run.add_argument("--model", default="/data/hf_models/MiniMax-H3")
    run.add_argument("--num-inference-steps", type=int, default=20)
    run.add_argument("--warmup-steps", type=int, default=0)
    run.add_argument("--short-edge", type=int, default=768)
    run.add_argument("--aspect-ratio", default="9:16")
    run.add_argument("--duration-seconds", type=float, default=15.0)
    run.add_argument("--flow-shift", type=float, default=12.0)
    run.add_argument("--audio-flow-shift", type=float, default=3.0)
    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=7200.0)
    run.add_argument("--output", type=Path, required=True)
    summarize = subparsers.add_parser("summarize")
    summarize.add_argument("--result-root", type=Path, required=True)
    summarize.add_argument("--output", type=Path, required=True)
    return parser.parse_args()


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


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