"""
Multi-GPU scaling benchmark for Indic Heritage Studio v2.

Measures:
  - Latency per image at 1 / 2 / 4 / 8 GPU configurations
  - Throughput (images/min)
  - Peak VRAM per GPU
  - Cold-start time for each pipeline (T2I / Style / I2V / Inpaint / ControlNet)
  - ROCm / CUDA backend detection

Outputs:
  - outputs/benchmark_<host>_<backend>.json
  - outputs/benchmark_<host>_<backend>.md  (human-readable summary)
  - A scaling chart PNG (matplotlib)
"""
from __future__ import annotations

import argparse
import json
import logging
import os
import socket
import time
from dataclasses import asdict, dataclass, field
from datetime import datetime
from pathlib import Path
from typing import List, Optional

log = logging.getLogger(__name__)


@dataclass
class BenchmarkResult:
    config: str          # "1-gpu", "2-gpu", etc.
    pipeline: str        # "t2i", "style", "i2v", etc.
    num_gpus: int
    num_samples: int
    total_seconds: float
    avg_latency_seconds: float
    throughput_images_per_min: float
    peak_vram_gb: float
    extra: dict = field(default_factory=dict)


def _detect_backend() -> str:
    try:
        import torch
        if not torch.cuda.is_available():
            return "cpu"
        if hasattr(torch.version, "hip") and torch.version.hip is not None:
            return f"rocm-{torch.version.hip}"
        return f"cuda-{torch.version.cuda}"
    except Exception:
        return "unknown"


def _get_gpu_info() -> dict:
    try:
        import torch
        if not torch.cuda.is_available():
            return {"count": 0}
        return {
            "count": torch.cuda.device_count(),
            "name": torch.cuda.get_device_name(0),
            "vram_total_gb": torch.cuda.get_device_properties(0).total_memory / 1e9,
            "torch_version": torch.__version__,
        }
    except Exception:
        return {}


def _peak_vram() -> float:
    try:
        import torch
        if not torch.cuda.is_available():
            return 0.0
        # max across all GPUs
        peak = 0.0
        for i in range(torch.cuda.device_count()):
            peak = max(peak, torch.cuda.max_memory_allocated(i) / 1e9)
        return round(peak, 2)
    except Exception:
        return 0.0


def _reset_vram() -> None:
    try:
        import torch
        for i in range(torch.cuda.device_count()):
            with torch.cuda.device(i):
                torch.cuda.reset_peak_memory_stats()
                torch.cuda.empty_cache()
    except Exception:
        pass


# ---------------------------------------------------------------------------
# Per-pipeline benchmarks
# ---------------------------------------------------------------------------
def bench_t2i(num_gpus: int, num_samples: int = 4) -> BenchmarkResult:
    """Benchmark SDXL T2I across `num_gpus` parallel workers."""
    from config.styles import get_style
    style = get_style("madhubani")
    prompts = [
        "a young woman reading under a banyan tree at sunset",
        "a temple festival at dawn with devotees",
        "a peacock dancing in a monsoon garden",
        "a sage meditating by the river",
        "a musician playing the sitar in a moonlit courtyard",
        "a wedding procession with dancers and drummers",
    ]

    if num_gpus == 1:
        # Single-GPU sequential
        from core.text_to_image import TextToImagePipeline
        _reset_vram()
        pipe = TextToImagePipeline().load()
        times = []
        for i in range(num_samples):
            t0 = time.time()
            pipe.generate(prompts[i % len(prompts)], style, seed=42 + i)
            times.append(time.time() - t0)
        total = sum(times)
        avg = total / num_samples
        return BenchmarkResult(
            config=f"{num_gpus}-gpu",
            pipeline="t2i",
            num_gpus=num_gpus,
            num_samples=num_samples,
            total_seconds=round(total, 2),
            avg_latency_seconds=round(avg, 2),
            throughput_images_per_min=round(num_samples / total * 60, 1),
            peak_vram_gb=_peak_vram(),
        )
    else:
        # Multi-GPU parallel via batch processor
        from core.batch_processor import BatchProcessor, BatchJob
        _reset_vram()
        jobs = [
            BatchJob(
                input_path=Path("/dev/null"),
                output_path=Path(f"outputs/bench/t2i_{i}.png"),
                style_id="madhubani",
                mode="t2i",
                prompt=prompts[i % len(prompts)],
                seed=42 + i,
            )
            for i in range(num_samples)
        ]
        proc = BatchProcessor(num_workers=num_gpus)
        t0 = time.time()
        results = proc.run(jobs)
        total = time.time() - t0
        succ = sum(1 for r in results if r.success)
        return BenchmarkResult(
            config=f"{num_gpus}-gpu",
            pipeline="t2i",
            num_gpus=num_gpus,
            num_samples=succ,
            total_seconds=round(total, 2),
            avg_latency_seconds=round(total / max(succ, 1), 2),
            throughput_images_per_min=round(succ / max(total, 1) * 60, 1),
            peak_vram_gb=_peak_vram(),
            extra={"failed": len(results) - succ},
        )


