"""
Armory System Metrics - GPU/CPU/Memory monitoring

This module provides system resource monitoring similar to W&B.
Automatically collects available metrics with graceful degradation.
"""

import threading
import time
import warnings
from typing import Dict, Optional, Callable, Set
from dataclasses import dataclass, field

from armory.logger import get_logger

logger = get_logger("system")

# Track which warnings have been shown (to avoid spam)
_warnings_shown: Set[str] = set()


def _warn_once(key: str, message: str):
    """Show a warning only once per session"""
    if key not in _warnings_shown:
        _warnings_shown.add(key)
        logger.warning(message)


@dataclass
class SystemMetrics:
    """System resource metrics"""
    gpu_memory_used: float = 0.0      # bytes
    gpu_memory_total: float = 0.0     # bytes
    gpu_utilization: float = 0.0      # Percentage (0-100)
    gpu_temperature: float = 0.0      # Celsius
    cpu_utilization: float = 0.0      # Percentage (0-100)
    ram_used: float = 0.0             # bytes
    ram_total: float = 0.0            # bytes
    disk_used: float = 0.0            # bytes
    disk_total: float = 0.0           # bytes

    def to_dict(self) -> Dict[str, float]:
        """Convert to dict for API submission (values in bytes/percentage)"""
        return {
            "gpu_memory_used": self.gpu_memory_used,
            "gpu_memory_total": self.gpu_memory_total,
            "gpu_utilization": self.gpu_utilization,
            "gpu_temperature": self.gpu_temperature,
            "cpu_utilization": self.cpu_utilization,
            "ram_used": self.ram_used,
            "ram_total": self.ram_total,
        }


@dataclass
class SystemCapabilities:
    """Track available system monitoring capabilities"""
    has_psutil: bool = False
    has_pynvml: bool = False
    has_torch_cuda: bool = False
    gpu_count: int = 0
    gpu_names: list = field(default_factory=list)

    def summary(self) -> str:
        parts = []
        if self.has_psutil:
            parts.append("CPU/RAM")
        if self.has_pynvml or self.has_torch_cuda:
            gpu_info = f"GPU x{self.gpu_count}" if self.gpu_count > 0 else "GPU"
            parts.append(gpu_info)
        return ", ".join(parts) if parts else "none"


# Global capabilities (detected once)
_capabilities: Optional[SystemCapabilities] = None


def detect_capabilities() -> SystemCapabilities:
    """Detect available system monitoring capabilities"""
    global _capabilities
    if _capabilities is not None:
        return _capabilities

    caps = SystemCapabilities()

    # Check psutil
    try:
        import psutil
        caps.has_psutil = True
    except ImportError:
        _warn_once("psutil", "psutil not installed. CPU/RAM metrics will not be collected. Install with: pip install psutil")

    # Check pynvml (NVIDIA GPU)
    try:
        import pynvml
        pynvml.nvmlInit()
        caps.has_pynvml = True
        caps.gpu_count = pynvml.nvmlDeviceGetCount()
        for i in range(caps.gpu_count):
            handle = pynvml.nvmlDeviceGetHandleByIndex(i)
            name = pynvml.nvmlDeviceGetName(handle)
            if isinstance(name, bytes):
                name = name.decode('utf-8')
            caps.gpu_names.append(name)
        pynvml.nvmlShutdown()
    except ImportError:
        # Try torch as fallback
        try:
            import torch
            if torch.cuda.is_available():
                caps.has_torch_cuda = True
                caps.gpu_count = torch.cuda.device_count()
                for i in range(caps.gpu_count):
                    caps.gpu_names.append(torch.cuda.get_device_name(i))
            else:
                _warn_once("gpu", "No GPU detected. GPU metrics will not be collected.")
        except ImportError:
            _warn_once("gpu", "pynvml/torch not installed. GPU metrics will not be collected. Install with: pip install pynvml")
    except Exception as e:
        _warn_once("pynvml_init", f"Failed to initialize GPU monitoring: {e}")

    _capabilities = caps

    # Log what's available
    if caps.has_psutil or caps.has_pynvml or caps.has_torch_cuda:
        gpu_info = ""
        if caps.gpu_count > 0:
            gpu_info = f" ({', '.join(caps.gpu_names)})"
        logger.debug(f"System metrics enabled: {caps.summary()}{gpu_info}")

    return caps


