#!/usr/bin/env python3
"""
CUDASlicer Benchmark Suite
==========================
Comprehensive performance benchmark comparing:
1. Baseline: PrusaSlicer CLI
2. CUDASlicer Matrix Configurations:
   - Backends: CUDA (GPU) vs CPU (Multi-threaded)
   - Interpolation: Sub-Pixel LERP vs No LERP (Discrete Marching Squares)
   - Grid Resolutions: 0.10 mm vs 0.05 mm Pitch
"""

from __future__ import annotations

import argparse
import os
import shutil
import struct
import subprocess
import sys
import time
from pathlib import Path
from typing import Any, Dict, List, Optional, Sequence, Tuple

import numpy as np

# Ensure local cudaslicer python package is discoverable
REPO_ROOT = Path(__file__).resolve().parent
PYTHON_PKG_DIR = REPO_ROOT / "python"
if PYTHON_PKG_DIR.exists() and str(PYTHON_PKG_DIR) not in sys.path:
    sys.path.insert(0, str(PYTHON_PKG_DIR))

try:
    import cudaslicer
except ImportError:
    cudaslicer = None


# ==============================================================================
# Fast Binary / ASCII STL Loader
# ==============================================================================
def load_stl(file_path: str | Path) -> np.ndarray:
    """
    Load an STL file into a NumPy array of shape (N, 3, 3) representing 3D triangles.
    Supports both Binary and ASCII STL files without external dependencies.
    """
    p = Path(file_path)
    if not p.exists():
        raise FileNotFoundError(f"STL file not found: {p}")

    file_size = p.stat().st_size
    if file_size < 84:
        raise ValueError(f"File {p} is too small to be a valid STL.")

    with open(p, "rb") as f:
        header = f.read(80)
        num_triangles_bytes = f.read(4)
        num_triangles = struct.unpack("<I", num_triangles_bytes)[0]
        expected_size = 84 + num_triangles * 50

        # Check if binary STL matches expected file size
        if expected_size == file_size:
            f.seek(84)
            record_dtype = np.dtype([
                ("normal", np.float32, 3),
                ("v0", np.float32, 3),
                ("v1", np.float32, 3),
                ("v2", np.float32, 3),
                ("attr", np.uint16),
            ])
            data = np.fromfile(f, dtype=record_dtype, count=num_triangles)
            return np.stack([data["v0"], data["v1"], data["v2"]], axis=1)

    # Fallback to ASCII parser
    vertices: list[list[float]] = []
    with open(p, "r", encoding="utf-8", errors="ignore") as f:
        for line in f:
            line_str = line.strip()
            if line_str.startswith("vertex"):
                parts = line_str.split()
                if len(parts) >= 4:
                    vertices.append([float(parts[1]), float(parts[2]), float(parts[3])])

    if len(vertices) % 3 != 0 or len(vertices) == 0:
        raise ValueError(f"Invalid or corrupted STL file: {p}")

    arr = np.array(vertices, dtype=np.float32)
    return arr.reshape(-1, 3, 3)


# ==============================================================================
# PrusaSlicer Baseline Runner
# ==============================================================================
def find_prusaslicer_bin(custom_path: Optional[str] = None) -> Optional[str]:
    """Find a functional PrusaSlicer executable."""
    candidates = [
        custom_path,
        str(REPO_ROOT / "tests_and_tools" / "prusa_binaries" / "squashfs-root" / "usr" / "bin" / "prusa-slicer"),
        str(REPO_ROOT / "tests_and_tools" / "prusa_binaries" / "prusa-slicer.AppImage"),
        str(REPO_ROOT / "bin" / "prusa-slicer"),
        "prusa-slicer",
        "prusa-slicer-console",
    ]
    for cand in candidates:
        if cand and (os.path.isfile(cand) or shutil.which(cand)):
            return cand
    return None