def bench_style_transfer(num_gpus: int, num_samples: int = 4) -> BenchmarkResult:
    """Benchmark IP-Adapter XL style transfer."""
    # Reuse the batch processor infrastructure
    from core.batch_processor import BatchProcessor, BatchJob
    _reset_vram()
    # Generate synthetic input images first (one per sample)
    from PIL import Image
    inputs_dir = Path("outputs/bench/inputs")
    inputs_dir.mkdir(parents=True, exist_ok=True)
    for i in range(num_samples):
        img = Image.new("RGB", (1024, 1024),
                        tuple(int(c * 255) for c in [0.5, 0.3, 0.2]))
        img.save(inputs_dir / f"input_{i}.png")

    jobs = [
        BatchJob(
            input_path=inputs_dir / f"input_{i}.png",
            output_path=Path(f"outputs/bench/style_{i}.png"),
            style_id="warli",
            mode="style_transfer",
        )
        for i in range(num_samples)
    ]
    proc = BatchProcessor(num_workers=num_gpus)
    t0 = time.time()
    results = proc.run(jobs)
    total = time.time() - t0
    succ = sum(1 for r in results if r.success)
    return BenchmarkResult(
        config=f"{num_gpus}-gpu",
        pipeline="style_transfer",
        num_gpus=num_gpus,
        num_samples=succ,
        total_seconds=round(total, 2),
        avg_latency_seconds=round(total / max(succ, 1), 2),
        throughput_images_per_min=round(succ / max(total, 1) * 60, 1),
        peak_vram_gb=_peak_vram(),
    )


def bench_i2v(num_samples: int = 1) -> BenchmarkResult:
    """Benchmark SVD image-to-video on a single GPU (SVD doesn't shard well)."""
    from PIL import Image
    from config.styles import get_style
    from core.image_to_video import ImageToVideoPipeline
    _reset_vram()
    pipe = ImageToVideoPipeline().load()
    img = Image.new("RGB", (1024, 576), (128, 96, 64))
    times = []
    for i in range(num_samples):
        t0 = time.time()
        pipe.generate(img, style=get_style("tanjore"), seed=42 + i)
        times.append(time.time() - t0)
    total = sum(times)
    return BenchmarkResult(
        config="1-gpu",
        pipeline="i2v",
        num_gpus=1,
        num_samples=num_samples,
        total_seconds=round(total, 2),
        avg_latency_seconds=round(total / num_samples, 2),
        throughput_images_per_min=round(num_samples / total * 60, 1),
        peak_vram_gb=_peak_vram(),
        extra={"frames_per_video": 25},
    )


