"""
Benchmark framework for FMUS-VID.

This module provides tools for measuring and comparing performance metrics.
"""

import json
import time
import statistics
from pathlib import Path
from typing import Dict, List, Optional, Any, Callable
from dataclasses import dataclass, asdict
import numpy as np
import psutil
import logging
from concurrent.futures import ThreadPoolExecutor
import fmusvid

logger = logging.getLogger(__name__)

@dataclass
class BenchmarkResult:
    """Single benchmark result."""
    name: str
    operation: str
    input_size: Dict[str, int]  # width, height, duration
    execution_time: float
    memory_usage: float
    cpu_usage: float
    fps: float
    backend: str
    hardware_info: Dict[str, str]

class Benchmark:
    """
    Base class for benchmarks.
    """

    def __init__(self, name: str):
        """
        Initialize benchmark.

        Args:
            name: Benchmark name
        """
        self.name = name
        self.results: List[BenchmarkResult] = []

    def run(self, iterations: int = 3) -> List[BenchmarkResult]:
        """
        Run benchmark multiple times.

        Args:
            iterations: Number of iterations

        Returns:
            List of benchmark results
        """
        for _ in range(iterations):
            try:
                result = self._run_single()
                self.results.append(result)
            except Exception as e:
                logger.error(f"Benchmark {self.name} failed: {e}")

        return self.results

    def _run_single(self) -> BenchmarkResult:
        """Run single benchmark iteration."""
        raise NotImplementedError

class BasicOperationsBenchmark(Benchmark):
    """
    Benchmark basic video operations.
    """

    def __init__(self, video_path: Path):
        """
        Initialize basic operations benchmark.

        Args:
            video_path: Path to test video
        """
        super().__init__("basic_operations")
        self.video_path = Path(video_path)

    def _run_single(self) -> BenchmarkResult:
        """Run basic operations benchmark."""
        video = fmusvid.load(self.video_path)
        start_time = time.time()
        start_memory = self._get_memory_usage()

        # Menjalankan operasi dasar
        processed = (video
            .trim(start=1.0, end=5.0)
            .resize(width=1280)
            .rotate(90)
            .grayscale()
        )

        processed.save("benchmark_output.mp4")

        end_time = time.time()
        end_memory = self._get_memory_usage()

        return BenchmarkResult(
            name=self.name,
            operation="basic_chain",
            input_size={
                "width": video.width,
                "height": video.height,
                "duration": video.duration
            },
            execution_time=end_time - start_time,
            memory_usage=end_memory - start_memory,
            cpu_usage=self._get_cpu_usage(),
            fps=video.duration / (end_time - start_time),
            backend=video.backend.name,
            hardware_info=self._get_hardware_info()
        )

    def _get_memory_usage(self) -> float:
        """Get current memory usage in MB."""
        process = psutil.Process()
        return process.memory_info().rss / 1024 / 1024

    def _get_cpu_usage(self) -> float:
        """Get CPU usage percentage."""
        return psutil.cpu_percent()

    def _get_hardware_info(self) -> Dict[str, str]:
        """Get system hardware information."""
        return {
            "cpu": self._get_cpu_info(),
            "memory": f"{psutil.virtual_memory().total / (1024**3):.1f}GB",
            "gpu": self._get_gpu_info()
        }

    def _get_cpu_info(self) -> str:
        """Get CPU information."""
        return f"{psutil.cpu_count()} cores @ {psutil.cpu_freq().current:.1f}MHz"

    def _get_gpu_info(self) -> str:
        """Get GPU information."""
        try:
            import torch
            if torch.cuda.is_available():
                return torch.cuda.get_device_name()
        except ImportError:
            pass
        return "N/A"