def run_prusaslicer_benchmark(
    stl_path: str,
    prusa_bin: str,
    layer_height: float = 0.20,
    perimeters: int = 2,
    layer_width: float = 0.45,
    runs: int = 3,
    expected_layers: Optional[int] = None,
) -> Optional[Dict[str, Any]]:
    """Execute PrusaSlicer CLI benchmark and measure execution time."""
    tmp_gcode = REPO_ROOT / "tests_and_tools" / "gcode_outputs" / "temp_bench_prusa.gcode"
    tmp_gcode.parent.mkdir(parents=True, exist_ok=True)

    cmd = [
        prusa_bin,
        "--export-gcode",
        f"--layer-height={layer_height}",
        f"--perimeters={perimeters}",
        "--fill-density=0%",
        f"--extrusion-width={layer_width}",
        f"--output={tmp_gcode}",
        stl_path,
    ]

    timings = []
    for r in range(runs):
        t0 = time.perf_counter()
        res = subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True)
        t1 = time.perf_counter()
        if res.returncode != 0:
            print(f"\n[Warning] PrusaSlicer failed with exit code {res.returncode}:")
            print(res.stderr)
            return None
        timings.append((t1 - t0) * 1000.0)

    # Clean up temp file
    if tmp_gcode.exists():
        tmp_gcode.unlink()

    return {
        "name": "PrusaSlicer CLI (Baseline)",
        "backend": "CPU",
        "pitch_mm": None,
        "subpixel_lerp": None,
        "time_ms": float(np.mean(timings)),
        "std_ms": float(np.std(timings)),
        "timings": timings,
        "layers_count": expected_layers,
        "total_vertices": None,
        "total_contours": None,
    }


# ==============================================================================
# CUDASlicer Benchmark Runner
# ==============================================================================
def run_cudaslicer_combination(
    tris: np.ndarray,
    backend: str,
    pitch_mm: float,
    subpixel_lerp: bool,
    layer_height: float = 0.20,
    perimeters_count: int = 2,
    layer_width: float = 0.45,
    pad_mm: float = 1.5,
    runs: int = 3,
    warmup: int = 1,
) -> Dict[str, Any]:
    """Benchmark a specific CUDASlicer parameter combination."""
    # Warmup
    for _ in range(warmup):
        _ = cudaslicer.slice_mesh_and_offset(
            mesh_or_triangles=tris,
            layer_height=layer_height,
            perimeters_count=perimeters_count,
            layer_width=layer_width,
            pitch_mm=pitch_mm,
            pad_mm=pad_mm,
            backend=backend,
            subpixel_lerp=subpixel_lerp,
        )

    timings = []
    last_layers = None
    for _ in range(runs):
        t0 = time.perf_counter()
        last_layers = cudaslicer.slice_mesh_and_offset(
            mesh_or_triangles=tris,
            layer_height=layer_height,
            perimeters_count=perimeters_count,
            layer_width=layer_width,
            pitch_mm=pitch_mm,
            pad_mm=pad_mm,
            backend=backend,
            subpixel_lerp=subpixel_lerp,
        )
        t1 = time.perf_counter()
        timings.append((t1 - t0) * 1000.0)

    total_vertices = 0
    total_contours = 0
    if last_layers:
        for layer in last_layers:
            for perim in getattr(layer, "perimeters", []):
                for cnt in perim.get("contours", []):
                    total_contours += 1
                    total_vertices += len(cnt)

    return {
        "name": f"CUDASlicer [{backend.upper()}] {pitch_mm:.2f}mm {'LERP' if subpixel_lerp else 'No LERP'}",
        "backend": backend.upper(),
        "pitch_mm": pitch_mm,
        "subpixel_lerp": subpixel_lerp,
        "time_ms": float(np.mean(timings)),
        "std_ms": float(np.std(timings)),
        "timings": timings,
        "layers_count": len(last_layers) if last_layers else 0,
        "total_vertices": total_vertices,
        "total_contours": total_contours,
    }


# ==============================================================================
# Console Reporting & Formatting
# ==============================================================================
def print_header(title: str) -> None:
    width = 98
    print("\n" + "=" * width)
    print(f" {title.center(width - 2)} ")
    print("=" * width)


