"""
Performance profiling infrastructure for Trax platform.

Provides comprehensive performance measurement, analysis, and optimization
capabilities for the transcription pipeline.
"""

import asyncio
import functools
import json
import logging
import platform
import time
from dataclasses import dataclass, asdict
from datetime import datetime, timezone
from pathlib import Path
from typing import Any, Callable, Dict, List, Optional, Protocol
from functools import lru_cache

import pandas as pd
import psutil
import torch

from .performance import PerformanceMetrics

logger = logging.getLogger(__name__)


@dataclass
class BenchmarkData:
    """Structured benchmark data for performance analysis."""
    operation_name: str
    batch_size: int
    duration_seconds: float
    peak_memory_mb: float
    throughput_items_per_second: float
    timestamp: datetime
    
    def to_dict(self) -> Dict[str, Any]:
        """Convert to dictionary for serialization."""
        data = asdict(self)
        data['timestamp'] = self.timestamp.isoformat()
        return data
    
    @classmethod
    def from_dict(cls, data: Dict[str, Any]) -> 'BenchmarkData':
        """Create from dictionary."""
        data['timestamp'] = datetime.fromisoformat(data['timestamp'])
        return cls(**data)


class MemoryTracker:
    """Track memory usage during operations."""
    
    def __init__(self):
        self.process = psutil.Process()
        self.peak_memory = 0.0
    
    def get_current_memory_mb(self) -> float:
        """Get current memory usage in MB."""
        memory_info = self.process.memory_info()
        current_memory = memory_info.rss / (1024 * 1024)
        self.peak_memory = max(self.peak_memory, current_memory)
        return current_memory
    
    def get_peak_memory_mb(self) -> float:
        """Get peak memory usage in MB."""
        return self.peak_memory
    
    def reset_peak(self) -> None:
        """Reset peak memory tracking."""
        self.peak_memory = 0.0


def timing_decorator(func: Callable) -> Callable:
    """Decorator to measure function execution time."""
    @functools.wraps(func)
    def wrapper(*args, **kwargs):
        start_time = time.time()
        result = func(*args, **kwargs)
        end_time = time.time()
        
        logger.debug(f"{func.__name__} took {end_time - start_time:.3f} seconds")
        return result
    
    return wrapper


class SystemInfoCollector:
    """Collect system information for benchmarking context."""
    
    def collect_system_info(self) -> Dict[str, Any]:
        """Collect comprehensive system information."""
        try:
            memory = psutil.virtual_memory()
            
            system_info = {
                'cpu_model': platform.processor(),
                'architecture': platform.machine(),
                'cpu_cores': psutil.cpu_count(logical=True),
                'total_memory_gb': memory.total / (1024**3),
                'available_memory_gb': memory.available / (1024**3),
                'python_version': platform.python_version(),
                'platform': platform.platform(),
                'timestamp': datetime.now(timezone.utc).isoformat()
            }
            
            # Add GPU information if available
            if torch.cuda.is_available():
                system_info.update({
                    'gpu_count': torch.cuda.device_count(),
                    'gpu_name': torch.cuda.get_device_name(0),
                    'gpu_memory_gb': torch.cuda.get_device_properties(0).total_memory / (1024**3)
                })
            
            return system_info
            
        except Exception as e:
            logger.error(f"Error collecting system info: {e}")
            return {'error': str(e)}


class BenchmarkDataStore:
    """Store and retrieve benchmark data."""
    
    def __init__(self, storage_path: Path):
        self.storage_path = Path(storage_path)
        self.storage_path.mkdir(parents=True, exist_ok=True)
    
    def store_benchmark_data(self, data: BenchmarkData) -> None:
        """Store benchmark data to file."""
        file_path = self.storage_path / f"{data.operation_name}_benchmarks.json"
        
        # Load existing data
        existing_data = []
        if file_path.exists():
            with open(file_path, 'r') as f:
                existing_data = json.load(f)
        
        # Add new data
        existing_data.append(data.to_dict())
        
        # Save updated data
        with open(file_path, 'w') as f:
            json.dump(existing_data, f, indent=2)
    
    def get_benchmark_data(self, operation_name: str) -> List[BenchmarkData]:
        """Retrieve benchmark data for an operation."""
        file_path = self.storage_path / f"{operation_name}_benchmarks.json"
        
        if not file_path.exists():
            return []
        
        with open(file_path, 'r') as f:
            data_list = json.load(f)
        
        return [BenchmarkData.from_dict(data) for data in data_list]


