"""
sichgate/system_health.py — System preflight checks.

Runs before model loading to catch resource problems early rather than
mid-run (after potentially waiting minutes for a model to load).

Three checks:
  1. Memory     — warn when available RAM < 4 GB; note swap risk < 8 GB
  2. Thermal    — warn when CPU is already throttled (macOS sysctl) or
                  heavily loaded (>80% CPU usage as a cross-platform proxy)
  3. Architecture — confirm MLX/MPS backend is active on Apple Silicon;
                    warn on Intel Mac where PyTorch CPU will be used

Returns a HealthReport dataclass so callers can gate on it programmatically.
"""

from __future__ import annotations

import platform
import subprocess
import dataclasses
from enum import Enum
from typing import Optional


# ── Severity levels ───────────────────────────────────────────────────────────

class CheckLevel(str, Enum):
    OK      = "ok"
    CAUTION = "caution"   # yellow — proceed, but be aware
    WARN    = "warn"      # orange — recommend action before proceeding
    BLOCK   = "block"     # red    — serious enough to recommend stopping


@dataclasses.dataclass
class CheckResult:
    name:    str
    level:   CheckLevel
    message: str
    detail:  Optional[str] = None   # extra context shown in verbose mode


@dataclasses.dataclass
class HealthReport:
    checks:     list[CheckResult]
    can_proceed: bool    # False only when a BLOCK-level finding is present
    mlx_ready:  bool     # True if mlx-lm is installed and on Apple Silicon
    mps_ready:  bool     # True if PyTorch MPS backend is available

    @property
    def worst_level(self) -> CheckLevel:
        order = [CheckLevel.OK, CheckLevel.CAUTION, CheckLevel.WARN, CheckLevel.BLOCK]
        return max(self.checks, key=lambda c: order.index(c.level)).level


# ── Individual checks ─────────────────────────────────────────────────────────

def _check_memory() -> CheckResult:
    try:
        import psutil
    except ImportError:
        return CheckResult(
            name="Memory",
            level=CheckLevel.CAUTION,
            message="psutil not installed — cannot check RAM",
            detail="pip install psutil",
        )

    vm = psutil.virtual_memory()
    available_gb = vm.available / (1024 ** 3)
    total_gb     = vm.total    / (1024 ** 3)
    used_pct     = vm.percent

    if available_gb < 2:
        return CheckResult(
            name="Memory",
            level=CheckLevel.BLOCK,
            message=f"{available_gb:.1f} GB available — critically low",
            detail=(
                f"Total: {total_gb:.0f} GB  |  Used: {used_pct:.0f}%\n"
                "Close Chrome, Slack, or other heavy apps before loading a model. "
                "Running with <2 GB free will cause extreme swap and may stall the run."
            ),
        )
    elif available_gb < 4:
        return CheckResult(
            name="Memory",
            level=CheckLevel.WARN,
            message=f"{available_gb:.1f} GB available — swap likely",
            detail=(
                f"Total: {total_gb:.0f} GB  |  Used: {used_pct:.0f}%\n"
                "A 4-bit 1.5B model needs ~1.5 GB; a 7B model needs ~4.5 GB. "
                "Close background apps to avoid heavy swap usage."
            ),
        )
    elif available_gb < 8:
        return CheckResult(
            name="Memory",
            level=CheckLevel.CAUTION,
            message=f"{available_gb:.1f} GB available — adequate for small models",
            detail=(
                f"Total: {total_gb:.0f} GB  |  Used: {used_pct:.0f}%\n"
                "7B+ models may use swap. Stick to 1.5B–3B models for smooth runs."
            ),
        )
    else:
        return CheckResult(
            name="Memory",
            level=CheckLevel.OK,
            message=f"{available_gb:.1f} GB available",
            detail=f"Total: {total_gb:.0f} GB  |  Used: {used_pct:.0f}%",
        )


