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

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

import psutil

from .base import BaseGPUMonitor, GPUInfo


class AMDNPUMonitor(BaseGPUMonitor):
    """AMD NPU monitor using xrt-smi and accel sysfs fallbacks."""

    AMD_VENDOR_IDS = {"0x1022", "0x1002"}
    _TOTAL_MEMORY_KEYS = (
        "npu_memory_total",
        "npu_memory_total_bytes",
        "npu_total_memory",
        "npu_total_memory_bytes",
        "mem_total",
    )

    def __init__(self) -> None:
        """Initialize AMD 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]:
        """Find xrt-smi when available, fallback to accel sysfs presence."""
        candidate = shutil.which("xrt-smi")
        if candidate:
            path = Path(candidate)
            if path.exists() and os.access(path, os.X_OK):
                try:
                    result = subprocess.run(
                        [str(path), "examine"],
                        capture_output=True,
                        text=True,
                        timeout=3.0,
                        check=False,
                    )
                    output = f"{result.stdout}\n{result.stderr}".lower()
                    if result.returncode == 0 and "npu" in output:
                        return path
                except Exception:
                    pass

        if self._iter_amd_accel_devices():
            return Path("/sys/class/accel")

        return None

    def _iter_amd_accel_devices(self) -> List[Path]:
        """List AMD accel devices exposed via sysfs."""
        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").lower()
            if vendor in self.AMD_VENDOR_IDS:
                devices.append(accel)
                continue

            driver = self._read_symlink_name(accel / "device" / "driver")
            if "amdxdna" in driver:
                devices.append(accel)
        return devices

    @staticmethod
    def _read_text(path: Path) -> str:
        """Read a file as stripped UTF-8 text."""
        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 = AMDNPUMonitor._read_text(path)
        if not raw:
            return None
        try:
            return int(raw, 0)
        except ValueError:
            return None

    @staticmethod
    def _read_symlink_name(path: Path) -> str:
        """Read symlink target name."""
        try:
            return path.resolve().name.lower()
        except OSError:
            return ""

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

    @staticmethod
    def _normalize_bdf(value: str) -> str:
        """Normalize BDF text for mapping across tool outputs."""
        normalized = value.strip().strip("[]").lower()
        return normalized

    @staticmethod
    def _parse_memory_to_mb(raw_value: str) -> float:
        """Convert memory strings such as ``64 KB`` to MB."""
        text = str(raw_value).strip().lower()
        if not text:
            return 0.0

        match = re.search(
            r"([0-9][0-9,_.]*(?:\.[0-9]+)?)\s*"
            r"(tib|tb|gib|gb|mib|mb|kib|kb|bytes|byte|b)?",
            text,
        )
        if not match:
            return 0.0

        token = match.group(1).replace(",", "").replace("_", "")
        try:
            value = float(token)
        except ValueError:
            return 0.0

        unit = (match.group(2) or "mb").lower()
        multipliers = {
            "tib": 1024.0 * 1024.0,
            "tb": 1024.0 * 1024.0,
            "gib": 1024.0,
            "gb": 1024.0,
            "mib": 1.0,
            "mb": 1.0,
            "kib": 1.0 / 1024.0,
            "kb": 1.0 / 1024.0,
            "bytes": 1.0 / (1024.0 * 1024.0),
            "byte": 1.0 / (1024.0 * 1024.0),
            "b": 1.0 / (1024.0 * 1024.0),
        }
        return max(0.0, value * multipliers.get(unit, 1.0))

    def _resolve_memory_budget_mb(
        self, device_path: Optional[Path]
    ) -> Tuple[float, str, str]:
        """Resolve shared-memory budget in MB with confidence metadata."""
        if device_path is not None:
            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 _compute_utilization(self, index: int, busy_time_us: int) -> float:
        """Compute utilization percentage from busy-time counter 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 _run_xrt_smi(self, args: List[str], timeout: float = 4.0) -> str:
        """Run xrt-smi command and return stdout."""
        if not self.smi_path or self.smi_path.name != "xrt-smi":
            return ""

        try:
            result = subprocess.run(
                [str(self.smi_path), *args],
                capture_output=True,
                text=True,
                timeout=timeout,
                check=False,
            )
            if result.returncode == 0:
                return result.stdout
            if result.stdout.strip():
                return result.stdout
        except Exception:
            return ""
        return ""

    @staticmethod
    def _parse_xrt_device_names(output: str) -> Dict[str, str]:
        """Parse xrt-smi examine output into BDF -> device name map."""
        devices: Dict[str, str] = {}
        row_pattern = re.compile(
            r"\|\s*\[?([0-9a-fA-F:.]+)\]?\s*\|\s*([^|]+?)\s*\|",
        )
        for line in output.splitlines():
            match = row_pattern.search(line)
            if not match:
                continue
            bdf = AMDNPUMonitor._normalize_bdf(match.group(1))
            name = match.group(2).strip()
            if bdf and name:
                devices[bdf] = name
        return devices

    @staticmethod
    def _parse_xrt_driver_version(output: str) -> str:
        """Parse driver version hints from xrt-smi examine output."""
        for line in output.splitlines():
            stripped = line.strip()
            lowered = stripped.lower()
            if lowered.startswith("amdxdna") and ":" in stripped:
                return stripped.split(":", 1)[1].strip()
            if "npu driver version" in lowered and ":" in stripped:
                return stripped.split(":", 1)[1].strip()
        return ""

    @staticmethod
    def _parse_xrt_platform_power(output: str) -> Dict[str, float]:
        """Parse BDF -> estimated power (W) from xrt-smi platform report."""
        power_by_bdf: Dict[str, float] = {}
        current_bdf = ""

        for line in output.splitlines():
            bdf_match = re.search(r"\[([0-9a-fA-F:.]+)\]\s*:", line)
            if bdf_match:
                current_bdf = AMDNPUMonitor._normalize_bdf(bdf_match.group(1))
                continue

            if "Estimated Power" not in line or not current_bdf:
                continue

            value_match = re.search(r"(-?\d+(?:\.\d+)?)", line)
            if value_match:
                power_by_bdf[current_bdf] = float(value_match.group(1))

        return power_by_bdf

    @staticmethod
    def _resolve_process_name(pid: int) -> str:
        """Resolve process name for a PID."""
        try:
            proc = psutil.Process(pid)
            return proc.name() or "unknown"
        except Exception:
            return "unknown"

    @staticmethod
    def _parse_xrt_aie_processes(
        output: str,
    ) -> Tuple[Dict[str, List[Dict[str, Any]]], Dict[str, int]]:
        """Parse BDF->processes and BDF->active_context_count from AIE report."""
        processes_by_bdf: Dict[str, Dict[int, Dict[str, Any]]] = {}
        active_counts: Dict[str, int] = {}
        current_bdf = ""

        for line in output.splitlines():
            bdf_match = re.search(r"\[([0-9a-fA-F:.]+)\]\s*:", line)
            if bdf_match:
                current_bdf = AMDNPUMonitor._normalize_bdf(bdf_match.group(1))
                processes_by_bdf.setdefault(current_bdf, {})
                active_counts.setdefault(current_bdf, 0)
                continue

            stripped = line.strip()
            if not current_bdf or not stripped.startswith("|"):
                continue
            if "PID" in stripped or "---" in stripped:
                continue

            columns = [value.strip() for value in stripped.strip("|").split("|")]
            if len(columns) < 4:
                continue

            try:
                pid = int(columns[0])
            except ValueError:
                continue

            status = columns[2].lower()
            if status == "active":
                active_counts[current_bdf] = active_counts.get(current_bdf, 0) + 1

            memory = AMDNPUMonitor._parse_memory_to_mb(columns[3])
            existing = processes_by_bdf[current_bdf].get(pid)
            if existing is None or memory > existing.get("memory", 0.0):
                processes_by_bdf[current_bdf][pid] = {
                    "pid": pid,
                    "name": AMDNPUMonitor._resolve_process_name(pid),
                    "memory": memory,
                }

        finalized = {
            bdf: sorted(processes.values(), key=lambda proc: int(proc["pid"]))
            for bdf, processes in processes_by_bdf.items()
        }
        return finalized, active_counts

    def _get_driver_version(self) -> str:
        """Get AMD NPU driver version from xrt-smi or sysfs (cached)."""
        if self._driver_version_initialized:
            return self._cached_driver_version

        version = ""
        if self.smi_path and self.smi_path.name == "xrt-smi":
            examine = self._run_xrt_smi(["examine"])
            version = self._parse_xrt_driver_version(examine)

        if not version:
            version = self._read_text(Path("/sys/module/amdxdna/version"))

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

    def _device_bdf(self, accel: Path) -> str:
        """Get BDF string for an accel device."""
        try:
            return self._normalize_bdf((accel / "device").resolve().name)
        except OSError:
            return ""

    def _quick_metrics_from_devices(
        self, active_contexts: Dict[str, int]
    ) -> Dict[int, Dict[str, float]]:
        """Collect quick metrics for currently visible accel devices."""
        metrics: Dict[int, Dict[str, float]] = {}

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

            busy_time_us = self._read_int(device / "npu_busy_time_us")
            utilization = 0.0
            if busy_time_us is not None:
                utilization = self._compute_utilization(index, busy_time_us)
            elif active_contexts.get(bdf, 0) > 0:
                utilization = 100.0

            memory_bytes = self._read_int(device / "npu_memory_utilization")
            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_quick_metrics(self) -> Dict[int, Dict[str, float]]:
        """Get lightweight utilization/memory metrics."""
        aie_output = self._run_xrt_smi(["examine", "--report", "aie-partitions"])
        _, active_contexts = self._parse_xrt_aie_processes(aie_output)
        return self._quick_metrics_from_devices(active_contexts)

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

        examine_output = self._run_xrt_smi(["examine"])
        platform_output = self._run_xrt_smi(["examine", "--report", "platform"])
        aie_output = self._run_xrt_smi(["examine", "--report", "aie-partitions"])

        device_names = self._parse_xrt_device_names(examine_output)
        power_by_bdf = self._parse_xrt_platform_power(platform_output)
        processes_by_bdf, active_contexts = self._parse_xrt_aie_processes(aie_output)
        quick_metrics = self._quick_metrics_from_devices(active_contexts)
        driver_version = self._get_driver_version()

        gpus: List[GPUInfo] = []
        for fallback, accel in enumerate(self._iter_amd_accel_devices()):
            index = self._accel_index(accel, fallback=fallback)
            bdf = self._device_bdf(accel)
            metrics = quick_metrics.get(
                index,
                {
                    "utilization": 0.0,
                    "memory_used": 0.0,
                    "memory_total": 0.0,
                },
            )
            memory_total_mb, memory_source, memory_confidence = (
                self._resolve_memory_budget_mb(accel / "device")
            )
            if metrics.get("memory_total", 0.0) <= 0 and memory_total_mb > 0:
                metrics = dict(metrics)
                metrics["memory_total"] = memory_total_mb

            name = device_names.get(bdf) or f"AMD NPU {index}"
            gpus.append(
                GPUInfo(
                    index=index,
                    name=name,
                    utilization=metrics.get("utilization", 0.0),
                    memory_used=metrics.get("memory_used", 0.0),
                    memory_total=metrics.get("memory_total", 0.0),
                    temperature=0.0,
                    power_draw=power_by_bdf.get(bdf, 0.0),
                    power_limit=0.0,
                    processes=processes_by_bdf.get(bdf, []),
                    driver_version=driver_version,
                    memory_model="shared",
                    memory_source=memory_source,
                    memory_confidence=memory_confidence,
                )
            )

        # xrt-smi can list devices even when accel sysfs is unavailable.
        if not gpus and device_names:
            for index, bdf in enumerate(sorted(device_names.keys())):
                memory_total_mb, memory_source, memory_confidence = (
                    self._resolve_memory_budget_mb(device_path=None)
                )
                gpus.append(
                    GPUInfo(
                        index=index,
                        name=device_names[bdf],
                        utilization=100.0 if active_contexts.get(bdf, 0) > 0 else 0.0,
                        memory_used=0.0,
                        memory_total=memory_total_mb,
                        temperature=0.0,
                        power_draw=power_by_bdf.get(bdf, 0.0),
                        power_limit=0.0,
                        processes=processes_by_bdf.get(bdf, []),
                        driver_version=driver_version,
                        memory_model="shared",
                        memory_source=memory_source,
                        memory_confidence=memory_confidence,
                    )
                )

        return gpus