def collect_system_metrics(gpu_index: int = 0) -> SystemMetrics:
    """
    Collect current system metrics.

    Args:
        gpu_index: Which GPU to monitor (default: 0)

    Returns:
        SystemMetrics with available metrics filled in
    """
    caps = detect_capabilities()
    metrics = SystemMetrics()

    # CPU/RAM metrics via psutil
    if caps.has_psutil:
        try:
            import psutil
            metrics.cpu_utilization = psutil.cpu_percent(interval=0.1)
            mem = psutil.virtual_memory()
            metrics.ram_used = mem.used
            metrics.ram_total = mem.total
        except Exception as e:
            _warn_once("psutil_collect", f"Failed to collect CPU/RAM metrics: {e}")

    # GPU metrics via pynvml (preferred)
    if caps.has_pynvml:
        try:
            import pynvml
            pynvml.nvmlInit()

            if gpu_index < caps.gpu_count:
                handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_index)

                # Memory
                mem_info = pynvml.nvmlDeviceGetMemoryInfo(handle)
                metrics.gpu_memory_used = mem_info.used
                metrics.gpu_memory_total = mem_info.total

                # Utilization
                util = pynvml.nvmlDeviceGetUtilizationRates(handle)
                metrics.gpu_utilization = util.gpu

                # Temperature
                try:
                    temp = pynvml.nvmlDeviceGetTemperature(handle, pynvml.NVML_TEMPERATURE_GPU)
                    metrics.gpu_temperature = temp
                except Exception:
                    pass  # Temperature might not be available on all GPUs

            pynvml.nvmlShutdown()
        except Exception as e:
            _warn_once("pynvml_collect", f"Failed to collect GPU metrics via pynvml: {e}")

    # GPU metrics via torch (fallback)
    elif caps.has_torch_cuda:
        try:
            import torch
            if gpu_index < caps.gpu_count:
                # Note: torch only provides memory info, not utilization/temperature
                metrics.gpu_memory_used = torch.cuda.memory_allocated(gpu_index)
                metrics.gpu_memory_total = torch.cuda.get_device_properties(gpu_index).total_memory
        except Exception as e:
            _warn_once("torch_collect", f"Failed to collect GPU metrics via torch: {e}")

    return metrics


class SystemMonitor:
    """
    Background thread for periodic system metrics collection.

    Automatically detects available metrics (CPU, RAM, GPU) and collects
    only what's available, with graceful degradation.

    Usage:
        monitor = SystemMonitor(callback=lambda m: run.log_system(m), interval=10)
        monitor.start()
        # ... training ...
        monitor.stop()
    """

    def __init__(
        self,
        callback: Callable[[SystemMetrics], None],
        interval: float = 10.0,  # seconds
        gpu_index: int = 0,
    ):
        self.callback = callback
        self.interval = interval
        self.gpu_index = gpu_index
        self._running = False
        self._thread: Optional[threading.Thread] = None
        self._step = 0

        # Detect capabilities on init
        self._caps = detect_capabilities()

    def _has_any_metrics(self) -> bool:
        """Check if any metrics can be collected"""
        return self._caps.has_psutil or self._caps.has_pynvml or self._caps.has_torch_cuda

    def _monitor_loop(self):
        while self._running:
            try:
                metrics = collect_system_metrics(self.gpu_index)
                # Only callback if we have meaningful data
                if self._has_any_metrics():
                    self.callback(metrics)
                self._step += 1
            except Exception as e:
                _warn_once("monitor_loop", f"System monitor error: {e}")
            time.sleep(self.interval)

    def start(self):
        """Start the monitoring thread"""
        if self._running:
            return

        # Skip if no metrics available
        if not self._has_any_metrics():
            _warn_once("no_metrics", "No system metrics available to collect. Skipping system monitoring.")
            return

        self._running = True
        self._thread = threading.Thread(target=self._monitor_loop, daemon=True, name="armory-system-monitor")
        self._thread.start()

    def stop(self):
        """Stop the monitoring thread"""
        self._running = False
        if self._thread:
            self._thread.join(timeout=self.interval + 1)
            self._thread = None
