"""
Monitoring and observability module for Toxo.
"""

import asyncio
import time
import psutil
import threading
from typing import Dict, Any, List, Optional, Callable
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from collections import defaultdict, deque
import json
from pathlib import Path
from contextlib import contextmanager

# Prometheus metrics
try:
    from prometheus_client import Counter, Histogram, Gauge, Info, CollectorRegistry, generate_latest, CONTENT_TYPE_LATEST
    PROMETHEUS_AVAILABLE = True
except ImportError:
    PROMETHEUS_AVAILABLE = False

# OpenTelemetry
try:
    from opentelemetry import trace, metrics
    from opentelemetry.exporter.prometheus import PrometheusMetricReader
    from opentelemetry.sdk.trace import TracerProvider
    from opentelemetry.sdk.metrics import MeterProvider
    OPENTELEMETRY_AVAILABLE = True
except ImportError:
    OPENTELEMETRY_AVAILABLE = False

from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError


@dataclass
class HealthCheck:
    """Represents a health check."""
    name: str
    status: str  # "healthy", "unhealthy", "degraded"
    message: str
    last_checked: datetime
    response_time: float
    metadata: Dict[str, Any] = None
    
    def __post_init__(self):
        if self.metadata is None:
            self.metadata = {}


@dataclass
class MetricData:
    """Represents a metric data point."""
    name: str
    value: float
    timestamp: datetime
    labels: Dict[str, str] = None
    
    def __post_init__(self):
        if self.labels is None:
            self.labels = {}


class MetricsCollector:
    """
    Collects and manages metrics for Toxo.
    """
    
    def __init__(self):
        self.logger = get_logger(__name__)
        self.metrics: Dict[str, deque] = defaultdict(lambda: deque(maxlen=1000))
        self.counters: Dict[str, int] = defaultdict(int)
        
        # Prometheus metrics
        if PROMETHEUS_AVAILABLE:
            self.registry = CollectorRegistry()
            self._init_prometheus_metrics()
        
        # OpenTelemetry
        if OPENTELEMETRY_AVAILABLE:
            self._init_opentelemetry()
    
    def _init_prometheus_metrics(self):
        """Initialize Prometheus metrics."""
        self.prom_request_counter = Counter(
            'toxo_requests_total',
            'Total number of requests processed',
            ['endpoint', 'status'],
            registry=self.registry
        )
        
        self.prom_request_duration = Histogram(
            'toxo_request_duration_seconds',
            'Request duration in seconds',
            ['endpoint'],
            registry=self.registry
        )
        
        self.prom_memory_usage = Gauge(
            'toxo_memory_usage_bytes',
            'Memory usage in bytes',
            ['type'],
            registry=self.registry
        )
        
        self.prom_active_sessions = Gauge(
            'toxo_active_sessions',
            'Number of active sessions',
            registry=self.registry
        )
        
        self.prom_model_inference_time = Histogram(
            'toxo_model_inference_seconds',
            'Model inference time in seconds',
            ['model_type'],
            registry=self.registry
        )
    
    def _init_opentelemetry(self):
        """Initialize OpenTelemetry."""
        # Set up tracing
        trace.set_tracer_provider(TracerProvider())
        self.tracer = trace.get_tracer(__name__)
        
        # Set up metrics
        metrics.set_meter_provider(MeterProvider())
        self.meter = metrics.get_meter(__name__)
    
    def record_metric(self, name: str, value: float, labels: Optional[Dict[str, str]] = None):
        """
        Record a metric value.
        
        Args:
            name: Metric name
            value: Metric value
            labels: Optional labels
        """
        metric = MetricData(
            name=name,
            value=value,
            timestamp=datetime.now(),
            labels=labels or {}
        )
        
        self.metrics[name].append(metric)
        
        # Update Prometheus metrics
        if PROMETHEUS_AVAILABLE:
            self._update_prometheus_metric(name, value, labels)
    
    def _update_prometheus_metric(self, name: str, value: float, labels: Optional[Dict[str, str]]):
        """Update Prometheus metrics."""
        try:
            if name == "requests_total":
                self.prom_request_counter.labels(**labels).inc(value)
            elif name == "request_duration":
                self.prom_request_duration.labels(**labels).observe(value)
            elif name == "memory_usage":
                self.prom_memory_usage.labels(**labels).set(value)
            elif name == "active_sessions":
                self.prom_active_sessions.set(value)
            elif name == "model_inference_time":
                self.prom_model_inference_time.labels(**labels).observe(value)
        except Exception as e:
            self.logger.warning(f"Failed to update Prometheus metric {name}: {e}")
    
    def increment_counter(self, name: str, labels: Optional[Dict[str, str]] = None):
        """Increment a counter metric."""
        self.counters[name] += 1
        self.record_metric(name, 1, labels)
    
    def get_metric_summary(self, name: str, window_minutes: int = 60) -> Dict[str, Any]:
        """
        Get summary statistics for a metric.
        
        Args:
            name: Metric name
            window_minutes: Time window in minutes
            
        Returns:
            Summary statistics
        """
        if name not in self.metrics:
            return {"error": f"Metric {name} not found"}
        
        cutoff_time = datetime.now() - timedelta(minutes=window_minutes)
        recent_metrics = [
            m for m in self.metrics[name]
            if m.timestamp >= cutoff_time
        ]
        
        if not recent_metrics:
            return {"count": 0}
        
        values = [m.value for m in recent_metrics]
        
        return {
            "count": len(values),
            "min": min(values),
            "max": max(values),
            "avg": sum(values) / len(values),
            "total": sum(values),
            "window_minutes": window_minutes
        }
    
    def get_prometheus_metrics(self) -> str:
        """Get Prometheus formatted metrics."""
        if not PROMETHEUS_AVAILABLE:
            return "# Prometheus not available\n"
        
        return generate_latest(self.registry).decode('utf-8')


