"""Performance metrics logging system."""

import asyncio
import json
import logging
import os
import psutil
import time
from contextlib import asynccontextmanager, contextmanager
from dataclasses import dataclass, field
from datetime import datetime, timezone
from enum import Enum
from typing import Any, Callable, Dict, List, Optional, TypeVar, Union
from collections import defaultdict, Counter
import threading

from ..errors import TraxError, ErrorCode, ErrorCategory, ErrorSeverity
from . import get_logger

logger = get_logger(__name__)

T = TypeVar('T')


class MetricType(Enum):
    """Types of metrics that can be collected."""
    TIMING = "timing"
    COUNTER = "counter"
    GAUGE = "gauge"
    HISTOGRAM = "histogram"
    MEMORY = "memory"
    HEALTH = "health"


class MetricSeverity(Enum):
    """Severity levels for metric alerts."""
    INFO = "info"
    WARNING = "warning"
    ERROR = "error"
    CRITICAL = "critical"


@dataclass
class MetricThreshold:
    """Threshold configuration for metric alerts."""
    metric_name: str
    threshold_value: float
    severity: MetricSeverity
    comparison: str  # 'gt', 'lt', 'eq', 'gte', 'lte'
    description: str = ""


@dataclass
class PerformanceMetrics:
    """Performance metrics data structure."""
    operation_name: str
    duration_ms: float
    start_time: datetime
    end_time: datetime
    metadata: Dict[str, Any] = field(default_factory=dict)
    memory_usage_mb: Optional[float] = None
    cpu_usage_percent: Optional[float] = None
    success: bool = True
    error_info: Optional[Dict[str, Any]] = None


class MetricsCollector:
    """Collects and manages performance metrics."""
    
    def __init__(self):
        self.metrics: Dict[str, List[PerformanceMetrics]] = defaultdict(list)
        self.counters: Counter = Counter()
        self.thresholds: List[MetricThreshold] = []
        self.logger = get_logger(f"{__name__}.MetricsCollector")
        self._lock = threading.Lock()
    
    def add_metric(self, metric: PerformanceMetrics):
        """Add a performance metric."""
        with self._lock:
            self.metrics[metric.operation_name].append(metric)
            self._check_thresholds(metric)
    
    def increment_counter(self, counter_name: str, value: int = 1):
        """Increment a counter metric."""
        with self._lock:
            self.counters[counter_name] += value
    
    def get_counter(self, counter_name: str) -> int:
        """Get the current value of a counter."""
        return self.counters[counter_name]
    
    def add_threshold(self, threshold: MetricThreshold):
        """Add a threshold for metric alerts."""
        self.thresholds.append(threshold)
    
    def _check_thresholds(self, metric: PerformanceMetrics):
        """Check if any thresholds have been exceeded."""
        for threshold in self.thresholds:
            if threshold.metric_name != metric.operation_name:
                continue
            
            value = metric.duration_ms
            exceeded = False
            
            if threshold.comparison == 'gt' and value > threshold.threshold_value:
                exceeded = True
            elif threshold.comparison == 'lt' and value < threshold.threshold_value:
                exceeded = True
            elif threshold.comparison == 'eq' and value == threshold.threshold_value:
                exceeded = True
            elif threshold.comparison == 'gte' and value >= threshold.threshold_value:
                exceeded = True
            elif threshold.comparison == 'lte' and value <= threshold.threshold_value:
                exceeded = True
            
            if exceeded:
                self.logger.warning(
                    f"Metric threshold exceeded: {threshold.metric_name} = {value} "
                    f"({threshold.comparison} {threshold.threshold_value}) - {threshold.description}",
                    extra={
                        "metric_name": threshold.metric_name,
                        "current_value": value,
                        "threshold_value": threshold.threshold_value,
                        "severity": threshold.severity.value,
                        "description": threshold.description
                    }
                )
    
    def get_metrics_summary(self, operation_name: Optional[str] = None) -> Dict[str, Any]:
        """Get a summary of collected metrics."""
        with self._lock:
            if operation_name:
                metrics = self.metrics.get(operation_name, [])
            else:
                metrics = [m for metrics_list in self.metrics.values() for m in metrics_list]
            
            if not metrics:
                return {}
            
            durations = [m.duration_ms for m in metrics]
            memory_usage = [m.memory_usage_mb for m in metrics if m.memory_usage_mb is not None]
            cpu_usage = [m.cpu_usage_percent for m in metrics if m.cpu_usage_percent is not None]
            
            summary = {
                "count": len(metrics),
                "success_count": sum(1 for m in metrics if m.success),
                "error_count": sum(1 for m in metrics if not m.success),
                "avg_duration_ms": sum(durations) / len(durations),
                "min_duration_ms": min(durations),
                "max_duration_ms": max(durations),
                "total_duration_ms": sum(durations)
            }
            
            if memory_usage:
                summary.update({
                    "avg_memory_mb": sum(memory_usage) / len(memory_usage),
                    "max_memory_mb": max(memory_usage)
                })
            
            if cpu_usage:
                summary.update({
                    "avg_cpu_percent": sum(cpu_usage) / len(cpu_usage),
                    "max_cpu_percent": max(cpu_usage)
                })
            
            return summary
    
    def export_metrics(self, format: str = "json") -> str:
        """Export metrics in the specified format."""
        with self._lock:
            data = {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "metrics": {
                    name: [m.__dict__ for m in metrics_list]
                    for name, metrics_list in self.metrics.items()
                },
                "counters": dict(self.counters),
                "summary": {
                    name: self.get_metrics_summary(name)
                    for name in self.metrics.keys()
                }
            }
            
            if format.lower() == "json":
                return json.dumps(data, indent=2, default=str)
            else:
                raise ValueError(f"Unsupported export format: {format}")