class MetricsAggregator:
    """Aggregate performance metrics for analysis."""
    
    def aggregate_metrics(self, metrics: List[PerformanceMetrics]) -> Dict[str, float]:
        """Aggregate performance metrics."""
        if not metrics:
            return {}
        
        total_operations = len(metrics)
        total_errors = sum(m.error_count for m in metrics)
        total_successes = sum(m.success_count for m in metrics)
        total_items = sum(m.total_count for m in metrics)
        
        return {
            'avg_duration_seconds': sum(m.duration_seconds for m in metrics) / total_operations,
            'avg_memory_peak_mb': sum(m.memory_peak_mb for m in metrics) / total_operations,
            'avg_cpu_peak_percent': sum(m.cpu_peak_percent for m in metrics) / total_operations,
            'avg_throughput_items_per_second': sum(m.throughput_items_per_second for m in metrics) / total_operations,
            'total_operations': total_operations,
            'total_errors': total_errors,
            'total_successes': total_successes,
            'success_rate': (total_successes / total_items * 100) if total_items > 0 else 0.0
        }


class PerformanceThresholdMonitor:
    """Monitor performance against defined thresholds."""
    
    def __init__(self, max_duration_seconds: float = 30.0, 
                 max_memory_mb: float = 2048.0, 
                 max_cpu_percent: float = 90.0):
        self.max_duration_seconds = max_duration_seconds
        self.max_memory_mb = max_memory_mb
        self.max_cpu_percent = max_cpu_percent
    
    def check_thresholds(self, metrics: PerformanceMetrics) -> List[str]:
        """Check metrics against thresholds and return violations."""
        violations = []
        
        if metrics.duration_seconds > self.max_duration_seconds:
            violations.append(f"Duration exceeded threshold: {metrics.duration_seconds:.2f}s > {self.max_duration_seconds}s")
        
        if metrics.memory_peak_mb > self.max_memory_mb:
            violations.append(f"Memory exceeded threshold: {metrics.memory_peak_mb:.2f}MB > {self.max_memory_mb}MB")
        
        if metrics.cpu_peak_percent > self.max_cpu_percent:
            violations.append(f"CPU exceeded threshold: {metrics.cpu_peak_percent:.2f}% > {self.max_cpu_percent}%")
        
        return violations