def _check_thermal() -> CheckResult:
    """
    macOS: reads kern.thermal_level via sysctl (no root required).
      0 = nominal, 1 = moderate throttling, 2 = heavy, 3 = trapping
    Fallback (Linux / Intel Mac): sample CPU utilisation over 0.5 s.
    """
    is_mac = platform.system() == "Darwin"

    if is_mac:
        try:
            result = subprocess.run(
                ["sysctl", "-n", "kern.thermal_level"],
                capture_output=True, text=True, timeout=2,
            )
            level = int(result.stdout.strip())
            labels = {0: "nominal", 1: "moderate throttle", 2: "heavy throttle", 3: "critical"}
            label  = labels.get(level, f"unknown ({level})")

            if level == 0:
                return CheckResult(
                    name="Thermal",
                    level=CheckLevel.OK,
                    message=f"Thermal state: {label}",
                )
            elif level == 1:
                return CheckResult(
                    name="Thermal",
                    level=CheckLevel.CAUTION,
                    message=f"Thermal state: {label}",
                    detail="CPU is mildly throttled. The run will still complete but may be slower.",
                )
            elif level == 2:
                return CheckResult(
                    name="Thermal",
                    level=CheckLevel.WARN,
                    message=f"Thermal state: {label}",
                    detail=(
                        "CPU is heavily throttled — inference will be significantly slower. "
                        "Let the Mac cool for 5 minutes before starting the run."
                    ),
                )
            else:
                return CheckResult(
                    name="Thermal",
                    level=CheckLevel.BLOCK,
                    message=f"Thermal state: {label} — Mac is at thermal limit",
                    detail=(
                        "The system is at critical thermal pressure. "
                        "Quit all heavy processes, let the Mac cool for 10+ minutes, "
                        "and ensure adequate airflow before running."
                    ),
                )
        except (ValueError, FileNotFoundError, subprocess.TimeoutExpired):
            pass  # fall through to CPU-usage proxy

    # Fallback: CPU utilisation as a thermal proxy
    try:
        import psutil
        cpu_pct = psutil.cpu_percent(interval=0.5)
        if cpu_pct >= 90:
            return CheckResult(
                name="Thermal",
                level=CheckLevel.WARN,
                message=f"CPU already at {cpu_pct:.0f}% — likely thermally constrained",
                detail="Close background processes. High baseline CPU = slow inference + swap.",
            )
        elif cpu_pct >= 60:
            return CheckResult(
                name="Thermal",
                level=CheckLevel.CAUTION,
                message=f"CPU at {cpu_pct:.0f}% baseline",
                detail="Some background load detected. Inference will compete for CPU time.",
            )
        else:
            return CheckResult(
                name="Thermal",
                level=CheckLevel.OK,
                message=f"CPU at {cpu_pct:.0f}% baseline — ready",
            )
    except ImportError:
        return CheckResult(
            name="Thermal",
            level=CheckLevel.CAUTION,
            message="psutil not installed — cannot check CPU/thermal",
            detail="pip install psutil",
        )