# Global metrics collector instance
_metrics_collector = MetricsCollector()


def get_metrics_collector() -> MetricsCollector:
    """Get the global metrics collector instance."""
    return _metrics_collector


@contextmanager
def timing_context(operation_name: str, metadata: Optional[Dict[str, Any]] = None):
    """Context manager for timing operations."""
    start_time = datetime.now(timezone.utc)
    start_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
    start_cpu = psutil.cpu_percent()
    
    try:
        yield
        success = True
        error_info = None
    except Exception as e:
        success = False
        error_info = {
            "error_type": type(e).__name__,
            "error_message": str(e)
        }
        raise
    finally:
        end_time = datetime.now(timezone.utc)
        end_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
        end_cpu = psutil.cpu_percent()
        
        duration_ms = (end_time - start_time).total_seconds() * 1000
        memory_usage_mb = end_memory - start_memory
        cpu_usage_percent = end_cpu - start_cpu
        
        metric = PerformanceMetrics(
            operation_name=operation_name,
            duration_ms=duration_ms,
            start_time=start_time,
            end_time=end_time,
            metadata=metadata or {},
            memory_usage_mb=memory_usage_mb,
            cpu_usage_percent=cpu_usage_percent,
            success=success,
            error_info=error_info
        )
        
        _metrics_collector.add_metric(metric)


@asynccontextmanager
async def async_timing_context(operation_name: str, metadata: Optional[Dict[str, Any]] = None):
    """Async context manager for timing operations."""
    start_time = datetime.now(timezone.utc)
    start_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
    start_cpu = psutil.cpu_percent()
    
    try:
        yield
        success = True
        error_info = None
    except Exception as e:
        success = False
        error_info = {
            "error_type": type(e).__name__,
            "error_message": str(e)
        }
        raise
    finally:
        end_time = datetime.now(timezone.utc)
        end_memory = psutil.Process().memory_info().rss / 1024 / 1024  # MB
        end_cpu = psutil.cpu_percent()
        
        duration_ms = (end_time - start_time).total_seconds() * 1000
        memory_usage_mb = end_memory - start_memory
        cpu_usage_percent = end_cpu - start_cpu
        
        metric = PerformanceMetrics(
            operation_name=operation_name,
            duration_ms=duration_ms,
            start_time=start_time,
            end_time=end_time,
            metadata=metadata or {},
            memory_usage_mb=memory_usage_mb,
            cpu_usage_percent=cpu_usage_percent,
            success=success,
            error_info=error_info
        )
        
        _metrics_collector.add_metric(metric)


def timing_decorator(operation_name: Optional[str] = None):
    """Decorator for timing function execution."""
    def decorator(func: Callable) -> Callable:
        def wrapper(*args, **kwargs):
            name = operation_name or f"{func.__module__}.{func.__name__}"
            with timing_context(name):
                return func(*args, **kwargs)
        return wrapper
    return decorator