class PerformanceBenchmark:
    """Comprehensive performance benchmarking system."""
    
    def __init__(self, model_manager: Any, diarization_manager: Any, domain_adapter: Any):
        self.model_manager = model_manager
        self.diarization_manager = diarization_manager
        self.domain_adapter = domain_adapter
        self.results: Dict[str, Any] = {}
        self.memory_tracker = MemoryTracker()
        self.system_info = SystemInfoCollector().collect_system_info()
    
    async def benchmark_transcription(self, audio_files: List[str], 
                                    batch_sizes: List[int] = [1, 2, 4, 8], 
                                    device: str = 'cuda') -> pd.DataFrame:
        """Benchmark transcription performance across different batch sizes."""
        metrics = []
        
        for batch_size in batch_sizes:
            start_time = time.time()
            peak_memory = 0.0
            
            # Process in batches
            for i in range(0, len(audio_files), batch_size):
                batch = audio_files[i:i + batch_size]
                
                # Reset memory tracking
                if device == 'cuda' and torch.cuda.is_available():
                    torch.cuda.reset_peak_memory_stats()
                    torch.cuda.empty_cache()
                
                self.memory_tracker.reset_peak()
                
                # Run transcription
                await self.model_manager.transcribe_batch(batch)
                
                # Track peak memory
                if device == 'cuda' and torch.cuda.is_available():
                    current_peak = torch.cuda.max_memory_allocated() / (1024 ** 3)  # GB
                else:
                    current_peak = self.memory_tracker.get_peak_memory_mb() / 1024  # GB
                
                peak_memory = max(peak_memory, current_peak)
            
            total_time = time.time() - start_time
            metrics.append({
                'batch_size': batch_size,
                'total_time': total_time,
                'throughput': len(audio_files) / total_time,
                'peak_memory_gb': peak_memory
            })
        
        self.results['transcription'] = pd.DataFrame(metrics)
        return self.results['transcription']
    
    async def benchmark_diarization(self, audio_files: List[str]) -> Dict[str, float]:
        """Benchmark speaker diarization performance."""
        start_time = time.time()
        process = psutil.Process()
        base_memory = process.memory_info().rss / (1024 ** 2)  # MB
        
        for audio_file in audio_files:
            await self.diarization_manager.process_audio(audio_file)
        
        total_time = time.time() - start_time
        peak_memory = (process.memory_info().rss / (1024 ** 2)) - base_memory
        
        self.results['diarization'] = {
            'total_time': total_time,
            'per_file_avg': total_time / len(audio_files),
            'peak_memory_mb': peak_memory
        }
        return self.results['diarization']
    
    async def benchmark_end_to_end(self, audio_files: List[str]) -> Dict[str, float]:
        """Benchmark complete pipeline performance."""
        start_time = time.time()
        
        # Transcription phase
        transcription_start = time.time()
        transcripts = await self.model_manager.transcribe_batch(audio_files)
        transcription_time = time.time() - transcription_start
        
        # Diarization phase
        diarization_start = time.time()
        diarization_results = []
        for audio_file in audio_files:
            result = await self.diarization_manager.process_audio(audio_file)
            diarization_results.append(result)
        diarization_time = time.time() - diarization_start
        
        # Adaptation phase
        adaptation_start = time.time()
        adapted_transcripts = []
        for transcript in transcripts:
            adapted = await self.domain_adapter.adapt_transcript(transcript)
            adapted_transcripts.append(adapted)
        adaptation_time = time.time() - adaptation_start
        
        total_time = time.time() - start_time
        
        self.results['end_to_end'] = {
            'total_processing_time': total_time,
            'transcription_time': transcription_time,
            'diarization_time': diarization_time,
            'adaptation_time': adaptation_time,
            'peak_memory_usage': self.memory_tracker.get_peak_memory_mb(),
            'throughput': len(audio_files) / total_time
        }
        
        return self.results['end_to_end']
    
    def generate_report(self, output_path: str = "benchmark_report.html") -> str:
        """Generate comprehensive performance report with visualizations."""
        # Create HTML report
        html_content = f"""
        <html>
        <head>
            <title>Performance Benchmark Report</title>
            <style>
                body {{ font-family: Arial, sans-serif; margin: 20px; }}
                .metric {{ margin-bottom: 20px; }}
                table {{ border-collapse: collapse; width: 100%; }}
                th, td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                th {{ background-color: #f2f2f2; }}
                .system-info {{ background-color: #f9f9f9; padding: 15px; }}
            </style>
        </head>
        <body>
            <h1>Performance Benchmark Report</h1>
            <div class="system-info">
                <h2>System Information</h2>
                <pre>{json.dumps(self.system_info, indent=2)}</pre>
            </div>
            <div class="metric">
                <h2>Transcription Performance</h2>
                {self.results.get('transcription', pd.DataFrame()).to_html() if 'transcription' in self.results else '<p>No transcription data available</p>'}
            </div>
            <div class="metric">
                <h2>Diarization Performance</h2>
                <pre>{json.dumps(self.results.get('diarization', {}), indent=2)}</pre>
            </div>
            <div class="metric">
                <h2>End-to-End Performance</h2>
                <pre>{json.dumps(self.results.get('end_to_end', {}), indent=2)}</pre>
            </div>
        </body>
        </html>
        """
        
        with open(output_path, 'w') as f:
            f.write(html_content)
        
        return output_path


def create_performance_benchmark(model_manager: Any, diarization_manager: Any, 
                                domain_adapter: Any) -> PerformanceBenchmark:
    """Create a performance benchmark instance."""
    return PerformanceBenchmark(model_manager, diarization_manager, domain_adapter)
