"""
System health monitoring: CPU temperature and memory pressure.

Reads CPU temperature from /sys/class/thermal/ and provides:
    - Thermal state classification: ok (<75C), warm (75-85C), hot (>85C)
    - Dynamic worker count reduction when CPU is warm
    - Cooldown pauses (2s) between heavy computation steps
    - wait_until_cool() blocks until temp drops below 70C
    - check_memory_pressure() checks both memory AND temperature

Integrated into the engine pipeline so every evolution step checks
system health before proceeding with heavy computation.
"""

import os
import time
import logging
import psutil
from enum import Enum
from dataclasses import dataclass
from typing import Optional, List

logger = logging.getLogger(__name__)


class ThermalState(Enum):
    """CPU thermal state classification."""
    OK = "ok"           # < 75C -- full speed
    WARM = "warm"       # 75-85C -- throttle (halve workers)
    HOT = "hot"         # > 85C -- pause until cool


# Thresholds (Celsius)
TEMP_OK_MAX = 75.0
TEMP_WARM_MAX = 85.0
TEMP_COOL_TARGET = 70.0

# Cooldown pause duration (seconds)
COOLDOWN_PAUSE_SECONDS = 2.0

# Memory pressure threshold (fraction of total RAM)
MEMORY_PRESSURE_THRESHOLD = 0.85


@dataclass
class HealthStatus:
    """Snapshot of system health at a point in time."""
    cpu_temp_celsius: Optional[float]
    thermal_state: ThermalState
    memory_used_fraction: float
    memory_available_mb: float
    memory_pressure: bool
    recommended_workers: int