class CompositionBenchmark(Benchmark):
    """
    Benchmark video composition operations.
    """

    def __init__(self, video_paths: List[Path]):
        """
        Initialize composition benchmark.

        Args:
            video_paths: List of test video paths
        """
        super().__init__("composition")
        self.video_paths = [Path(p) for p in video_paths]

    def _run_single(self) -> BenchmarkResult:
        """Run composition benchmark."""
        videos = [fmusvid.load(path) for path in self.video_paths]
        start_time = time.time()
        start_memory = self._get_memory_usage()

        # Menguji operasi komposisi
        grid = fmusvid.grid(videos, rows=2, cols=2)
        grid.save("benchmark_grid.mp4")

        end_time = time.time()
        end_memory = self._get_memory_usage()

        return BenchmarkResult(
            name=self.name,
            operation="grid_composition",
            input_size={
                "num_videos": len(videos),
                "avg_width": sum(v.width for v in videos) // len(videos),
                "avg_height": sum(v.height for v in videos) // len(videos),
                "total_duration": sum(v.duration for v in videos)
            },
            execution_time=end_time - start_time,
            memory_usage=end_memory - start_memory,
            cpu_usage=self._get_cpu_usage(),
            fps=sum(v.duration for v in videos) / (end_time - start_time),
            backend=videos[0].backend.name,
            hardware_info=self._get_hardware_info()
        )

class AIFeaturesBenchmark(Benchmark):
    """
    Benchmark AI-enhanced features.
    """

    def __init__(self, video_path: Path):
        """
        Initialize AI features benchmark.

        Args:
            video_path: Path to test video
        """
        super().__init__("ai_features")
        self.video_path = Path(video_path)

    def _run_single(self) -> BenchmarkResult:
        """Run AI features benchmark."""
        from fmusvid.ai import ObjectDetector, ObjectTracker

        video = fmusvid.load(self.video_path)
        detector = ObjectDetector()
        tracker = ObjectTracker()

        start_time = time.time()
        start_memory = self._get_memory_usage()

        # Menguji fitur AI
        for frame in video.frames():
            detections = detector.detect(frame)
            tracks = tracker.update(detections)
            frame = tracker.annotate_frame(frame, tracks)

        end_time = time.time()
        end_memory = self._get_memory_usage()

        return BenchmarkResult(
            name=self.name,
            operation="object_tracking",
            input_size={
                "width": video.width,
                "height": video.height,
                "duration": video.duration,
                "frames": len(list(video.frames()))
            },
            execution_time=end_time - start_time,
            memory_usage=end_memory - start_memory,
            cpu_usage=self._get_cpu_usage(),
            fps=video.duration / (end_time - start_time),
            backend=video.backend.name,
            hardware_info=self._get_hardware_info()
        )

class BenchmarkRunner:
    """
    Runs and manages benchmarks.
    """

    def __init__(self, output_dir: Path):
        """
        Initialize benchmark runner.

        Args:
            output_dir: Output directory for results
        """
        self.output_dir = Path(output_dir)
        self.output_dir.mkdir(parents=True, exist_ok=True)
        self.benchmarks: List[Benchmark] = []

    def add_benchmark(self, benchmark: Benchmark) -> None:
        """Add benchmark to run."""
        self.benchmarks.append(benchmark)

    def run_all(self, iterations: int = 3) -> None:
        """
        Run all benchmarks.

        Args:
            iterations: Number of iterations per benchmark
        """
        results = []

        for benchmark in self.benchmarks:
            logger.info(f"Running benchmark: {benchmark.name}")
            benchmark_results = benchmark.run(iterations)
            results.extend(benchmark_results)

        self._save_results(results)
        self._generate_report(results)

    def _save_results(self, results: List[BenchmarkResult]) -> None:
        """Save benchmark results to JSON."""
        output_file = self.output_dir / "benchmark_results.json"

        with open(output_file, "w") as f:
            json.dump(
                [asdict(result) for result in results],
                f,
                indent=2
            )

    def _generate_report(self, results: List[BenchmarkResult]) -> None:
        """Generate benchmark report."""
        report = ["# Benchmark Report\n"]

        # Mengelompokkan hasil berdasarkan nama benchmark
        grouped_results = {}
        for result in results:
            if result.name not in grouped_results:
                grouped_results[result.name] = []
            grouped_results[result.name].append(result)

        for name, benchmark_results in grouped_results.items():
            report.append(f"## {name}\n")

            # Menghitung statistik
            times = [r.execution_time for r in benchmark_results]
            memories = [r.memory_usage for r in benchmark_results]
            fpses = [r.fps for r in benchmark_results]

            report.extend([
                f"- Average execution time: {statistics.mean(times):.3f}s",
                f"- Average memory usage: {statistics.mean(memories):.1f}MB",
                f"- Average FPS: {statistics.mean(fpses):.1f}",
                f"- Backend: {benchmark_results[0].backend}",
                ""
            ])

        report_file = self.output_dir / "benchmark_report.md"
        report_file.write_text("\n".join(report))