class HealthChecker:
    """
    Manages health checks for Toxo components.
    """
    
    def __init__(self):
        self.logger = get_logger(__name__)
        self.health_checks: Dict[str, Callable] = {}
        self.health_status: Dict[str, HealthCheck] = {}
        self.check_interval = 30  # seconds
        self.running = False
        self._check_task = None
    
    def register_health_check(self, name: str, check_func: Callable):
        """
        Register a health check function.
        
        Args:
            name: Health check name
            check_func: Health check function (should return bool)
        """
        self.health_checks[name] = check_func
        self.logger.info(f"Registered health check: {name}")
    
    async def start_monitoring(self):
        """Start continuous health monitoring."""
        self.running = True
        self._check_task = asyncio.create_task(self._monitor_loop())
        self.logger.info("Health monitoring started")
    
    async def stop_monitoring(self):
        """Stop health monitoring."""
        self.running = False
        if self._check_task:
            self._check_task.cancel()
        self.logger.info("Health monitoring stopped")
    
    async def _monitor_loop(self):
        """Main monitoring loop."""
        while self.running:
            try:
                await self.run_all_checks()
                await asyncio.sleep(self.check_interval)
            except Exception as e:
                self.logger.error(f"Health monitoring error: {e}")
                await asyncio.sleep(5)
    
    async def run_all_checks(self):
        """Run all registered health checks."""
        for name, check_func in self.health_checks.items():
            await self.run_single_check(name, check_func)
    
    async def run_single_check(self, name: str, check_func: Callable):
        """Run a single health check."""
        start_time = time.time()
        
        try:
            if asyncio.iscoroutinefunction(check_func):
                result = await check_func()
            else:
                result = check_func()
            
            response_time = time.time() - start_time
            
            if isinstance(result, dict):
                status = result.get("status", "healthy" if result.get("success", True) else "unhealthy")
                message = result.get("message", "")
                metadata = result.get("metadata", {})
            else:
                status = "healthy" if result else "unhealthy"
                message = f"Check {'passed' if result else 'failed'}"
                metadata = {}
            
            self.health_status[name] = HealthCheck(
                name=name,
                status=status,
                message=message,
                last_checked=datetime.now(),
                response_time=response_time,
                metadata=metadata
            )
            
        except Exception as e:
            self.health_status[name] = HealthCheck(
                name=name,
                status="unhealthy",
                message=f"Check failed: {str(e)}",
                last_checked=datetime.now(),
                response_time=time.time() - start_time
            )
    
    def get_health_status(self) -> Dict[str, Any]:
        """Get overall health status."""
        if not self.health_status:
            return {"status": "unknown", "checks": []}
        
        all_healthy = all(
            check.status == "healthy"
            for check in self.health_status.values()
        )
        
        any_degraded = any(
            check.status == "degraded"
            for check in self.health_status.values()
        )
        
        overall_status = "healthy"
        if not all_healthy:
            overall_status = "degraded" if any_degraded else "unhealthy"
        
        return {
            "status": overall_status,
            "timestamp": datetime.now().isoformat(),
            "checks": [asdict(check) for check in self.health_status.values()]
        }


