#!/usr/bin/env python3
"""Intel NPU monitoring implementation (experimental)."""

import re
import subprocess
import time
from pathlib import Path
from typing import Dict, List, Optional, Tuple

import psutil

from .base import BaseGPUMonitor, GPUInfo


class IntelNPUMonitor(BaseGPUMonitor):
    """Intel NPU monitor backed by accel sysfs telemetry."""

    INTEL_VENDOR_ID = "0x8086"
    _TOTAL_MEMORY_KEYS = (
        "npu_memory_total",
        "npu_memory_total_bytes",
        "npu_total_memory",
        "npu_total_memory_bytes",
        "mem_total",
    )

    def __init__(self) -> None:
        """Initialize Intel NPU monitor."""
        super().__init__()
        self._busy_counters: Dict[int, Tuple[float, int]] = {}
        self._cached_driver_version = ""
        self._driver_version_initialized = False

    def _find_smi(self) -> Optional[Path]:
        """Detect Intel NPU availability through accel sysfs."""
        devices = self._iter_intel_accel_devices()
        if devices:
            return Path("/sys/class/accel")
        return None

    def _iter_intel_accel_devices(self) -> List[Path]:
        """List Intel accel devices exposed by the kernel."""
        accel_root = Path("/sys/class/accel")
        if not accel_root.exists():
            return []

        devices: List[Path] = []
        for accel in sorted(accel_root.glob("accel*")):
            vendor = self._read_text(accel / "device" / "vendor")
            if vendor.lower() == self.INTEL_VENDOR_ID:
                devices.append(accel)
        return devices

    @staticmethod
    def _read_text(path: Path) -> str:
        """Read a text file and return a stripped value."""
        try:
            return path.read_text(encoding="utf-8", errors="ignore").strip()
        except (OSError, UnicodeError):
            return ""

    @staticmethod
    def _read_int(path: Path) -> Optional[int]:
        """Read integer content from sysfs file."""
        raw = IntelNPUMonitor._read_text(path)
        if not raw:
            return None
        try:
            return int(raw, 0)
        except ValueError:
            return None

    @staticmethod
    def _accel_index(path: Path, fallback: int) -> int:
        """Parse accel index from accel path name."""
        match = re.search(r"accel(\d+)", path.name)
        if match:
            return int(match.group(1))
        return fallback

    def _get_driver_version(self) -> str:
        """Get Intel NPU driver version from sysfs/modinfo (cached)."""
        if self._driver_version_initialized:
            return self._cached_driver_version

        version = ""
        for sysfs_path in (
            Path("/sys/module/ivpu/version"),
            Path("/sys/module/intel_vpu/version"),
        ):
            value = self._read_text(sysfs_path)
            if value:
                version = value
                break

        if not version:
            for module in ("ivpu", "intel_vpu"):
                try:
                    result = subprocess.run(
                        ["modinfo", module],
                        capture_output=True,
                        text=True,
                        timeout=2.0,
                        check=False,
                    )
                    if result.returncode != 0:
                        continue
                    for line in result.stdout.splitlines():
                        if line.startswith("version:"):
                            candidate = line.split(":", 1)[1].strip()
                            if candidate:
                                version = candidate
                                break
                except Exception:
                    continue
                if version:
                    break

        self._cached_driver_version = version
        self._driver_version_initialized = True
        return version

    def _compute_utilization(self, index: int, busy_time_us: int) -> float:
        """Compute utilization percentage from npu_busy_time_us deltas."""
        now = time.monotonic()
        utilization = 0.0

        previous = self._busy_counters.get(index)
        if previous:
            prev_time, prev_busy = previous
            elapsed = now - prev_time
            busy_delta = busy_time_us - prev_busy
            if elapsed > 0 and busy_delta >= 0:
                utilization = (busy_delta / (elapsed * 1_000_000.0)) * 100.0

        self._busy_counters[index] = (now, busy_time_us)
        return max(0.0, min(100.0, utilization))

    def _build_device_name(self, device_path: Path, index: int) -> str:
        """Build a readable device label from sysfs fields."""
        device_id = self._read_text(device_path / "device")
        if not device_id:
            device_id = self._read_text(device_path)
        if device_id:
            return f"Intel AI Boost ({device_id})"
        return f"Intel AI Boost {index}"

    def _resolve_memory_budget_mb(self, device_path: Path) -> Tuple[float, str, str]:
        """Resolve shared-memory budget in MB with confidence metadata."""
        for key in self._TOTAL_MEMORY_KEYS:
            total_bytes = self._read_int(device_path / key)
            if total_bytes and total_bytes > 0:
                return (
                    float(total_bytes) / (1024.0 * 1024.0),
                    f"sysfs:{key}",
                    "direct",
                )

        try:
            host_total_bytes = int(psutil.virtual_memory().total)
        except Exception:
            return 0.0, "sysfs:npu_memory_utilization", "unknown"

        if host_total_bytes <= 0:
            return 0.0, "sysfs:npu_memory_utilization", "unknown"

        return (
            float(host_total_bytes) / (1024.0 * 1024.0),
            "sysfs:npu_memory_utilization+host_ram",
            "estimated",
        )

    def get_quick_metrics(self) -> Dict[int, Dict[str, float]]:
        """Get lightweight utilization and memory metrics."""
        metrics: Dict[int, Dict[str, float]] = {}
        for fallback, accel in enumerate(self._iter_intel_accel_devices()):
            index = self._accel_index(accel, fallback=fallback)
            device = accel / "device"
            busy_time_us = self._read_int(device / "npu_busy_time_us")
            memory_bytes = self._read_int(device / "npu_memory_utilization")

            utilization = (
                self._compute_utilization(index, busy_time_us)
                if busy_time_us is not None
                else 0.0
            )
            memory_used_mb = (
                float(memory_bytes) / (1024.0 * 1024.0) if memory_bytes else 0.0
            )
            memory_total_mb, _memory_source, _memory_confidence = (
                self._resolve_memory_budget_mb(device)
            )

            metrics[index] = {
                "utilization": utilization,
                "memory_used": memory_used_mb,
                "memory_total": memory_total_mb,
                "memory_percent": (
                    memory_used_mb / memory_total_mb * 100.0
                    if memory_total_mb > 0
                    else 0.0
                ),
            }

        return metrics

    def get_gpu_info(self) -> List[GPUInfo]:
        """Get Intel NPU data mapped into GPUInfo for UI compatibility."""
        if not self.smi_path:
            return []

        driver_version = self._get_driver_version()
        gpus: List[GPUInfo] = []

        for fallback, accel in enumerate(self._iter_intel_accel_devices()):
            index = self._accel_index(accel, fallback=fallback)
            device = accel / "device"

            busy_time_us = self._read_int(device / "npu_busy_time_us")
            memory_bytes = self._read_int(device / "npu_memory_utilization")
            current_freq = self._read_int(device / "npu_current_frequency_mhz")
            max_freq = self._read_int(device / "npu_max_frequency_mhz")

            utilization = (
                self._compute_utilization(index, busy_time_us)
                if busy_time_us is not None
                else 0.0
            )
            memory_used_mb = (
                float(memory_bytes) / (1024.0 * 1024.0) if memory_bytes else 0.0
            )
            memory_total_mb, memory_source, memory_confidence = (
                self._resolve_memory_budget_mb(device)
            )

            name = self._build_device_name(device, index)
            if current_freq and max_freq:
                name = f"{name} {current_freq}/{max_freq}MHz"

            gpus.append(
                GPUInfo(
                    index=index,
                    name=name,
                    utilization=utilization,
                    memory_used=memory_used_mb,
                    memory_total=memory_total_mb,
                    temperature=0.0,
                    power_draw=0.0,
                    power_limit=0.0,
                    processes=[],
                    driver_version=driver_version,
                    memory_model="shared",
                    memory_source=memory_source,
                    memory_confidence=memory_confidence,
                )
            )

        return gpus