def _check_architecture() -> tuple[CheckResult, bool, bool]:
    """
    Returns (CheckResult, mlx_ready, mps_ready).
    """
    is_mac        = platform.system() == "Darwin"
    is_arm        = platform.machine() == "arm64"
    is_apple_si   = is_mac and is_arm

    # Check MLX
    mlx_ready = False
    try:
        import mlx_lm  # noqa: F401
        mlx_ready = True
    except ImportError:
        pass

    # Check MPS (PyTorch fallback)
    mps_ready = False
    try:
        import torch
        mps_ready = torch.backends.mps.is_available()
    except ImportError:
        pass

    if is_apple_si and mlx_ready:
        return (
            CheckResult(
                name="Architecture",
                level=CheckLevel.OK,
                message="Apple Silicon + mlx-lm — GPU/Neural Engine active",
                detail=(
                    f"Chip: {platform.processor() or 'Apple Silicon (arm64)'}\n"
                    "Inference runs on Metal GPU — CPU cores stay free."
                ),
            ),
            mlx_ready, mps_ready,
        )
    elif is_apple_si and mps_ready:
        return (
            CheckResult(
                name="Architecture",
                level=CheckLevel.CAUTION,
                message="Apple Silicon — PyTorch MPS backend (no mlx-lm)",
                detail=(
                    "MPS is better than CPU but mlx-lm is 3–10× faster and uses less RAM.\n"
                    "Install: pip install mlx-lm"
                ),
            ),
            mlx_ready, mps_ready,
        )
    elif is_apple_si:
        return (
            CheckResult(
                name="Architecture",
                level=CheckLevel.WARN,
                message="Apple Silicon — running on CPU only",
                detail=(
                    "Neither mlx-lm nor PyTorch MPS is available.\n"
                    "Install mlx-lm for GPU inference: pip install mlx-lm"
                ),
            ),
            mlx_ready, mps_ready,
        )
    elif is_mac:
        return (
            CheckResult(
                name="Architecture",
                level=CheckLevel.CAUTION,
                message="Intel Mac — CPU inference only",
                detail=(
                    "mlx-lm requires Apple Silicon (M1/M2/M3/M4). "
                    "Inference will run on CPU. Stick to 1.5B models to avoid swap."
                ),
            ),
            mlx_ready, mps_ready,
        )
    else:
        # Linux / Windows
        try:
            import torch
            if torch.cuda.is_available():
                gpu = torch.cuda.get_device_name(0)
                return (
                    CheckResult(
                        name="Architecture",
                        level=CheckLevel.OK,
                        message=f"CUDA GPU detected: {gpu}",
                    ),
                    mlx_ready, mps_ready,
                )
        except ImportError:
            pass
        return (
            CheckResult(
                name="Architecture",
                level=CheckLevel.CAUTION,
                message=f"{platform.system()} {platform.machine()} — CPU inference",
                detail="No GPU backend detected. Inference will be slow on large models.",
            ),
            mlx_ready, mps_ready,
        )


# ── Public API ────────────────────────────────────────────────────────────────

def run_checks() -> HealthReport:
    """Run all three preflight checks and return a HealthReport."""
    mem_check             = _check_memory()
    thermal_check         = _check_thermal()
    arch_check, mlx, mps  = _check_architecture()

    checks = [mem_check, thermal_check, arch_check]
    can_proceed = all(c.level != CheckLevel.BLOCK for c in checks)

    return HealthReport(
        checks=checks,
        can_proceed=can_proceed,
        mlx_ready=mlx,
        mps_ready=mps,
    )


def print_report(report: HealthReport, verbose: bool = False) -> None:
    """Print a Rich-formatted health report to stdout."""
    from rich.console import Console
    from rich.table import Table
    from rich.text import Text

    console = Console()

    _ICONS = {
        CheckLevel.OK:      ("[green]✓[/green]",  "[green]OK[/green]"),
        CheckLevel.CAUTION: ("[yellow]⚠[/yellow]", "[yellow]CAUTION[/yellow]"),
        CheckLevel.WARN:    ("[orange1]⚠[/orange1]", "[orange1]WARN[/orange1]"),
        CheckLevel.BLOCK:   ("[red]✗[/red]",      "[red]BLOCK[/red]"),
    }

    table = Table(
        title="[bold]SichGate — System Health Check[/bold]",
        show_header=True,
        header_style="bold dim",
        border_style="dim",
        expand=False,
    )
    table.add_column("Check",   style="bold", min_width=14)
    table.add_column("Status",  min_width=9)
    table.add_column("Detail")

    for check in report.checks:
        icon, badge = _ICONS[check.level]
        detail_text = check.message
        if verbose and check.detail:
            detail_text += f"\n[dim]{check.detail}[/dim]"
        table.add_row(check.name, badge, detail_text)

    console.print()
    console.print(table)

    if report.can_proceed:
        if report.worst_level in (CheckLevel.OK, CheckLevel.CAUTION):
            console.print("  [green]✓[/green] System ready — starting red team run\n")
        else:
            console.print("  [orange1]⚠[/orange1] Warnings above — proceeding, but expect slower inference\n")
    else:
        console.print(
            "  [red bold]✗ BLOCK-level issue detected.[/red bold]\n"
            "    Resolve the issue above before running. "
            "Use [cyan]sichgate check --verbose[/cyan] for details.\n"
        )