class SystemMonitor:
    """
    Monitors system resources and performance.
    """
    
    def __init__(self, metrics_collector: MetricsCollector):
        self.metrics = metrics_collector
        self.logger = get_logger(__name__)
        self.monitoring = False
        self._monitor_task = None
    
    async def start_monitoring(self, interval: int = 30):
        """
        Start system monitoring.
        
        Args:
            interval: Monitoring interval in seconds
        """
        self.monitoring = True
        self._monitor_task = asyncio.create_task(self._monitor_loop(interval))
        self.logger.info("System monitoring started")
    
    async def stop_monitoring(self):
        """Stop system monitoring."""
        self.monitoring = False
        if self._monitor_task:
            self._monitor_task.cancel()
        self.logger.info("System monitoring stopped")
    
    async def _monitor_loop(self, interval: int):
        """Main monitoring loop."""
        while self.monitoring:
            try:
                await self._collect_system_metrics()
                await asyncio.sleep(interval)
            except Exception as e:
                self.logger.error(f"System monitoring error: {e}")
                await asyncio.sleep(5)
    
    async def _collect_system_metrics(self):
        """Collect system metrics."""
        # CPU usage
        cpu_percent = psutil.cpu_percent()
        self.metrics.record_metric("cpu_usage_percent", cpu_percent)
        
        # Memory usage
        memory = psutil.virtual_memory()
        self.metrics.record_metric("memory_usage_bytes", memory.used, {"type": "used"})
        self.metrics.record_metric("memory_usage_bytes", memory.available, {"type": "available"})
        self.metrics.record_metric("memory_usage_percent", memory.percent)
        
        # Disk usage
        disk = psutil.disk_usage('/')
        self.metrics.record_metric("disk_usage_bytes", disk.used, {"type": "used"})
        self.metrics.record_metric("disk_usage_bytes", disk.free, {"type": "free"})
        self.metrics.record_metric("disk_usage_percent", (disk.used / disk.total) * 100)
        
        # Network I/O
        network = psutil.net_io_counters()
        self.metrics.record_metric("network_bytes", network.bytes_sent, {"direction": "sent"})
        self.metrics.record_metric("network_bytes", network.bytes_recv, {"direction": "received"})


