#!/usr/bin/env python3
"""Apple/macOS GPU monitoring implementation (experimental)."""

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

from .base import BaseGPUMonitor, GPUInfo


class AppleGPUMonitor(BaseGPUMonitor):
    """Apple GPU monitor using ``system_profiler`` inventory data.

    Notes:
        - macOS does not expose a stable, non-privileged equivalent to vendor
          SMI tools for per-process GPU telemetry across all hardware.
        - This monitor focuses on safe hardware inventory. Live utilization is
          opt-in via ``powermetrics`` to avoid expensive/privileged probes in
          the default lightweight path.
    """

    def __init__(
        self,
        cache_ttl_seconds: float = 30.0,
        enable_live_metrics: Optional[bool] = None,
    ) -> None:
        """Initialize Apple GPU monitor with a lightweight cache."""
        super().__init__()
        self._cache_ttl_seconds = max(5.0, float(cache_ttl_seconds))
        self._cached_gpu_info: List[GPUInfo] = []
        self._cached_at = 0.0
        self._powermetrics_path = self._find_powermetrics()
        if enable_live_metrics is None:
            raw_env = os.environ.get("AITOP_ENABLE_MACOS_POWERMETRICS", "")
            enable_live_metrics = raw_env.strip().lower() in {
                "1",
                "true",
                "yes",
                "on",
            }
        self._enable_live_metrics = bool(
            enable_live_metrics and self._powermetrics_path
        )

    def _find_smi(self) -> Optional[Path]:
        """Locate ``system_profiler`` executable on macOS."""
        resolved = shutil.which("system_profiler")
        if resolved:
            return Path(resolved)

        fallback = Path("/usr/sbin/system_profiler")
        if fallback.exists():
            return fallback
        return None

    @staticmethod
    def _find_powermetrics() -> Optional[Path]:
        """Locate optional ``powermetrics`` binary for live utilization."""
        resolved = shutil.which("powermetrics")
        if resolved:
            return Path(resolved)
        fallback = Path("/usr/bin/powermetrics")
        if fallback.exists() and os.access(fallback, os.X_OK):
            return fallback
        return None

    @staticmethod
    def _parse_memory_mb(raw_value: Any) -> float:
        """Parse memory strings such as ``8 GB`` or ``1536 MB`` into MB."""
        if raw_value is None:
            return 0.0

        if isinstance(raw_value, (int, float)):
            numeric = float(raw_value)
            # Large integers are likely bytes from lower-level APIs.
            if numeric > 1_000_000:
                return max(0.0, numeric / (1024.0 * 1024.0))
            return max(0.0, numeric)

        text = str(raw_value).strip()
        if not text:
            return 0.0

        match = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*(tb|gb|mb|kb|b)\b", text.lower())
        if not match:
            # Fallback: parse first number and assume MB.
            fallback = re.search(r"([0-9]+(?:\.[0-9]+)?)", text)
            if not fallback:
                return 0.0
            return max(0.0, float(fallback.group(1)))

        value = float(match.group(1))
        unit = match.group(2)
        multipliers = {
            "tb": 1024.0 * 1024.0,
            "gb": 1024.0,
            "mb": 1.0,
            "kb": 1.0 / 1024.0,
            "b": 1.0 / (1024.0 * 1024.0),
        }
        return max(0.0, value * multipliers[unit])

    @staticmethod
    def _extract_display_items(payload: Any) -> List[Dict[str, Any]]:
        """Extract display entries from ``system_profiler -xml`` payload."""
        if not isinstance(payload, list):
            return []

        items: List[Dict[str, Any]] = []
        for section in payload:
            if not isinstance(section, dict):
                continue
            raw_items = section.get("_items")
            if not isinstance(raw_items, list):
                continue
            for raw_item in raw_items:
                if isinstance(raw_item, dict):
                    items.append(raw_item)
        return items

    @staticmethod
    def _looks_like_gpu_entry(item: Dict[str, Any]) -> bool:
        """Return whether a display entry likely represents a GPU adapter."""
        text_candidates = (
            str(item.get("sppci_model", "")),
            str(item.get("spdisplays_vendor", "")),
            str(item.get("_name", "")),
            str(item.get("spdisplays_metal", "")),
        )
        blob = " ".join(text_candidates).lower()
        if not blob.strip():
            return False
        return any(
            token in blob
            for token in (
                "apple",
                "gpu",
                "radeon",
                "nvidia",
                "geforce",
                "intel",
                "metal",
            )
        )

    @staticmethod
    def _build_name(item: Dict[str, Any]) -> str:
        """Build readable adapter name from profiler fields."""
        model = str(
            item.get("sppci_model")
            or item.get("_name")
            or item.get("spdisplays_display-product-name")
            or "Apple GPU"
        ).strip()
        vendor = str(item.get("spdisplays_vendor", "")).strip()

        if vendor and vendor.lower() not in model.lower():
            return f"{vendor} {model}"
        return model

    @staticmethod
    def _clone_gpu_info(info: GPUInfo) -> GPUInfo:
        """Clone ``GPUInfo`` to avoid cross-thread/shared mutations."""
        return GPUInfo(
            index=info.index,
            name=info.name,
            utilization=float(info.utilization),
            memory_used=float(info.memory_used),
            memory_total=float(info.memory_total),
            temperature=float(info.temperature),
            power_draw=float(info.power_draw),
            power_limit=float(info.power_limit),
            processes=[dict(process) for process in info.processes],
            driver_version=info.driver_version,
            cuda_version=info.cuda_version,
            rocm_version=info.rocm_version,
        )

    @staticmethod
    def _parse_powermetrics_utilization(output: str) -> Optional[float]:
        """Parse GPU utilization percentage from powermetrics output."""
        candidates: List[float] = []
        for line in output.splitlines():
            lowered = line.lower()
            if "gpu" not in lowered:
                continue
            if not any(
                token in lowered
                for token in ("active", "residency", "utilization", "busy")
            ):
                continue

            match = re.search(r"([0-9]+(?:\.[0-9]+)?)\s*%", lowered)
            if not match:
                continue
            value = max(0.0, min(100.0, float(match.group(1))))
            candidates.append(value)

        if not candidates:
            return None
        return max(candidates)

    def _collect_powermetrics_utilization(self) -> Optional[float]:
        """Collect one-shot utilization from powermetrics when enabled."""
        if not self._enable_live_metrics or not self._powermetrics_path:
            return None

        try:
            result = subprocess.run(
                [
                    str(self._powermetrics_path),
                    "--samplers",
                    "gpu_power",
                    "-n",
                    "1",
                    "-i",
                    "1000",
                ],
                capture_output=True,
                text=True,
                timeout=4.0,
                check=False,
            )
        except Exception as error:
            self.logger.debug("powermetrics execution failed: %s", error)
            return None

        output = "\n".join(
            part for part in (result.stdout, result.stderr) if isinstance(part, str)
        )
        if not output.strip():
            return None
        return self._parse_powermetrics_utilization(output)

    def _load_inventory(self) -> List[GPUInfo]:
        """Load GPU inventory from ``system_profiler`` plist output."""
        smi_path = self.smi_path
        if not smi_path:
            return []

        try:
            result = subprocess.run(
                [str(smi_path), "SPDisplaysDataType", "-xml"],
                capture_output=True,
                text=False,
                timeout=8.0,
                check=False,
            )
        except Exception as error:
            self.logger.debug("system_profiler execution failed: %s", error)
            return []

        if result.returncode != 0 or not result.stdout:
            self.logger.debug(
                "system_profiler returned no display data (rc=%s)",
                result.returncode,
            )
            return []

        try:
            payload = plistlib.loads(result.stdout)
        except Exception as error:
            self.logger.debug("Failed to parse system_profiler plist: %s", error)
            return []

        inventory: List[GPUInfo] = []
        for item in self._extract_display_items(payload):
            if not self._looks_like_gpu_entry(item):
                continue

            memory_total = 0.0
            for key in (
                "spdisplays_vram_shared",
                "spdisplays_vram",
                "spdisplays_vram_dynamic",
                "spdisplays_mem",
            ):
                parsed_mb = self._parse_memory_mb(item.get(key))
                if parsed_mb > 0:
                    memory_total = parsed_mb
                    break

            inventory.append(
                GPUInfo(
                    index=len(inventory),
                    name=self._build_name(item),
                    utilization=0.0,
                    memory_used=0.0,
                    memory_total=memory_total,
                    temperature=0.0,
                    power_draw=0.0,
                    power_limit=0.0,
                    processes=[],
                )
            )

        return inventory

    def get_gpu_info(self) -> List[GPUInfo]:
        """Get Apple GPU inventory (best-effort, cached)."""
        now = time.time()
        if self._cached_gpu_info and now - self._cached_at < self._cache_ttl_seconds:
            return [self._clone_gpu_info(gpu) for gpu in self._cached_gpu_info]

        self._cached_gpu_info = self._load_inventory()
        self._cached_at = now
        return [self._clone_gpu_info(gpu) for gpu in self._cached_gpu_info]

    def get_quick_metrics(self) -> Dict[int, Dict[str, float]]:
        """Return lightweight runtime metrics.

        Live metrics are optional via ``powermetrics`` and disabled by default.
        """
        utilization = self._collect_powermetrics_utilization()
        if utilization is None:
            return {}
        if not self._cached_gpu_info:
            return {}

        return {
            gpu.index: {"utilization": utilization} for gpu in self._cached_gpu_info
        }