def compare_results(
    current_results: Path,
    baseline_results: Path,
    threshold: float = 0.1
) -> bool:
    """
    Compare benchmark results with baseline.

    Args:
        current_results: Path to current results
        baseline_results: Path to baseline results
        threshold: Maximum allowed performance regression

    Returns:
        True if results are within threshold
    """
    with open(current_results) as f:
        current = json.load(f)
    with open(baseline_results) as f:
        baseline = json.load(f)

    # Mengelompokkan hasil berdasarkan nama dan operasi
    def group_results(results):
        grouped = {}
        for result in results:
            key = (result["name"], result["operation"])
            if key not in grouped:
                grouped[key] = []
            grouped[key].append(result)
        return grouped

    current_grouped = group_results(current)
    baseline_grouped = group_results(baseline)

    # Membandingkan metrik
    all_within_threshold = True
    comparisons = []

    for key in current_grouped:
        if key not in baseline_grouped:
            continue

        current_times = [r["execution_time"] for r in current_grouped[key]]
        baseline_times = [r["execution_time"] for r in baseline_grouped[key]]

        current_avg = statistics.mean(current_times)
        baseline_avg = statistics.mean(baseline_times)

        regression = (current_avg - baseline_avg) / baseline_avg
        within_threshold = abs(regression) <= threshold

        if not within_threshold:
            all_within_threshold = False

        comparisons.append({
            "name": key[0],
            "operation": key[1],
            "baseline_time": baseline_avg,
            "current_time": current_avg,
            "regression": regression,
            "passed": within_threshold
        })

    # Membuat laporan perbandingan
    report = ["# Benchmark Comparison Report\n"]

    for comp in comparisons:
        status = "✅" if comp["passed"] else "❌"
        report.extend([
            f"## {comp['name']} - {comp['operation']} {status}\n",
            f"- Baseline: {comp['baseline_time']:.3f}s",
            f"- Current: {comp['current_time']:.3f}s",
            f"- Change: {comp['regression']*100:+.1f}%\n"
        ])

    report_file = Path("benchmark_comparison.md")
    report_file.write_text("\n".join(report))

    return all_within_threshold

def run_benchmarks(
    test_data_dir: Path,
    output_dir: Path,
    iterations: int = 3
) -> None:
    """
    Run all benchmarks.

    Args:
        test_data_dir: Directory containing test videos
        output_dir: Output directory for results
        iterations: Number of iterations per benchmark
    """
    runner = BenchmarkRunner(output_dir)

    # Menambahkan benchmark
    runner.add_benchmark(
        BasicOperationsBenchmark(test_data_dir / "sample.mp4")
    )

    runner.add_benchmark(
        CompositionBenchmark([
            test_data_dir / "clip1.mp4",
            test_data_dir / "clip2.mp4",
            test_data_dir / "clip3.mp4",
            test_data_dir / "clip4.mp4"
        ])
    )

    runner.add_benchmark(
        AIFeaturesBenchmark(test_data_dir / "sample.mp4")
    )

    # Menjalankan benchmark
    runner.run_all(iterations)
