"""Benchmarking utilities: FPS, latency, throughput per pipeline stage.

Usage:
    from src.utils.benchmark import benchmark_stage, benchmark_full_pipeline
"""

from __future__ import annotations

import logging
import time
from dataclasses import dataclass, field
from typing import Any, Callable

import numpy as np
import torch
from torch.utils.data import DataLoader, Dataset
from tqdm import tqdm

logger = logging.getLogger(__name__)


# ---------------------------------------------------------------------------
# Result container
# ---------------------------------------------------------------------------

@dataclass
class BenchmarkResult:
    """Benchmark measurement for a single stage or full pipeline."""

    stage: str
    num_images: int
    total_time_s: float
    batch_times_s: list[float] = field(default_factory=list)

    @property
    def fps(self) -> float:
        return self.num_images / self.total_time_s if self.total_time_s > 0 else 0.0

    @property
    def avg_latency_ms(self) -> float:
        return (self.total_time_s / self.num_images * 1000) if self.num_images > 0 else 0.0

    @property
    def median_batch_latency_ms(self) -> float:
        if not self.batch_times_s:
            return 0.0
        return float(np.median(self.batch_times_s)) * 1000

    @property
    def p95_batch_latency_ms(self) -> float:
        if not self.batch_times_s:
            return 0.0
        return float(np.percentile(self.batch_times_s, 95)) * 1000

    def log_summary(self) -> None:
        logger.info("📊 Benchmark [%s]:", self.stage)
        logger.info("   Images: %d", self.num_images)
        logger.info("   Total:  %.2f s", self.total_time_s)
        logger.info("   FPS:    %.1f img/s", self.fps)
        logger.info("   Avg latency:    %.1f ms/img", self.avg_latency_ms)
        logger.info("   Median batch:   %.1f ms", self.median_batch_latency_ms)
        logger.info("   P95 batch:      %.1f ms", self.p95_batch_latency_ms)


# ---------------------------------------------------------------------------
# Synthetic dataset
# ---------------------------------------------------------------------------

class SyntheticImageDataset(Dataset):
    """Generate random RGB images for benchmarking."""

    def __init__(self, num_images: int, image_size: int = 256) -> None:
        self.num_images = num_images
        self.image_size = image_size

    def __len__(self) -> int:
        return self.num_images

    def __getitem__(self, idx: int) -> dict[str, Any]:
        return {
            "image_raw": torch.rand(3, self.image_size, self.image_size),
            "rel_path": f"synthetic/img_{idx:05d}.png",
            "stem": f"img_{idx:05d}",
            "output_dir": "",
        }


# ---------------------------------------------------------------------------
# Stage benchmarks
# ---------------------------------------------------------------------------

def benchmark_depth(
    model: Any,
    device: torch.device,
    num_images: int = 100,
    batch_size: int = 4,
    image_size: int = 256,
    num_workers: int = 0,
    warmup_batches: int = 3,
) -> BenchmarkResult:
    """Benchmark depth inference stage."""
    from src.augmentor.inference import infer_depth_batch

    ds = SyntheticImageDataset(num_images, image_size)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        num_workers=num_workers, pin_memory=torch.cuda.is_available())

    # Warmup
    for i, batch in enumerate(loader):
        if i >= warmup_batches:
            break
        infer_depth_batch(model, batch["image_raw"], device)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    batch_times = []
    processed = 0
    t_start = time.perf_counter()

    for batch in tqdm(loader, desc="⏱️  bench depth", unit="batch", colour="green"):
        t0 = time.perf_counter()
        infer_depth_batch(model, batch["image_raw"], device)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        batch_times.append(time.perf_counter() - t0)
        processed += batch["image_raw"].shape[0]

    total = time.perf_counter() - t_start
    result = BenchmarkResult("🌊 depth", processed, total, batch_times)
    result.log_summary()
    return result


def benchmark_chmv2(
    model: Any,
    processor: Any,
    device: torch.device,
    num_images: int = 100,
    batch_size: int = 4,
    image_size: int = 256,
    num_workers: int = 0,
    warmup_batches: int = 3,
) -> BenchmarkResult:
    """Benchmark CHMv2 inference stage."""
    from src.augmentor.inference import infer_chmv2_batch

    ds = SyntheticImageDataset(num_images, image_size)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        num_workers=num_workers, pin_memory=torch.cuda.is_available())

    # Warmup
    for i, batch in enumerate(loader):
        if i >= warmup_batches:
            break
        infer_chmv2_batch(model, processor, batch["image_raw"], device)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    batch_times = []
    processed = 0
    t_start = time.perf_counter()

    for batch in tqdm(loader, desc="⏱️  bench chmv2", unit="batch", colour="blue"):
        t0 = time.perf_counter()
        infer_chmv2_batch(model, processor, batch["image_raw"], device)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        batch_times.append(time.perf_counter() - t0)
        processed += batch["image_raw"].shape[0]

    total = time.perf_counter() - t_start
    result = BenchmarkResult("🏔️  chmv2", processed, total, batch_times)
    result.log_summary()
    return result