def async_timing_decorator(operation_name: Optional[str] = None):
    """Async decorator for timing function execution."""
    def decorator(func: Callable) -> Callable:
        async def wrapper(*args, **kwargs):
            name = operation_name or f"{func.__module__}.{func.__name__}"
            async with async_timing_context(name):
                return await func(*args, **kwargs)
        return wrapper
    return decorator


class SystemHealthMonitor:
    """Monitors system health and logs periodic metrics."""
    
    def __init__(self, interval_seconds: int = 60):
        self.interval_seconds = interval_seconds
        self.logger = get_logger(f"{__name__}.SystemHealthMonitor")
        self._running = False
        self._task: Optional[asyncio.Task] = None
    
    async def start(self):
        """Start the health monitoring loop."""
        if self._running:
            return
        
        self._running = True
        self._task = asyncio.create_task(self._monitor_loop())
        self.logger.info(f"Started system health monitoring (interval: {self.interval_seconds}s)")
    
    async def stop(self):
        """Stop the health monitoring loop."""
        if not self._running:
            return
        
        self._running = False
        if self._task:
            self._task.cancel()
            try:
                await self._task
            except asyncio.CancelledError:
                pass
        
        self.logger.info("Stopped system health monitoring")
    
    async def _monitor_loop(self):
        """Main monitoring loop."""
        while self._running:
            try:
                await self._collect_health_metrics()
                await asyncio.sleep(self.interval_seconds)
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"Error in health monitoring loop: {e}")
                await asyncio.sleep(self.interval_seconds)
    
    async def _collect_health_metrics(self):
        """Collect and log system health metrics."""
        try:
            # CPU usage
            cpu_percent = psutil.cpu_percent(interval=1)
            
            # Memory usage
            memory = psutil.virtual_memory()
            
            # Disk usage
            disk = psutil.disk_usage('/')
            
            # Process-specific metrics
            process = psutil.Process()
            process_memory = process.memory_info()
            process_cpu = process.cpu_percent()
            
            health_data = {
                "timestamp": datetime.now(timezone.utc).isoformat(),
                "system": {
                    "cpu_percent": cpu_percent,
                    "memory_percent": memory.percent,
                    "memory_available_gb": memory.available / 1024 / 1024 / 1024,
                    "disk_percent": disk.percent,
                    "disk_free_gb": disk.free / 1024 / 1024 / 1024
                },
                "process": {
                    "cpu_percent": process_cpu,
                    "memory_mb": process_memory.rss / 1024 / 1024,
                    "memory_percent": process.memory_percent(),
                    "threads": process.num_threads(),
                    "open_files": len(process.open_files()),
                    "connections": len(process.connections())
                }
            }
            
            self.logger.info("System health metrics", extra=health_data)
            
            # Check for critical thresholds
            if cpu_percent > 90:
                self.logger.warning(f"High CPU usage: {cpu_percent}%")
            
            if memory.percent > 90:
                self.logger.warning(f"High memory usage: {memory.percent}%")
            
            if disk.percent > 90:
                self.logger.warning(f"High disk usage: {disk.percent}%")
                
        except Exception as e:
            self.logger.error(f"Failed to collect health metrics: {e}")


# Global health monitor instance
_health_monitor = SystemHealthMonitor()


def get_health_monitor() -> SystemHealthMonitor:
    """Get the global health monitor instance."""
    return _health_monitor


# Convenience functions
def log_operation_timing(operation_name: str, duration_ms: float, metadata: Optional[Dict[str, Any]] = None):
    """Log operation timing manually."""
    metric = PerformanceMetrics(
        operation_name=operation_name,
        duration_ms=duration_ms,
        start_time=datetime.now(timezone.utc),
        end_time=datetime.now(timezone.utc),
        metadata=metadata or {}
    )
    _metrics_collector.add_metric(metric)


def increment_operation_counter(operation_name: str, value: int = 1):
    """Increment operation counter."""
    _metrics_collector.increment_counter(operation_name, value)


def get_operation_metrics(operation_name: str) -> Dict[str, Any]:
    """Get metrics for a specific operation."""
    return _metrics_collector.get_metrics_summary(operation_name)


def export_all_metrics(format: str = "json") -> str:
    """Export all collected metrics."""
    return _metrics_collector.export_metrics(format)


async def start_health_monitoring(interval_seconds: int = 60):
    """Start system health monitoring."""
    _health_monitor.interval_seconds = interval_seconds
    await _health_monitor.start()


async def stop_health_monitoring():
    """Stop system health monitoring."""
    await _health_monitor.stop()