class SystemHealthMonitor:
    """
    Monitors CPU temperature and memory to protect hardware during
    heavy quantum simulation workloads.

    Reads temperature from /sys/class/thermal/thermal_zone*/temp.
    Falls back gracefully if thermal sensors are unavailable.

    Usage:
        monitor = SystemHealthMonitor(max_workers=8)

        # Before heavy computation
        monitor.check_memory_pressure()  # checks memory AND temperature

        # Get current recommended parallelism
        n_workers = monitor.get_recommended_workers()

        # Block until safe to proceed with heavy work
        monitor.wait_until_cool()
    """

    def __init__(
        self,
        max_workers: Optional[int] = None,
        temp_ok_max: float = TEMP_OK_MAX,
        temp_warm_max: float = TEMP_WARM_MAX,
        cool_target: float = TEMP_COOL_TARGET,
        cooldown_pause: float = COOLDOWN_PAUSE_SECONDS,
        memory_threshold: float = MEMORY_PRESSURE_THRESHOLD,
    ):
        """
        Args:
            max_workers: Maximum parallel workers at full speed.
                         Default: os.cpu_count() or 4.
            temp_ok_max: Temperature below which state is OK (Celsius).
            temp_warm_max: Temperature below which state is WARM (Celsius).
            cool_target: Temperature target for wait_until_cool (Celsius).
            cooldown_pause: Seconds to pause between heavy steps.
            memory_threshold: Fraction of RAM usage that triggers pressure.
        """
        self.max_workers = max_workers or (os.cpu_count() or 4)
        self.temp_ok_max = temp_ok_max
        self.temp_warm_max = temp_warm_max
        self.cool_target = cool_target
        self.cooldown_pause = cooldown_pause
        self.memory_threshold = memory_threshold

        # Discover thermal zone paths once at init
        self._thermal_paths = self._discover_thermal_zones()
        if self._thermal_paths:
            logger.info(
                "SystemHealthMonitor: found %d thermal zone(s)",
                len(self._thermal_paths),
            )
        else:
            logger.info(
                "SystemHealthMonitor: no thermal zones found, "
                "temperature monitoring disabled"
            )

    # ------------------------------------------------------------------
    # Thermal zone discovery
    # ------------------------------------------------------------------

    @staticmethod
    def _discover_thermal_zones() -> List[str]:
        """Find all /sys/class/thermal/thermal_zone*/temp files."""
        base = "/sys/class/thermal"
        paths = []
        if not os.path.isdir(base):
            return paths
        try:
            for entry in sorted(os.listdir(base)):
                if entry.startswith("thermal_zone"):
                    temp_file = os.path.join(base, entry, "temp")
                    if os.path.isfile(temp_file):
                        paths.append(temp_file)
        except OSError:
            pass
        return paths

    # ------------------------------------------------------------------
    # Temperature reading
    # ------------------------------------------------------------------

    def read_cpu_temperature(self) -> Optional[float]:
        """
        Read the highest CPU temperature across all thermal zones.

        The kernel reports temperature in millidegrees Celsius (e.g. 72000
        means 72.0C).  We take the max across all zones to be conservative.

        Returns:
            Temperature in Celsius, or None if sensors are unavailable.
        """
        if not self._thermal_paths:
            return None

        temps = []
        for path in self._thermal_paths:
            try:
                with open(path, "r") as f:
                    raw = f.read().strip()
                temp_c = float(raw) / 1000.0
                temps.append(temp_c)
            except (OSError, ValueError):
                continue

        return max(temps) if temps else None

    # ------------------------------------------------------------------
    # Thermal state classification
    # ------------------------------------------------------------------

    def classify_thermal_state(
        self, temp_celsius: Optional[float] = None
    ) -> ThermalState:
        """
        Classify current thermal state.

        Args:
            temp_celsius: Override temperature. If None, reads from sensors.

        Returns:
            ThermalState enum value.
        """
        if temp_celsius is None:
            temp_celsius = self.read_cpu_temperature()

        if temp_celsius is None:
            # Sensors unavailable -- assume OK
            return ThermalState.OK

        if temp_celsius < self.temp_ok_max:
            return ThermalState.OK
        elif temp_celsius < self.temp_warm_max:
            return ThermalState.WARM
        else:
            return ThermalState.HOT

    # ------------------------------------------------------------------
    # Memory pressure
    # ------------------------------------------------------------------

    @staticmethod
    def _get_memory_info() -> dict:
        """Return memory usage info via psutil."""
        mem = psutil.virtual_memory()
        return {
            "used_fraction": mem.percent / 100.0,
            "available_mb": mem.available / (1024 * 1024),
            "total_mb": mem.total / (1024 * 1024),
        }

    def is_memory_pressure(self) -> bool:
        """Check if system is under memory pressure."""
        info = self._get_memory_info()
        return info["used_fraction"] >= self.memory_threshold

    # ------------------------------------------------------------------
    # Combined health check
    # ------------------------------------------------------------------

    def check_health(self) -> HealthStatus:
        """
        Full health snapshot: temperature + memory.

        Returns:
            HealthStatus with all metrics and recommendations.
        """
        temp = self.read_cpu_temperature()
        state = self.classify_thermal_state(temp)
        mem = self._get_memory_info()
        pressure = mem["used_fraction"] >= self.memory_threshold

        workers = self._compute_workers(state, pressure)

        return HealthStatus(
            cpu_temp_celsius=temp,
            thermal_state=state,
            memory_used_fraction=mem["used_fraction"],
            memory_available_mb=mem["available_mb"],
            memory_pressure=pressure,
            recommended_workers=workers,
        )

    def _compute_workers(
        self, state: ThermalState, memory_pressure: bool
    ) -> int:
        """Compute recommended worker count given current conditions."""
        workers = self.max_workers

        # Halve on warm CPU
        if state == ThermalState.WARM:
            workers = max(1, workers // 2)
            logger.debug(
                "CPU warm -- reducing workers to %d", workers
            )

        # Single worker on hot CPU
        if state == ThermalState.HOT:
            workers = 1
            logger.warning(
                "CPU hot (>%.0fC) -- reducing to 1 worker",
                self.temp_warm_max,
            )

        # Further halve under memory pressure
        if memory_pressure:
            workers = max(1, workers // 2)
            logger.debug(
                "Memory pressure (>%.0f%%) -- reducing workers to %d",
                self.memory_threshold * 100,
                workers,
            )

        return workers

    # ------------------------------------------------------------------
    # Dynamic worker recommendation
    # ------------------------------------------------------------------

    def get_recommended_workers(self) -> int:
        """
        Return the recommended number of parallel workers right now.

        - OK state + no memory pressure: max_workers
        - WARM state: max_workers // 2
        - HOT state: 1
        - Memory pressure: additional halving
        """
        status = self.check_health()
        return status.recommended_workers

    # ------------------------------------------------------------------
    # check_memory_pressure -- combined memory + thermal gate
    # ------------------------------------------------------------------

    def check_memory_pressure(self) -> HealthStatus:
        """
        Check both memory AND temperature.

        If the CPU is HOT, pauses for cooldown_pause seconds.
        If memory is under pressure, logs a warning.

        This should be called before every heavy computation step.

        Returns:
            Current HealthStatus.
        """
        status = self.check_health()

        if status.thermal_state == ThermalState.HOT:
            logger.warning(
                "CPU temperature %.1fC > %.0fC -- "
                "pausing %.1fs for cooldown",
                status.cpu_temp_celsius or 0.0,
                self.temp_warm_max,
                self.cooldown_pause,
            )
            time.sleep(self.cooldown_pause)

        if status.memory_pressure:
            logger.warning(
                "Memory pressure: %.1f%% used (%.0f MB available)",
                status.memory_used_fraction * 100,
                status.memory_available_mb,
            )

        return status

    # ------------------------------------------------------------------
    # wait_until_cool -- blocking cooldown
    # ------------------------------------------------------------------

    def wait_until_cool(
        self,
        poll_interval: float = 2.0,
        max_wait: float = 300.0,
    ) -> float:
        """
        Block until CPU temperature drops below cool_target (default 70C).

        Used before launching heavy parallel workloads. If sensors are
        unavailable, returns immediately.

        Args:
            poll_interval: Seconds between temperature checks.
            max_wait: Maximum seconds to wait before giving up.

        Returns:
            Total seconds waited.
        """
        if not self._thermal_paths:
            return 0.0

        total_waited = 0.0
        temp = self.read_cpu_temperature()

        while temp is not None and temp > self.cool_target:
            if total_waited >= max_wait:
                logger.warning(
                    "wait_until_cool: timed out after %.0fs "
                    "(temp still %.1fC > %.0fC)",
                    total_waited,
                    temp,
                    self.cool_target,
                )
                break

            logger.info(
                "CPU at %.1fC > %.0fC target -- "
                "waiting %.1fs (%.0fs elapsed)",
                temp,
                self.cool_target,
                poll_interval,
                total_waited,
            )
            time.sleep(poll_interval)
            total_waited += poll_interval
            temp = self.read_cpu_temperature()

        if total_waited > 0:
            logger.info(
                "CPU cooled to %.1fC after %.1fs",
                temp if temp is not None else 0.0,
                total_waited,
            )

        return total_waited

    # ------------------------------------------------------------------
    # Cooldown pause for between computation steps
    # ------------------------------------------------------------------

    def cooldown_between_steps(self) -> None:
        """
        Insert a cooldown pause between heavy computation steps.

        Only pauses if CPU is WARM or HOT.  Duration is cooldown_pause
        seconds (default 2s).
        """
        state = self.classify_thermal_state()
        if state in (ThermalState.WARM, ThermalState.HOT):
            logger.debug(
                "Cooldown pause (%.1fs) -- CPU state: %s",
                self.cooldown_pause,
                state.value,
            )
            time.sleep(self.cooldown_pause)


# ------------------------------------------------------------------
# Module-level convenience: shared monitor instance
# ------------------------------------------------------------------

_default_monitor: Optional[SystemHealthMonitor] = None


def get_default_monitor() -> SystemHealthMonitor:
    """Get or create the module-level default SystemHealthMonitor."""
    global _default_monitor
    if _default_monitor is None:
        _default_monitor = SystemHealthMonitor()
    return _default_monitor