class ToxoMonitor:
    """
    Main monitoring system for Toxo.
    """
    
    def __init__(self):
        self.logger = get_logger(__name__)
        self.metrics = MetricsCollector()
        self.health_checker = HealthChecker()
        self.system_monitor = SystemMonitor(self.metrics)
        
        # Application metrics
        self.session_count = 0
        self.request_count = 0
        
        # Register default health checks
        self._register_default_health_checks()
    
    def _register_default_health_checks(self):
        """Register default health checks."""
        self.health_checker.register_health_check("memory", self._check_memory_usage)
        self.health_checker.register_health_check("disk", self._check_disk_space)
        self.health_checker.register_health_check("system", self._check_system_health)
    
    def _check_memory_usage(self) -> Dict[str, Any]:
        """Check memory usage."""
        memory = psutil.virtual_memory()
        
        if memory.percent > 90:
            return {
                "success": False,
                "status": "unhealthy",
                "message": f"High memory usage: {memory.percent}%",
                "metadata": {"memory_percent": memory.percent}
            }
        elif memory.percent > 80:
            return {
                "success": True,
                "status": "degraded",
                "message": f"Moderate memory usage: {memory.percent}%",
                "metadata": {"memory_percent": memory.percent}
            }
        else:
            return {
                "success": True,
                "status": "healthy",
                "message": f"Memory usage normal: {memory.percent}%",
                "metadata": {"memory_percent": memory.percent}
            }
    
    def _check_disk_space(self) -> Dict[str, Any]:
        """Check disk space."""
        disk = psutil.disk_usage('/')
        usage_percent = (disk.used / disk.total) * 100
        
        if usage_percent > 95:
            return {
                "success": False,
                "status": "unhealthy",
                "message": f"Critical disk usage: {usage_percent:.1f}%",
                "metadata": {"disk_percent": usage_percent}
            }
        elif usage_percent > 85:
            return {
                "success": True,
                "status": "degraded",
                "message": f"High disk usage: {usage_percent:.1f}%",
                "metadata": {"disk_percent": usage_percent}
            }
        else:
            return {
                "success": True,
                "status": "healthy",
                "message": f"Disk usage normal: {usage_percent:.1f}%",
                "metadata": {"disk_percent": usage_percent}
            }
    
    def _check_system_health(self) -> Dict[str, Any]:
        """Check overall system health."""
        cpu_percent = psutil.cpu_percent()
        
        if cpu_percent > 95:
            return {
                "success": False,
                "status": "unhealthy",
                "message": f"Critical CPU usage: {cpu_percent}%",
                "metadata": {"cpu_percent": cpu_percent}
            }
        elif cpu_percent > 80:
            return {
                "success": True,
                "status": "degraded",
                "message": f"High CPU usage: {cpu_percent}%",
                "metadata": {"cpu_percent": cpu_percent}
            }
        else:
            return {
                "success": True,
                "status": "healthy",
                "message": f"System running normally: {cpu_percent}% CPU",
                "metadata": {"cpu_percent": cpu_percent}
            }
    
    async def start_monitoring(self):
        """Start all monitoring components."""
        await self.health_checker.start_monitoring()
        await self.system_monitor.start_monitoring()
        self.logger.info("Toxo monitoring started")
    
    async def stop_monitoring(self):
        """Stop all monitoring components."""
        await self.health_checker.stop_monitoring()
        await self.system_monitor.stop_monitoring()
        self.logger.info("Toxo monitoring stopped")
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "session_count": getattr(self, 'session_count', 0),
            "request_count": getattr(self, 'request_count', 0),
            "metrics_collector_available": getattr(self, 'metrics', None) is not None,
            "health_checker_available": getattr(self, 'health_checker', None) is not None,
            "system_monitor_available": getattr(self, 'system_monitor', None) is not None,
            "system_type": "ToxoMonitor"
        }
    
    def record_request(self, endpoint: str, duration: float, status: str = "success"):
        """Record a request metric."""
        self.request_count += 1
        self.metrics.record_metric("requests_total", 1, {"endpoint": endpoint, "status": status})
        self.metrics.record_metric("request_duration", duration, {"endpoint": endpoint})
    
    def record_session_start(self):
        """Record session start."""
        self.session_count += 1
        self.metrics.record_metric("active_sessions", self.session_count)
    
    def record_session_end(self):
        """Record session end."""
        self.session_count = max(0, self.session_count - 1)
        self.metrics.record_metric("active_sessions", self.session_count)
    
    def record_model_inference(self, model_type: str, duration: float):
        """Record model inference time."""
        self.metrics.record_metric("model_inference_time", duration, {"model_type": model_type})
    
    def get_dashboard_data(self) -> Dict[str, Any]:
        """Get data for monitoring dashboard."""
        return {
            "health": self.health_checker.get_health_status(),
            "metrics": {
                "requests": self.metrics.get_metric_summary("requests_total"),
                "response_time": self.metrics.get_metric_summary("request_duration"),
                "memory": self.metrics.get_metric_summary("memory_usage_percent"),
                "cpu": self.metrics.get_metric_summary("cpu_usage_percent"),
                "active_sessions": self.session_count,
                "total_requests": self.request_count
            },
            "timestamp": datetime.now().isoformat()
        }
    
    def export_metrics(self, format: str = "prometheus") -> str:
        """
        Export metrics in specified format.
        
        Args:
            format: Export format ("prometheus", "json")
            
        Returns:
            Formatted metrics data
        """
        if format == "prometheus":
            return self.metrics.get_prometheus_metrics()
        elif format == "json":
            return json.dumps(self.get_dashboard_data(), indent=2)
        else:
            raise ToxoError(f"Unsupported export format: {format}")