# ---------------------------------------------------------------------------
# Main
# ---------------------------------------------------------------------------
def run_full_benchmark(out_dir: Path, gpu_configs: List[int],
                       num_samples: int = 4) -> Path:
    """Run a full multi-config benchmark and write results."""
    backend = _detect_backend()
    host = socket.gethostname()
    timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
    gpu_info = _get_gpu_info()

    log.info("=== Benchmark start ===")
    log.info("Host: %s | Backend: %s | GPUs: %s", host, backend, gpu_info)
    log.info("Configs to test: %s", gpu_configs)

    results: List[BenchmarkResult] = []

    for n in gpu_configs:
        log.info("--- T2I with %d GPU(s) ---", n)
        try:
            r = bench_t2i(num_gpus=n, num_samples=num_samples)
            results.append(r)
            log.info("Result: %s", r)
        except Exception as exc:
            log.error("T2I %d-gpu failed: %s", n, exc)

    for n in gpu_configs:
        log.info("--- Style transfer with %d GPU(s) ---", n)
        try:
            r = bench_style_transfer(num_gpus=n, num_samples=num_samples)
            results.append(r)
            log.info("Result: %s", r)
        except Exception as exc:
            log.error("Style %d-gpu failed: %s", n, exc)

    log.info("--- I2V (single GPU, SVD doesn't shard) ---")
    try:
        r = bench_i2v(num_samples=max(1, num_samples // 2))
        results.append(r)
        log.info("Result: %s", r)
    except Exception as exc:
        log.error("I2V failed: %s", exc)

    # Write JSON
    out_dir.mkdir(parents=True, exist_ok=True)
    json_path = out_dir / f"benchmark_{host}_{backend}_{timestamp}.json"
    payload = {
        "host": host,
        "backend": backend,
        "timestamp": timestamp,
        "gpu_info": gpu_info,
        "num_samples_per_config": num_samples,
        "results": [asdict(r) for r in results],
    }
    json_path.write_text(json.dumps(payload, indent=2, default=str))

    # Write Markdown summary
    md_path = out_dir / f"benchmark_{host}_{backend}_{timestamp}.md"
    md_lines = [
        f"# Indic Heritage Studio v2 — Benchmark Report",
        "",
        f"- **Host:** `{host}`",
        f"- **Backend:** `{backend}`",
        f"- **GPUs:** {gpu_info.get('count', 0)} × {gpu_info.get('name', 'unknown')} "
        f"({gpu_info.get('vram_total_gb', 0):.1f} GB each)",
        f"- **Date:** {timestamp}",
        f"- **Samples per config:** {num_samples}",
        "",
        "## Results",
        "",
        "| Config | Pipeline | Samples | Total (s) | Avg latency (s) | Throughput (img/min) | Peak VRAM (GB) |",
        "|---|---|---|---|---|---|---|",
    ]
    for r in results:
        md_lines.append(
            f"| {r.config} | {r.pipeline} | {r.num_samples} | "
            f"{r.total_seconds} | {r.avg_latency_seconds} | "
            f"{r.throughput_images_per_min} | {r.peak_vram_gb} |"
        )
    md_lines.append("")
    md_path.write_text("\n".join(md_lines))

    # Scaling chart
    try:
        _plot_scaling(results, out_dir / f"scaling_{host}_{backend}_{timestamp}.png")
    except Exception as exc:
        log.warning("Scaling chart failed: %s", exc)

    log.info("=== Benchmark done. Reports: %s, %s ===", json_path, md_path)
    return json_path


def _plot_scaling(results: List[BenchmarkResult], out_path: Path) -> None:
    import matplotlib.font_manager as fm
    fm.fontManager.addfont('/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf')
    import matplotlib.pyplot as plt
    plt.rcParams['font.sans-serif'] = ['DejaVu Sans']
    plt.rcParams['axes.unicode_minus'] = False

    # Group by pipeline
    by_pipe: dict = {}
    for r in results:
        by_pipe.setdefault(r.pipeline, []).append(r)

    fig, ax = plt.subplots(figsize=(10, 6), constrained_layout=True)
    for pipe, rs in by_pipe.items():
        rs_sorted = sorted(rs, key=lambda r: r.num_gpus)
        xs = [r.num_gpus for r in rs_sorted]
        ys = [r.throughput_images_per_min for r in rs_sorted]
        ax.plot(xs, ys, marker="o", linewidth=2, label=pipe)

    ax.set_xlabel("Number of GPUs")
    ax.set_ylabel("Throughput (images / minute)")
    ax.set_title("Multi-GPU Scaling — Indic Heritage Studio v2")
    ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1.0))
    ax.grid(True, alpha=0.3)
    ax.set_xticks([1, 2, 4, 8])
    fig.savefig(out_path, dpi=200)
    plt.close(fig)


def _cli():
    p = argparse.ArgumentParser(description="Indic Heritage Studio v2 — Multi-GPU Benchmark")
    p.add_argument("--configs", nargs="+", type=int, default=[1, 2, 4, 8],
                   help="GPU counts to benchmark (default: 1 2 4 8)")
    p.add_argument("--samples", type=int, default=4,
                   help="Number of samples per config")
    p.add_argument("--out-dir", type=Path, default=Path("outputs/benchmarks"))
    args = p.parse_args()

    logging.basicConfig(level=logging.INFO,
                        format="%(asctime)s | %(levelname)s | %(message)s")
    run_full_benchmark(args.out_dir, args.configs, args.samples)


if __name__ == "__main__":
    _cli()