import asyncio
import time
import psutil
import numpy as np
from datetime import datetime, timedelta
from typing import Dict, List, Optional
from dataclasses import dataclass, field
import torch
import json
import platform
import websockets
from urllib.parse import urlparse

from config_loader import config
from solana_logger import logger

if platform.system() == 'Windows':
    asyncio.set_event_loop_policy(asyncio.WindowsSelectorEventLoopPolicy())

@dataclass
class SystemMetrics:
    timestamp: datetime = field(default_factory=datetime.now)
    cpu_percent: float = 0.0
    memory_percent: float = 0.0
    gpu_memory_used: float = 0.0
    gpu_utilization: float = 0.0
    network_latency: float = 0.0
    disk_usage: float = 0.0
    active_threads: int = 0
    heap_memory: float = 0.0

@dataclass
class NetworkMetrics:
    ws_latency: float = 0.0
    last_trade_time: Optional[datetime] = None
    connection_drops: int = 0
    successful_requests: int = 0
    failed_requests: int = 0

class SystemMonitor:
    def __init__(self):
        self._setup_configuration()
        self._initialize_metrics()
        self._setup_alerting()
        logger.info("System Monitor initialized")

    def _setup_configuration(self):
        """Initialize monitoring configuration"""
        network_config = config.get_network_config()
        self.WS_URL = network_config["ws_url"]
        self.API_KEY = network_config["api_key"]
        
        self.CPU_THRESHOLD = 90.0
        self.MEMORY_THRESHOLD = 90.0
        self.GPU_MEMORY_THRESHOLD = 90.0
        self.LATENCY_THRESHOLD = 500.0
        
        self.SYSTEM_CHECK_INTERVAL = 5
        self.NETWORK_CHECK_INTERVAL = 10
        self.METRICS_LOG_INTERVAL = 300
        self.METRIC_WINDOW = 1000
        
        # WebSocket configuration
        self.WS_CONFIG = {
            "ping_interval": 0.5,
            "ping_timeout": 5,
            "compression": None,
            "max_size": 2**23,
            "close_timeout": 1.0,
            "max_queue": 2**16,
            "read_limit": 2**16,
        }

    def _initialize_metrics(self):
        """Initialize metric tracking"""
        self.system_metrics: List[SystemMetrics] = []
        self.network_metrics = NetworkMetrics()
        self.start_time = datetime.now()
        self.last_alert_time = {}
        self.alert_cooldown = 300
        self.ws = None

    def _setup_alerting(self):
        """Configure alerting thresholds"""
        self.alert_levels = {
            'INFO': 0,
            'WARNING': 1,
            'CRITICAL': 2,
            'EMERGENCY': 3
        }

    async def start(self):
        """Start system monitoring"""
        logger.info("Starting System Monitor")
        try:
            await asyncio.gather(
                self._monitor_system_resources(),
                self._monitor_network_health(),
                self._log_metrics_periodically()
            )
        except Exception as e:
            logger.error(f"Monitor error: {str(e)}")

    async def _monitor_system_resources(self):
        """Monitor system resources"""
        while True:
            try:
                metrics = SystemMetrics(
                    timestamp=datetime.now(),
                    cpu_percent=psutil.cpu_percent(interval=1),
                    memory_percent=psutil.virtual_memory().percent,
                    gpu_memory_used=self._get_gpu_memory_usage(),
                    gpu_utilization=self._get_gpu_utilization(),
                    disk_usage=psutil.disk_usage('/').percent,
                    active_threads=len(psutil.Process().threads()),
                    heap_memory=self._get_heap_memory()
                )

                self.system_metrics.append(metrics)
                if len(self.system_metrics) > self.METRIC_WINDOW:
                    self.system_metrics.pop(0)

                await self._check_system_alerts(metrics)
                await asyncio.sleep(self.SYSTEM_CHECK_INTERVAL)

            except Exception as e:
                logger.error(f"Resource monitoring error: {str(e)}")
                await asyncio.sleep(1)

    def _get_gpu_memory_usage(self) -> float:
        """Get GPU memory usage percentage"""
        try:
            if not torch.cuda.is_available():
                return 0.0
            
            memory_allocated = torch.cuda.memory_allocated(0)
            memory_total = torch.cuda.get_device_properties(0).total_memory
            return (memory_allocated / memory_total) * 100
        except Exception as e:
            logger.error(f"GPU memory check error: {str(e)}")
            return 0.0

    def _get_gpu_utilization(self) -> float:
        """Get GPU utilization percentage"""
        try:
            import subprocess
            result = subprocess.run(
                ['nvidia-smi', '--query-gpu=utilization.gpu', '--format=csv,noheader,nounits'],
                capture_output=True,
                text=True
            )
            return float(result.stdout.strip())
        except Exception:
            return 0.0

    def _get_heap_memory(self) -> float:
        """Get Python heap memory usage"""
        try:
            process = psutil.Process()
            return process.memory_info().rss / 1024 / 1024  # MB
        except Exception:
            return 0.0

    async def _monitor_network_health(self):
        """Monitor network health through WebSocket"""
        while True:
            try:
                async with websockets.connect(
                    self.WS_URL,
                    extra_headers={'api-key': self.API_KEY},
                    **self.WS_CONFIG
                ) as ws:
                    # Test WebSocket latency
                    start_time = time.time()
                    await ws.ping()
                    pong_waiter = await ws.ping()
                    await pong_waiter
                    self.network_metrics.ws_latency = (time.time() - start_time) * 1000
                    self.network_metrics.successful_requests += 1

                    # Subscribe to slot updates for continuous monitoring
                    await ws.send(json.dumps({
                        "jsonrpc": "2.0",
                        "id": 1,
                        "method": "slotSubscribe"
                    }))

                    while True:
                        try:
                            await asyncio.wait_for(ws.recv(), timeout=self.WS_CONFIG["ping_timeout"])
                        except asyncio.TimeoutError:
                            pong_waiter = await ws.ping()
                            await asyncio.wait_for(
                                pong_waiter,
                                timeout=self.WS_CONFIG["ping_timeout"]
                            )
                        except Exception:
                            break

            except Exception as e:
                self.network_metrics.failed_requests += 1
                self.network_metrics.connection_drops += 1
                logger.error(f"WebSocket connection error: {str(e)}")
                await asyncio.sleep(1)

            await asyncio.sleep(self.NETWORK_CHECK_INTERVAL)

    async def _check_system_alerts(self, metrics: SystemMetrics):
        """Check system metrics against thresholds"""
        alerts = []
        
        if metrics.cpu_percent > self.CPU_THRESHOLD:
            alerts.append(('WARNING', f'High CPU usage: {metrics.cpu_percent:.1f}%'))
        
        if metrics.memory_percent > self.MEMORY_THRESHOLD:
            alerts.append(('WARNING', f'High memory usage: {metrics.memory_percent:.1f}%'))
        
        if metrics.gpu_memory_used > self.GPU_MEMORY_THRESHOLD:
            alerts.append(('WARNING', f'High GPU memory usage: {metrics.gpu_memory_used:.1f}%'))

        for level, message in alerts:
            await self._send_alert(level, message)

    async def _send_alert(self, level: str, message: str):
        """Send alert with cooldown"""
        current_time = datetime.now()
        alert_key = f"{level}:{message}"

        if alert_key in self.last_alert_time:
            time_since_last = (current_time - self.last_alert_time[alert_key]).total_seconds()
            if time_since_last < self.alert_cooldown:
                return

        self.last_alert_time[alert_key] = current_time
        logger.warning(f"⚠️ {level} Alert: {message}")

    async def _log_metrics_periodically(self):
        """Log system metrics periodically"""
        while True:
            try:
                await self._log_current_metrics()
                await asyncio.sleep(self.METRICS_LOG_INTERVAL)
            except Exception as e:
                logger.error(f"Metrics logging error: {str(e)}")
                await asyncio.sleep(1)

    async def _log_current_metrics(self):
        """Log current system status"""
        try:
            if not self.system_metrics:
                return

            recent_metrics = self.system_metrics[-100:]
            
            avg_metrics = {
                'cpu': np.mean([m.cpu_percent for m in recent_metrics]),
                'memory': np.mean([m.memory_percent for m in recent_metrics]),
                'gpu_memory': np.mean([m.gpu_memory_used for m in recent_metrics]),
                'gpu_util': np.mean([m.gpu_utilization for m in recent_metrics]),
                'ws_latency': self.network_metrics.ws_latency,
            }

            uptime = datetime.now() - self.start_time
            success_rate = self.network_metrics.successful_requests / max(
                self.network_metrics.successful_requests + self.network_metrics.failed_requests, 1
            )

            logger.info(f"""
🖥️ System Status:
Uptime: {str(uptime).split('.')[0]}
CPU: {avg_metrics['cpu']:.1f}%
Memory: {avg_metrics['memory']:.1f}%
GPU Memory: {avg_metrics['gpu_memory']:.1f}%
GPU Utilization: {avg_metrics['gpu_util']:.1f}%
WS Latency: {avg_metrics['ws_latency']:.1f}ms
Network Success: {success_rate:.1%}
Connection Drops: {self.network_metrics.connection_drops}
""")

        except Exception as e:
            logger.error(f"Metrics logging error: {str(e)}")

    async def get_system_health(self) -> Dict:
        """Get current system health status"""
        try:
            if not self.system_metrics:
                return {"status": "initializing"}

            latest = self.system_metrics[-1]
            return {
                "status": "healthy" if all([
                    latest.cpu_percent < self.CPU_THRESHOLD,
                    latest.memory_percent < self.MEMORY_THRESHOLD,
                    latest.gpu_memory_used < self.GPU_MEMORY_THRESHOLD,
                    self.network_metrics.ws_latency < self.LATENCY_THRESHOLD,
                ]) else "degraded",
                "metrics": {
                    "cpu": latest.cpu_percent,
                    "memory": latest.memory_percent,
                    "gpu_memory": latest.gpu_memory_used,
                    "gpu_utilization": latest.gpu_utilization,
                    "ws_latency": self.network_metrics.ws_latency,
                }
            }

        except Exception as e:
            logger.error(f"Health check error: {str(e)}")
            return {"status": "error", "message": str(e)}

    async def cleanup(self):
        """Cleanup resources"""
        pass

async def test_system_monitor():
    """Test system monitor functionality"""
    monitor = SystemMonitor()
    try:
        await monitor.start()
    except KeyboardInterrupt:
        await monitor.cleanup()
        logger.info("System monitor stopped")

if __name__ == "__main__":
    try:
        asyncio.run(test_system_monitor())
    except KeyboardInterrupt:
        pass