# Global monitor instance
_monitor = None

def get_monitor() -> ToxoMonitor:
    """Get global monitor instance."""
    global _monitor
    if _monitor is None:
        _monitor = ToxoMonitor()
    return _monitor 


class PerformanceMonitor:
    """
    Core performance monitoring for Toxo operations.
    
    This is a simplified performance monitor that tracks basic
    metrics for system operations and provides monitoring capabilities.
    """
    
    def __init__(self):
        self.logger = get_logger(__name__)
        self.metrics = MetricsCollector()
        self._start_time = datetime.now()
        self.operation_counts: Dict[str, int] = defaultdict(int)
        self.operation_times: Dict[str, List[float]] = defaultdict(list)
        self.lock = threading.Lock()
        
    def record_operation(self, operation_name: str, duration: float, success: bool = True):
        """Record an operation with its duration and success status."""
        with self.lock:
            self.operation_counts[operation_name] += 1
            self.operation_times[operation_name].append(duration)
            
            # Keep only recent times (last 1000 operations)
            if len(self.operation_times[operation_name]) > 1000:
                self.operation_times[operation_name] = self.operation_times[operation_name][-1000:]
            
            # Record in metrics collector
            self.metrics.record_metric(
                f"{operation_name}_duration",
                duration,
                {"status": "success" if success else "failure"}
            )
            
            if success:
                self.metrics.increment_counter(f"{operation_name}_success")
            else:
                self.metrics.increment_counter(f"{operation_name}_failure")
    
    def get_operation_stats(self, operation_name: str) -> Dict[str, Any]:
        """Get statistics for a specific operation."""
        with self.lock:
            if operation_name not in self.operation_times:
                return {"error": "Operation not found"}
            
            times = self.operation_times[operation_name]
            if not times:
                return {"error": "No timing data available"}
            
            return {
                "operation": operation_name,
                "count": self.operation_counts[operation_name],
                "avg_duration": sum(times) / len(times),
                "min_duration": min(times),
                "max_duration": max(times),
                "recent_duration": times[-1] if times else 0,
                "p95_duration": sorted(times)[int(len(times) * 0.95)] if len(times) > 1 else times[0]
            }
    
    def get_all_stats(self) -> Dict[str, Any]:
        """Get statistics for all operations."""
        with self.lock:
            stats = {}
            for operation in self.operation_counts.keys():
                stats[operation] = self.get_operation_stats(operation)
            
            return {
                "uptime_seconds": (datetime.now() - self._start_time).total_seconds(),
                "operations": stats,
                "total_operations": sum(self.operation_counts.values())
            }
    
    def reset_stats(self):
        """Reset all statistics."""
        with self.lock:
            self.operation_counts.clear()
            self.operation_times.clear()
            self._start_time = datetime.now()
    
    @contextmanager
    def measure_operation(self, operation_name: str):
        """Context manager to measure operation duration."""
        start_time = time.time()
        success = True
        
        try:
            yield
        except Exception:
            success = False
            raise
        finally:
            duration = time.time() - start_time
            self.record_operation(operation_name, duration, success)
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state for serialization."""
        return {
            "health_checks_count": len(getattr(self, 'health_checks', {})),
            "metrics_collector_available": getattr(self, 'metrics_collector', None) is not None,
            "performance_monitor_available": getattr(self, 'performance_monitor', None) is not None,
            "start_time": str(getattr(self, '_start_time', '')),
            "system_type": "ToxoMonitor"
        } 