from __future__ import annotations

import argparse
import asyncio
import dataclasses
import json
import math
import os
import random
import statistics
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional

import torch

# Make the repo root importable when running directly from the benchmarks/ dir.
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

import comfy.model_management
import comfy.sd


# -----------------------------
# Data models
# -----------------------------

@dataclasses.dataclass
class RequestSpec:
    profile_name: str
    batch_size: int
    width: int
    height: int
    num_frames: int
    steps: int
    cfg_scale: float
    seed: int
    timeout_s: float = 180.0
    extra: Dict[str, Any] = dataclasses.field(default_factory=dict)


@dataclasses.dataclass
class RequestResult:
    request_id: int
    profile_name: str
    ok: bool
    error: Optional[str]
    latency_ms: float
    queue_wait_ms: float
    step_latencies_ms: List[float]
    ttfs_ms: float          # time to first (denoising) step
    peak_vram_mb: float
    est_mem_mb: Optional[float]
    started_at: float
    ended_at: float


@dataclasses.dataclass
class RunSummary:
    total_requests: int
    success: int
    failed: int
    throughput_req_s: float
    p50_ms: float
    p90_ms: float
    p95_ms: float
    p99_ms: float
    mean_ms: float
    ttfs_p50_ms: float
    ttfs_p99_ms: float
    step_mean_ms: float
    step_p99_ms: float
    max_vram_mb: float


# -----------------------------
# Helpers
# -----------------------------

def percentile(values: List[float], p: float) -> float:
    if not values:
        return float("nan")
    values = sorted(values)
    k = (len(values) - 1) * (p / 100.0)
    f = math.floor(k)
    c = math.ceil(k)
    if f == c:
        return values[int(k)]
    return values[f] * (c - k) + values[c] * (k - f)


def now() -> float:
    return time.perf_counter()


def gpu_peak_mb() -> float:
    if not torch.cuda.is_available():
        return 0.0
    return torch.cuda.max_memory_allocated() / (1024 ** 2)


def reset_gpu_peak() -> None:
    if torch.cuda.is_available():
        torch.cuda.reset_peak_memory_stats()


def sync_cuda() -> None:
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def build_request_stream(
    num_requests: int,
    base_seed: int,
    profiles: List[RequestSpec],
    weighted: Optional[List[float]] = None,
) -> List[RequestSpec]:
    rnd = random.Random(base_seed)
    out: List[RequestSpec] = []
    for i in range(num_requests):
        p = rnd.choices(profiles, weights=weighted, k=1)[0]
        out.append(dataclasses.replace(p, seed=base_seed + i))
    return out


# -----------------------------
# Model adapter
# -----------------------------

class WanRunner:
    """
    Thin adapter around ComfyUI model loading + the BaseModel.apply_model call path.

    Only the DiT denoiser is timed — no VAE encode/decode, no CLIP, no scheduler
    overhead — so measurements reflect true model inference cost.

    Latent shape convention (WAN):  [B, 16, T, H//8, W//8]
    Text conditioning shape (UMT5): [B, text_seq_len, text_dim]  (zeros for benchmarking)
    Sigma schedule (flow-matching):  linspace(1.0 → 1/steps, steps)
    """

    def __init__(
        self,
        checkpoint: str,
        device: str,
        dtype_str: str,
        text_seq_len: int = 512,
        text_dim: int = 4096,
    ):
        self.checkpoint = checkpoint
        self.device_str = device
        self.dtype_str = dtype_str
        self.text_seq_len = text_seq_len
        self.text_dim = text_dim
        self.patcher, self.model = self._load_model()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _load_model(self):
        dtype_map = {
            "fp16": torch.float16,
            "bf16": torch.bfloat16,
            "fp32": torch.float32,
        }
        dtype = dtype_map.get(self.dtype_str)
        model_opts = {"dtype": dtype} if dtype is not None else {}

        patcher = comfy.sd.load_diffusion_model(self.checkpoint, model_options=model_opts)
        # force_full_load=True keeps the whole model resident on GPU rather than
        # streaming weights on demand (important for latency benchmarking).
        comfy.model_management.load_models_gpu([patcher], force_full_load=True)
        return patcher, patcher.model

    def _estimate_mem_mb(self, latent_shape: tuple, text_seq_len: int) -> Optional[float]:
        cond_shapes = {
            "c_crossattn": [(latent_shape[0], text_seq_len, self.text_dim)],
        }
        try:
            return self.model.memory_required(latent_shape, cond_shapes) / (1024 ** 2)
        except Exception:
            return None

    # ------------------------------------------------------------------
    # Single-request execution
    # ------------------------------------------------------------------

    @torch.inference_mode()
    def run_one(self, req: RequestSpec) -> RequestResult:
        start = now()
        reset_gpu_peak()

        step_latencies: List[float] = []
        ttfs_ms = float("nan")
        est_mem_mb: Optional[float] = None
        ok = True
        err = None

        try:
            device = comfy.model_management.get_torch_device()
            dtype = self.model.get_dtype_inference()

            # Latent noise tensor: [B, 16 channels, T frames, H/8, W/8]
            latent_shape = (
                req.batch_size, 16,
                req.num_frames,
                req.height // 8,
                req.width // 8,
            )
            x = torch.randn(latent_shape, dtype=dtype, device=device)
            est_mem_mb = self._estimate_mem_mb(latent_shape, self.text_seq_len)

            # Fake text conditioning — zeros have the right shape, non-zero
            # values are not needed for throughput/latency benchmarking.
            cross_attn = torch.zeros(
                req.batch_size, self.text_seq_len, self.text_dim,
                dtype=dtype, device=device,
            )

            # Linear sigma schedule: 1.0 → 1/steps  (flow-matching, noise→clean)
            sigmas = torch.linspace(1.0, 1.0 / req.steps, req.steps, device=device)

            for step_i, sigma_val in enumerate(sigmas):
                sigma_t = sigma_val.expand(req.batch_size)
                t0 = now()
                x = self.model.apply_model(x, sigma_t, c_crossattn=cross_attn)
                sync_cuda()
                elapsed_ms = (now() - t0) * 1000.0
                step_latencies.append(elapsed_ms)
                if step_i == 0:
                    ttfs_ms = elapsed_ms

        except Exception as e:
            ok = False
            err = repr(e)

        end = now()
        return RequestResult(
            request_id=-1,
            profile_name=req.profile_name,
            ok=ok,
            error=err,
            latency_ms=(end - start) * 1000.0,
            queue_wait_ms=0.0,     # filled in by the scheduler
            step_latencies_ms=step_latencies,
            ttfs_ms=ttfs_ms,
            peak_vram_mb=gpu_peak_mb(),
            est_mem_mb=est_mem_mb,
            started_at=start,
            ended_at=end,
        )


