import time
import psutil
import threading
import os
from dataclasses import dataclass
from typing import Dict, Any, List
from src.utils.logger import get_logger

logger = get_logger(__name__)

@dataclass
class PerformanceMetrics:
    start_time: float
    end_time: float
    duration: float
    cpu_usage: List[float]
    memory_usage: List[float]
    peak_memory: float
    operation_name: str

class PerformanceMonitor:
    def __init__(self, operation_name: str):
        self.operation_name = operation_name
        self.start_time = None
        self.end_time = None
        self.cpu_usage = []
        self.memory_usage = []
        self.monitoring = False
        self.monitor_thread = None
        # Get current process for accurate memory tracking
        self.process = psutil.Process(os.getpid())
        self.baseline_memory = self.process.memory_info().rss / 1024 / 1024  # MB
    
    def __enter__(self):
        """Context manager entry - start monitoring"""
        self.start_monitoring()
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        """Context manager exit - stop monitoring"""
        return self.stop_monitoring()
    
    def start_monitoring(self):
        """Start monitoring system resources"""
        self.start_time = time.time()
        self.monitoring = True
        self.cpu_usage = []
        self.memory_usage = []
        
        # Start monitoring thread
        self.monitor_thread = threading.Thread(target=self._monitor_resources)
        self.monitor_thread.daemon = True
        self.monitor_thread.start()
        
        logger.info(f"Started monitoring: {self.operation_name}")
    
    def stop_monitoring(self):
        """Stop resource monitoring and return metrics"""
        self.end_time = time.time()
        self.monitoring = False
        
        if self.monitor_thread:
            self.monitor_thread.join()
        
        # Calculate memory increase from baseline
        if self.memory_usage:
            current_memory = self.process.memory_info().rss / 1024 / 1024
            memory_increase = max(self.memory_usage) - self.baseline_memory
            peak_memory = max(memory_increase, 0)  # Ensure non-negative
        else:
            peak_memory = 0
        
        metrics = PerformanceMetrics(
            start_time=self.start_time,
            end_time=self.end_time,
            duration=self.end_time - self.start_time,
            cpu_usage=self.cpu_usage,
            memory_usage=self.memory_usage,
            peak_memory=peak_memory,  # Memory increase from baseline
            operation_name=self.operation_name
        )
        
        self._log_metrics(metrics)
        return metrics
    
    def _monitor_resources(self):
        """Monitor CPU and memory usage in background"""
        while self.monitoring:
            try:
                # Get process-specific CPU usage
                cpu_percent = self.process.cpu_percent()
                
                # Get process-specific memory usage (RSS - Resident Set Size)
                memory_info = self.process.memory_info()
                memory_mb = memory_info.rss / 1024 / 1024  # Convert to MB
                
                self.cpu_usage.append(cpu_percent)
                self.memory_usage.append(memory_mb)
                
                time.sleep(0.1)  # Monitor every 0.1 seconds for more precision
            except Exception as e:
                logger.error(f"Error monitoring resources: {str(e)}")
                break
    
    def _log_metrics(self, metrics: PerformanceMetrics):
        """Log performance metrics"""
        logger.info(f"Performance Report for: {metrics.operation_name}")
        logger.info(f"  Duration: {metrics.duration:.2f} seconds")
        logger.info(f"  Peak Memory: {metrics.peak_memory:.2f} MB")
        if metrics.cpu_usage:
            avg_cpu = sum(metrics.cpu_usage) / len(metrics.cpu_usage)
            logger.info(f"  Average CPU: {avg_cpu:.1f}%")
    
    def get_current_memory_usage(self):
        """Get current process memory usage in MB"""
        return self.process.memory_info().rss / 1024 / 1024

def get_system_info():
    """Get system information for performance context"""
    process = psutil.Process()
    return {
        "cpu_count": psutil.cpu_count(),
        "memory_total_gb": psutil.virtual_memory().total / 1024 / 1024 / 1024,
        "process_memory_mb": process.memory_info().rss / 1024 / 1024
    }