def print_table(
    results: List[Dict[str, Any]],
    baseline_time_ms: Optional[float] = None,
    cpu_ref_map: Optional[Dict[Tuple[float, bool], float]] = None,
) -> None:
    """Print a clean formatted benchmark summary table to console."""
    # Header format
    hdr = f"{'Configuration':<38} | {'Pitch':<6} | {'LERP':<7} | {'Time (ms)':<11} | {'ms/Layer':<9} | {'Speedup (Base)':<15} | {'GPU vs CPU':<11}"
    sep = "-" * len(hdr)

    print("\n" + sep)
    print(hdr)
    print(sep)

    for r in results:
        name = r["name"]
        pitch_str = f"{r['pitch_mm']:.2f} mm" if r["pitch_mm"] is not None else "-"
        lerp_str = "Yes" if r["subpixel_lerp"] is True else ("No" if r["subpixel_lerp"] is False else "-")
        t_ms = r["time_ms"]
        num_layers = r["layers_count"] or 1
        ms_per_layer = t_ms / num_layers if num_layers > 0 else 0.0

        # Baseline speedup
        if baseline_time_ms is not None and baseline_time_ms > 0:
            speedup_base = baseline_time_ms / t_ms
            speedup_base_str = f"{speedup_base:6.2f}x"
        else:
            speedup_base_str = "-"

        # GPU vs CPU speedup for same pitch & lerp
        gpu_vs_cpu_str = "-"
        if r["backend"] == "CUDA" and cpu_ref_map:
            key = (r["pitch_mm"], r["subpixel_lerp"])
            if key in cpu_ref_map and cpu_ref_map[key] > 0:
                gpu_speedup = cpu_ref_map[key] / t_ms
                gpu_vs_cpu_str = f"{gpu_speedup:6.2f}x"

        print(
            f"{name:<38} | {pitch_str:<6} | {lerp_str:<7} | {t_ms:9.2f} ± {r['std_ms']:3.1f} | {ms_per_layer:8.2f} | {speedup_base_str:^15} | {gpu_vs_cpu_str:^11}"
        )

    print(sep + "\n")