# -----------------------------
# Serving-style scheduler
# -----------------------------

async def run_closed_loop(
    runner: WanRunner,
    requests: List[RequestSpec],
    concurrency: int,
    request_rate: float = float("inf"),
) -> List[RequestResult]:
    """
    Closed-loop scheduler (default) or Poisson open-loop when request_rate is finite.

    Each request is dispatched to a thread so the asyncio event loop stays
    free to issue the next request while the GPU is busy.
    """
    sem = asyncio.Semaphore(concurrency)
    results: List[Optional[RequestResult]] = [None] * len(requests)

    async def worker(i: int, req: RequestSpec) -> None:
        async with sem:
            t_enq = now()
            res = await asyncio.to_thread(runner.run_one, req)
            res.request_id = i
            res.queue_wait_ms = max(0.0, (res.started_at - t_enq) * 1000.0)
            results[i] = res

    if request_rate == float("inf") or request_rate <= 0:
        await asyncio.gather(*(worker(i, r) for i, r in enumerate(requests)))
    else:
        tasks: List[asyncio.Task] = []
        for i, req in enumerate(requests):
            if i > 0:
                await asyncio.sleep(random.expovariate(request_rate))
            tasks.append(asyncio.create_task(worker(i, req)))
        await asyncio.gather(*tasks)

    return [r for r in results if r is not None]


def summarize(results: List[RequestResult], wall_s: float) -> RunSummary:
    lat = [r.latency_ms for r in results if r.ok]
    ttfs = [r.ttfs_ms for r in results if r.ok and math.isfinite(r.ttfs_ms)]
    all_steps = [s for r in results if r.ok for s in r.step_latencies_ms]
    succ = sum(1 for r in results if r.ok)
    fail = len(results) - succ
    return RunSummary(
        total_requests=len(results),
        success=succ,
        failed=fail,
        throughput_req_s=(succ / wall_s) if wall_s > 0 else 0.0,
        p50_ms=percentile(lat, 50),
        p90_ms=percentile(lat, 90),
        p95_ms=percentile(lat, 95),
        p99_ms=percentile(lat, 99),
        mean_ms=(statistics.mean(lat) if lat else float("nan")),
        ttfs_p50_ms=percentile(ttfs, 50),
        ttfs_p99_ms=percentile(ttfs, 99),
        step_mean_ms=(statistics.mean(all_steps) if all_steps else float("nan")),
        step_p99_ms=percentile(all_steps, 99),
        max_vram_mb=max((r.peak_vram_mb for r in results), default=0.0),
    )