def benchmark_edges(
    num_images: int = 100,
    batch_size: int = 32,
    image_size: int = 256,
    warmup_batches: int = 3,
) -> BenchmarkResult:
    """Benchmark edges (Sobel) computation stage."""
    from src.augmentor.inference import compute_edges_from_depth

    total_batches = (num_images + batch_size - 1) // batch_size

    # Warmup
    for _ in range(warmup_batches):
        dummy = torch.rand(batch_size, 1, image_size, image_size)
        compute_edges_from_depth(dummy)

    batch_times = []
    processed = 0
    t_start = time.perf_counter()

    for i in tqdm(range(total_batches), desc="⏱️  bench edges", unit="batch", colour="cyan"):
        cur_bs = min(batch_size, num_images - i * batch_size)
        depth = torch.rand(cur_bs, 1, image_size, image_size)
        t0 = time.perf_counter()
        compute_edges_from_depth(depth)
        batch_times.append(time.perf_counter() - t0)
        processed += cur_bs

    total = time.perf_counter() - t_start
    result = BenchmarkResult("🔪 edges", processed, total, batch_times)
    result.log_summary()
    return result


def benchmark_segmentation(
    model: Any,
    seg_config: dict[str, Any],
    device: torch.device,
    num_images: int = 100,
    batch_size: int = 4,
    image_size: int = 256,
    num_workers: int = 0,
    warmup_batches: int = 3,
) -> BenchmarkResult:
    """Benchmark segmentation inference stage."""
    from src.augmentor.inference import infer_segmentation_batch

    ds = SyntheticImageDataset(num_images, image_size)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        num_workers=num_workers, pin_memory=torch.cuda.is_available())

    # Warmup
    for i, batch in enumerate(loader):
        if i >= warmup_batches:
            break
        infer_segmentation_batch(model, seg_config, batch["image_raw"], device)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    batch_times = []
    processed = 0
    t_start = time.perf_counter()

    for batch in tqdm(loader, desc="⏱️  bench segmentation", unit="batch", colour="yellow"):
        t0 = time.perf_counter()
        infer_segmentation_batch(model, seg_config, batch["image_raw"], device)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        batch_times.append(time.perf_counter() - t0)
        processed += batch["image_raw"].shape[0]

    total = time.perf_counter() - t_start
    result = BenchmarkResult("🗺️  segmentation", processed, total, batch_times)
    result.log_summary()
    return result


def benchmark_full_pipeline(
    depth_model: Any,
    seg_model: Any,
    seg_config: dict[str, Any],
    device: torch.device,
    num_images: int = 100,
    batch_size: int = 4,
    image_size: int = 256,
    num_workers: int = 0,
    warmup_batches: int = 3,
) -> dict[str, BenchmarkResult]:
    """Benchmark all stages and total pipeline throughput."""
    from src.augmentor.inference import (
        infer_depth_batch, compute_edges_from_depth, infer_segmentation_batch,
    )

    ds = SyntheticImageDataset(num_images, image_size)
    loader = DataLoader(ds, batch_size=batch_size, shuffle=False,
                        num_workers=num_workers, pin_memory=torch.cuda.is_available())

    # Warmup
    for i, batch in enumerate(loader):
        if i >= warmup_batches:
            break
        imgs = batch["image_raw"]
        depths = infer_depth_batch(depth_model, imgs, device)
        compute_edges_from_depth(depths)
        infer_segmentation_batch(seg_model, seg_config, imgs, device)
    if torch.cuda.is_available():
        torch.cuda.synchronize()

    depth_times, edge_times, seg_times, total_times = [], [], [], []
    processed = 0
    t_pipeline_start = time.perf_counter()

    for batch in tqdm(loader, desc="⏱️  bench full pipeline", unit="batch", colour="magenta"):
        imgs = batch["image_raw"]
        cur_bs = imgs.shape[0]

        t_total_0 = time.perf_counter()

        # Depth
        t0 = time.perf_counter()
        depths = infer_depth_batch(depth_model, imgs, device)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        depth_times.append(time.perf_counter() - t0)

        # Edges
        t0 = time.perf_counter()
        compute_edges_from_depth(depths)
        edge_times.append(time.perf_counter() - t0)

        # Segmentation
        t0 = time.perf_counter()
        infer_segmentation_batch(seg_model, seg_config, imgs, device)
        if torch.cuda.is_available():
            torch.cuda.synchronize()
        seg_times.append(time.perf_counter() - t0)

        total_times.append(time.perf_counter() - t_total_0)
        processed += cur_bs

    t_pipeline_total = time.perf_counter() - t_pipeline_start

    results = {
        "depth": BenchmarkResult("🌊 depth", processed, sum(depth_times), depth_times),
        "edges": BenchmarkResult("🔪 edges", processed, sum(edge_times), edge_times),
        "segmentation": BenchmarkResult("🗺️  segmentation", processed, sum(seg_times), seg_times),
        "total": BenchmarkResult("🏁 total pipeline", processed, t_pipeline_total, total_times),
    }

    logger.info("=" * 60)
    for r in results.values():
        r.log_summary()
    logger.info("=" * 60)

    return results