# ==============================================================================
# Main Benchmark CLI
# ==============================================================================
def main() -> None:
    parser = argparse.ArgumentParser(
        description="CUDASlicer Performance Benchmark Suite (PrusaSlicer Baseline vs CPU/GPU, LERP/No LERP, 0.05mm/0.10mm pitch)"
    )
    default_stl = REPO_ROOT / "benchy.stl"
    parser.add_argument("--stl", type=str, default=str(default_stl), help="Path to STL mesh model (default: benchy.stl)")
    parser.add_argument("--layer-height", type=float, default=0.20, help="Slicing layer height in mm (default: 0.20)")
    parser.add_argument("--perimeters", type=int, default=2, help="Number of perimeter passes (default: 2)")
    parser.add_argument("--layer-width", type=float, default=0.45, help="Extrusion line width in mm (default: 0.45)")
    parser.add_argument("--pitches", nargs="+", type=float, default=[0.10, 0.05], help="Grid pixel pitches in mm (default: 0.10 0.05)")
    parser.add_argument("--runs", type=int, default=3, help="Number of benchmark iterations to average (default: 3)")
    parser.add_argument("--warmup", type=int, default=1, help="Number of warmup runs (default: 1)")
    parser.add_argument("--no-prusa", action="store_true", help="Skip PrusaSlicer baseline benchmark")
    parser.add_argument("--prusa-bin", type=str, default=None, help="Path to PrusaSlicer executable")
    parser.add_argument("--gpu-only", action="store_true", help="Run only GPU (CUDA) benchmarks")
    parser.add_argument("--cpu-only", action="store_true", help="Run only CPU benchmarks")
    parser.add_argument("--lerp-only", action="store_true", help="Run only Sub-pixel LERP passes")
    parser.add_argument("--no-lerp-only", action="store_true", help="Run only No-LERP (discrete) passes")

    args = parser.parse_args()

    print_header("CUDASLICER PERFORMANCE BENCHMARK SUITE")

    # 1. System & Hardware Details
    stl_path = Path(args.stl).resolve()
    if not stl_path.exists():
        print(f"[Error] Specified STL file does not exist: {stl_path}")
        sys.exit(1)

    print(f"Loading Model: {stl_path.name}")
    tris = load_stl(stl_path)
    num_tris = len(tris)
    bounds_min = tris.min(axis=(0, 1))
    bounds_max = tris.max(axis=(0, 1))
    dimensions = bounds_max - bounds_min

    print(f"  • Triangles:      {num_tris:,}")
    print(f"  • Dimensions:     {dimensions[0]:.2f} x {dimensions[1]:.2f} x {dimensions[2]:.2f} mm")
    print(f"  • Layer Height:   {args.layer_height:.2f} mm")
    print(f"  • Perimeters:     {args.perimeters} passes @ {args.layer_width:.2f} mm line width")
    print(f"  • Iterations:     {args.runs} runs (warmup: {args.warmup})")

    # Check cudaslicer availability
    if not cudaslicer or not cudaslicer.is_available():
        print("[Error] cudaslicer native backend is not loaded or available.")
        sys.exit(1)

    has_gpu = cudaslicer.has_cuda()
    print(f"  • CUDA GPU:       {'Available' if has_gpu else 'Not Available'}")
    if not has_gpu and not args.cpu_only:
        print("    (Notice: No CUDA GPU detected, forcing CPU execution)")

    results: List[Dict[str, Any]] = []

    expected_layers = int(np.ceil(dimensions[2] / args.layer_height))

    # 2. Benchmark Baseline: PrusaSlicer
    prusa_res = None
    if not args.no_prusa:
        prusa_bin = find_prusaslicer_bin(args.prusa_bin)
        if prusa_bin:
            print(f"\n[1/3] Benchmarking Baseline: PrusaSlicer CLI ({Path(prusa_bin).name})...")
            prusa_res = run_prusaslicer_benchmark(
                stl_path=str(stl_path),
                prusa_bin=prusa_bin,
                layer_height=args.layer_height,
                perimeters=args.perimeters,
                layer_width=args.layer_width,
                runs=args.runs,
                expected_layers=expected_layers,
            )
            if prusa_res:
                print(f"  -> PrusaSlicer Baseline: {prusa_res['time_ms']:.2f} ± {prusa_res['std_ms']:.2f} ms")
                results.append(prusa_res)
        else:
            print("\n[Warning] PrusaSlicer binary not found. Skipping baseline comparison.")

    # 3. Build combination matrix
    backends: list[str] = []
    if args.cpu_only:
        backends = ["cpu"]
    elif args.gpu_only:
        backends = ["cuda"] if has_gpu else []
    else:
        backends = (["cuda"] if has_gpu else []) + ["cpu"]

    lerp_modes: list[bool] = []
    if args.lerp_only:
        lerp_modes = [True]
    elif args.no_lerp_only:
        lerp_modes = [False]
    else:
        lerp_modes = [True, False]

    pitches = args.pitches

    combo_idx = 1
    cpu_ref_map: Dict[Tuple[float, bool], float] = {}

    for backend in backends:
        for pitch in pitches:
            if backend.lower() == "cpu" and pitch < 0.09 and not args.cpu_only:
                # Skip 0.05mm on CPU by default as it operates on 200M+ dense grid on CPU
                continue

            for lerp in lerp_modes:
                lerp_tag = "LERP" if lerp else "No LERP"
                print(
                    f"  [{combo_idx}] Running {backend.upper()} | Pitch: {pitch:.2f}mm | {lerp_tag} ...",
                    end="",
                    flush=True,
                )
                res = run_cudaslicer_combination(
                    tris=tris,
                    backend=backend,
                    pitch_mm=pitch,
                    subpixel_lerp=lerp,
                    layer_height=args.layer_height,
                    perimeters_count=args.perimeters,
                    layer_width=args.layer_width,
                    runs=args.runs,
                    warmup=args.warmup,
                )
                print(f" {res['time_ms']:.2f} ms ({res['time_ms']/res['layers_count']:.2f} ms/layer)")
                results.append(res)
                if backend.lower() == "cpu":
                    cpu_ref_map[(pitch, lerp)] = res["time_ms"]
                combo_idx += 1

    # 4. Print Summary Report
    print_header("BENCHMARK RESULTS & SPEEDUP COMPARISON")
    baseline_time = prusa_res["time_ms"] if prusa_res else None
    print_table(results, baseline_time_ms=baseline_time, cpu_ref_map=cpu_ref_map)

    # Key Highlights
    gpu_results = [r for r in results if r["backend"] == "CUDA"]
    if gpu_results and baseline_time:
        fastest_gpu = min(gpu_results, key=lambda x: x["time_ms"])
        speedup = baseline_time / fastest_gpu["time_ms"]
        print("Performance Highlights:")
        print(f"  ★ Fastest GPU Configuration: {fastest_gpu['name']} in {fastest_gpu['time_ms']:.2f} ms")
        print(f"  ★ GPU Speedup vs PrusaSlicer Baseline: {speedup:.2f}x faster")
        if cpu_ref_map:
            key = (fastest_gpu["pitch_mm"], fastest_gpu["subpixel_lerp"])
            if key in cpu_ref_map:
                gpu_vs_cpu = cpu_ref_map[key] / fastest_gpu["time_ms"]
                print(f"  ★ GPU Acceleration vs Multi-threaded CPU: {gpu_vs_cpu:.2f}x faster")


if __name__ == "__main__":
    main()