import psutil
import py3nvml
import asyncio
from typing import Dict, Any, List
import numpy as np
from datetime import datetime, timedelta

class HardwareMonitor:
    def __init__(self):
        self.nvml_initialized = False
        try:
            py3nvml.nvmlInit()
            self.nvml_initialized = True
        except:
            pass

    async def get_cpu_thermal_data(self) -> Dict[str, Any]:
        """Get CPU thermal data"""
        try:
            temps = psutil.sensors_temperatures()
            cpu_temps = temps.get('coretemp', [])
            return {
                'current': [temp.current for temp in cpu_temps],
                'high': [temp.high for temp in cpu_temps],
                'critical': [temp.critical for temp in cpu_temps],
                'core_count': len(cpu_temps)
            }
        except:
            return {'error': 'CPU thermal data unavailable'}

    async def get_gpu_thermal_data(self) -> Dict[str, Any]:
        """Get GPU thermal data"""
        if not self.nvml_initialized:
            return {'error': 'NVIDIA tools not available'}

        try:
            gpu_temps = []
            device_count = py3nvml.nvmlDeviceGetCount()
            
            for i in range(device_count):
                handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
                temp = py3nvml.nvmlDeviceGetTemperature(handle, py3nvml.NVML_TEMPERATURE_GPU)
                gpu_temps.append({
                    'index': i,
                    'temperature': temp,
                    'threshold': py3nvml.nvmlDeviceGetTemperatureThreshold(handle, 0)
                })
            
            return {'devices': gpu_temps}
        except:
            return {'error': 'GPU thermal data unavailable'}

    async def get_memory_stats(self) -> Dict[str, Any]:
        """Get comprehensive memory statistics"""
        virtual_memory = psutil.virtual_memory()
        swap = psutil.swap_memory()
        
        memory_stats = {
            'virtual': {
                'total': virtual_memory.total,
                'available': virtual_memory.available,
                'used': virtual_memory.used,
                'percent': virtual_memory.percent
            },
            'swap': {
                'total': swap.total,
                'used': swap.used,
                'free': swap.free,
                'percent': swap.percent
            }
        }
        
        if self.nvml_initialized:
            try:
                device_count = py3nvml.nvmlDeviceGetCount()
                gpu_memory = []
                
                for i in range(device_count):
                    handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
                    info = py3nvml.nvmlDeviceGetMemoryInfo(handle)
                    gpu_memory.append({
                        'index': i,
                        'total': info.total,
                        'used': info.used,
                        'free': info.free
                    })
                
                memory_stats['gpu'] = gpu_memory
            except:
                memory_stats['gpu'] = {'error': 'GPU memory data unavailable'}
        
        return memory_stats

    async def get_power_stats(self) -> Dict[str, Any]:
        """Get power consumption statistics"""
        power_stats = {
            'cpu': await self._get_cpu_power(),
            'gpu': await self._get_gpu_power()
        }
        return power_stats

    async def _get_cpu_power(self) -> Dict[str, Any]:
        """Get CPU power consumption"""
        try:
            # This is a simplified estimation based on CPU usage
            cpu_percent = psutil.cpu_percent(interval=1, percpu=True)
            return {
                'percent_per_core': cpu_percent,
                'average_percent': sum(cpu_percent) / len(cpu_percent)
            }
        except:
            return {'error': 'CPU power data unavailable'}

    async def _get_gpu_power(self) -> Dict[str, Any]:
        """Get GPU power consumption"""
        if not self.nvml_initialized:
            return {'error': 'NVIDIA tools not available'}

        try:
            device_count = py3nvml.nvmlDeviceGetCount()
            gpu_power = []
            
            for i in range(device_count):
                handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
                power = py3nvml.nvmlDeviceGetPowerUsage(handle) / 1000.0  # Convert to Watts
                power_limit = py3nvml.nvmlDeviceGetEnforcedPowerLimit(handle) / 1000.0
                gpu_power.append({
                    'index': i,
                    'power_usage': power,
                    'power_limit': power_limit,
                    'percent': (power / power_limit) * 100 if power_limit > 0 else 0
                })
            
            return {'devices': gpu_power}
        except:
            return {'error': 'GPU power data unavailable'}

    async def get_cooling_status(self) -> Dict[str, Any]:
        """Get cooling system status"""
        try:
            fans = psutil.sensors_fans()
            return {
                'fans': fans,
                'cpu_fans': await self._get_cpu_fan_speeds(),
                'gpu_fans': await self._get_gpu_fan_speeds()
            }
        except:
            return {'error': 'Cooling status unavailable'}

    async def _get_cpu_fan_speeds(self) -> List[Dict[str, Any]]:
        """Get CPU fan speeds"""
        try:
            fans = psutil.sensors_fans()
            return [{'label': fan.label, 'speed': fan.current} for fan in fans.get('cpu_fan', [])]
        except:
            return []

    async def _get_gpu_fan_speeds(self) -> List[Dict[str, Any]]:
        """Get GPU fan speeds"""
        if not self.nvml_initialized:
            return []

        try:
            device_count = py3nvml.nvmlDeviceGetCount()
            gpu_fans = []
            
            for i in range(device_count):
                handle = py3nvml.nvmlDeviceGetHandleByIndex(i)
                speed = py3nvml.nvmlDeviceGetFanSpeed(handle)
                gpu_fans.append({
                    'index': i,
                    'speed': speed
                })
            
            return gpu_fans
        except:
            return []

    def cleanup(self):
        """Cleanup monitoring resources"""
        if self.nvml_initialized:
            try:
                py3nvml.nvmlShutdown()
            except:
                pass

class PerformancePredictor:
    def __init__(self, history_data: Dict[str, List[Dict[str, Any]]]):
        self.history_data = history_data

    async def predict_hash_rate(self, hours_ahead: int = 24) -> Dict[str, Any]:
        """Predict hash rate for the next period"""
        try:
            hash_rates = [d['data']['hash_rate'] for d in self.history_data.get('performance', [])]
            if not hash_rates:
                return {'error': 'Insufficient data for prediction'}

            # Simple linear regression for prediction
            x = np.arange(len(hash_rates))
            y = np.array(hash_rates)
            z = np.polyfit(x, y, 1)
            p = np.poly1d(z)
            
            # Predict future values
            future_x = np.arange(len(hash_rates), len(hash_rates) + hours_ahead)
            predictions = p(future_x)
            
            return {
                'current_trend': z[0],  # Slope of the trend
                'predictions': predictions.tolist(),
                'confidence': self._calculate_prediction_confidence(y, p(x))
            }
        except:
            return {'error': 'Prediction failed'}

    def _calculate_prediction_confidence(self, actual: np.ndarray, predicted: np.ndarray) -> float:
        """Calculate prediction confidence based on R-squared value"""
        try:
            residuals = actual - predicted
            ss_res = np.sum(residuals ** 2)
            ss_tot = np.sum((actual - np.mean(actual)) ** 2)
            r_squared = 1 - (ss_res / ss_tot)
            return float(r_squared)
        except:
            return 0.0