def print_summary(
    args: argparse.Namespace,
    summ: RunSummary,
    total_requests: int,
    wall_s: float,
) -> None:
    w = 60
    sep = "-" * w
    print("\n" + "=" * w)
    print("{s:^{n}}".format(s=" WAN Benchmark Result ", n=w))
    print("=" * w)
    print("{:<40} {:<}".format("Checkpoint:", Path(args.checkpoint).name))
    print("{:<40} {:<}".format("Device / dtype:", f"{args.device}/{args.dtype}"))
    print("{:<40} {:<}".format("Concurrency:", args.concurrency))
    rate_str = f"{args.request_rate:.1f} req/s" if args.request_rate != float("inf") else "inf (closed-loop)"
    print("{:<40} {:<}".format("Request rate:", rate_str))
    print(sep)
    print("{:<40} {:<.2f}".format("Benchmark duration (s):", wall_s))
    print("{:<40} {}/{}".format("Successful requests:", summ.success, total_requests))
    if summ.failed:
        print("{:<40} {:<}".format("Failed requests:", summ.failed))
    print(sep)
    print("{:<40} {:<.3f}".format("Throughput (req/s):", summ.throughput_req_s))
    print("{:<40} {:<.1f}".format("Latency mean (ms):", summ.mean_ms))
    print("{:<40} {:<.1f}".format("Latency p50  (ms):", summ.p50_ms))
    print("{:<40} {:<.1f}".format("Latency p90  (ms):", summ.p90_ms))
    print("{:<40} {:<.1f}".format("Latency p95  (ms):", summ.p95_ms))
    print("{:<40} {:<.1f}".format("Latency p99  (ms):", summ.p99_ms))
    print(sep)
    print("{:<40} {:<.1f}".format("TTFS p50 (ms):", summ.ttfs_p50_ms))
    print("{:<40} {:<.1f}".format("TTFS p99 (ms):", summ.ttfs_p99_ms))
    print("{:<40} {:<.1f}".format("Step latency mean (ms):", summ.step_mean_ms))
    print("{:<40} {:<.1f}".format("Step latency p99  (ms):", summ.step_p99_ms))
    print(sep)
    print("{:<40} {:<.1f}".format("Peak VRAM (MB):", summ.max_vram_mb))
    print("=" * w)


# -----------------------------
# CLI
# -----------------------------

def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Benchmark ComfyUI WAN diffusion model denoising throughput and latency."
    )
    p.add_argument(
        "--checkpoint", required=True,
        help="Path to the WAN diffusion-model checkpoint (.safetensors / .pt).",
    )
    p.add_argument("--device", default="cuda")
    p.add_argument("--dtype", default="fp16", choices=["fp16", "bf16", "fp32"])
    p.add_argument("--num-requests", type=int, default=100)
    p.add_argument("--concurrency", type=int, default=4,
                   help="Max number of in-flight requests (semaphore width).")
    p.add_argument(
        "--request-rate", type=float, default=float("inf"),
        help="Poisson arrival rate in req/s.  inf = closed-loop (default).",
    )
    p.add_argument("--warmup-requests", type=int, default=2,
                   help="Warmup iterations excluded from metrics.")
    p.add_argument("--seed", type=int, default=1234)
    p.add_argument("--text-seq-len", type=int, default=512,
                   help="Cross-attention sequence length (UMT5 default: 512).")
    p.add_argument("--text-dim", type=int, default=4096,
                   help="Text embedding width (UMT5-XXL: 4096).")
    p.add_argument("--out-dir", type=Path, default=Path("benchmarks/out"))
    p.add_argument("--output-file", type=Path, default=None,
                   help="Override path for the summary JSON output.")
    return p.parse_args()


def default_profiles() -> List[RequestSpec]:
    return [
        RequestSpec("wan21_t2v_720p_16f_30s", 1, 1280, 720, 16, 30, 6.0, 0),
        RequestSpec("wan21_t2v_720p_32f_30s", 1, 1280, 720, 32, 30, 6.0, 0),
        RequestSpec("wan21_t2v_480p_32f_20s", 1,  854, 480, 32, 20, 6.0, 0),
    ]


async def main_async() -> None:
    args = parse_args()
    args.out_dir.mkdir(parents=True, exist_ok=True)

    runner = WanRunner(
        checkpoint=args.checkpoint,
        device=args.device,
        dtype_str=args.dtype,
        text_seq_len=args.text_seq_len,
        text_dim=args.text_dim,
    )

    all_reqs = build_request_stream(
        args.num_requests + args.warmup_requests,
        args.seed,
        default_profiles(),
    )
    warmup_reqs = all_reqs[: args.warmup_requests]
    bench_reqs = all_reqs[args.warmup_requests :]

    if warmup_reqs:
        print(f"Running {len(warmup_reqs)} warmup request(s)...")
        for req in warmup_reqs:
            runner.run_one(req)
        print("Warmup complete.")

    print(f"Benchmarking {len(bench_reqs)} requests (concurrency={args.concurrency})...")
    t0 = now()
    results = await run_closed_loop(runner, bench_reqs, args.concurrency, args.request_rate)
    wall_s = now() - t0

    summ = summarize(results, wall_s)
    print_summary(args, summ, len(bench_reqs), wall_s)

    out_file = args.output_file or (args.out_dir / "summary.json")
    with open(args.out_dir / "requests.jsonl", "w") as f:
        for r in results:
            f.write(json.dumps(dataclasses.asdict(r)) + "\n")
    with open(out_file, "w") as f:
        json.dump(dataclasses.asdict(summ), f, indent=2)
    print(f"\nResults written to {args.out_dir}/")


if __name__ == "__main__":
    asyncio.run(main_async())