#!/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,_.]*(?:\.[0-9]+)?)\s*"
            r"(tib|tb|gib|gb|mib|mb|kib|kb|bytes|byte|b)\b",
            text.lower(),
        )
        if not match:
            # Fallback: parse first number and assume MB.
            fallback = re.search(r"([0-9][0-9,_.]*(?:\.[0-9]+)?)", text)
            if not fallback:
                return 0.0
            parsed_fallback = AppleGPUMonitor._parse_number_token(fallback.group(1))
            if parsed_fallback is None:
                return 0.0
            return max(0.0, parsed_fallback)

        value = AppleGPUMonitor._parse_number_token(match.group(1))
        if value is None:
            return 0.0
        unit = match.group(2)
        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[unit])

    @staticmethod
    def _parse_number_token(token: str) -> Optional[float]:
        """Parse locale-flavored numeric tokens into float."""
        cleaned = str(token).strip().replace("_", "")
        if not cleaned:
            return None

        if "," in cleaned and "." not in cleaned:
            parts = cleaned.split(",")
            # Treat "1,5" as decimal comma, but "1,536" as thousands separator.
            if len(parts) == 2 and len(parts[1]) in {1, 2}:
                cleaned = cleaned.replace(",", ".")
            else:
                cleaned = cleaned.replace(",", "")
        else:
            cleaned = cleaned.replace(",", "")

        try:
            return float(cleaned)
        except ValueError:
            return None

    @staticmethod
    def _extract_memory_total_mb(item: Dict[str, Any]) -> float:
        """Extract best-effort VRAM/unified-memory total from display payload."""
        preferred_keys = (
            "spdisplays_vram_shared",
            "spdisplays_vram",
            "spdisplays_vram_dynamic",
            "spdisplays_mem",
            "_spdisplays_vram",
            "_spdisplays_vram_shared",
            "_spdisplays_vram_dynamic",
        )
        for key in preferred_keys:
            parsed_mb = AppleGPUMonitor._parse_memory_mb(item.get(key))
            if parsed_mb > 0:
                return parsed_mb

        fallback = 0.0
        for raw_key, raw_value in item.items():
            key = str(raw_key).lower()
            if not any(token in key for token in ("vram", "memory", "mem")):
                continue
            parsed_mb = AppleGPUMonitor._parse_memory_mb(raw_value)
            if parsed_mb > 0:
                fallback = max(fallback, parsed_mb)
        return fallback

    @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,
            memory_model=info.memory_model,
            memory_source=info.memory_source,
            memory_confidence=info.memory_confidence,
        )

    @staticmethod
    def _parse_powermetrics_utilization(output: str) -> Optional[float]:
        """Parse GPU utilization percentage from powermetrics output."""
        candidates: List[float] = []
        fallback_residency: List[float] = []

        for line in output.splitlines():
            lowered = line.lower()
            if "gpu" not in lowered:
                continue

            is_active_residency = "active" in lowered and "residency" in lowered
            is_idle_residency = "idle" in lowered and "residency" in lowered
            is_utilization = "utilization" in lowered or "utilisation" in lowered
            is_busy = "busy" in lowered
            is_any_residency = "residency" in lowered

            if not (
                is_active_residency
                or is_utilization
                or is_busy
                or is_any_residency
            ):
                continue
            if is_idle_residency and not (is_utilization or is_busy):
                continue

            match = re.search(r"([0-9][0-9,_.]*(?:\.[0-9]+)?|\.[0-9]+)\s*%", lowered)
            if not match:
                continue
            parsed = AppleGPUMonitor._parse_number_token(match.group(1))
            if parsed is None:
                continue
            value = max(0.0, min(100.0, parsed))

            if is_active_residency or is_utilization or is_busy:
                candidates.append(value)
            elif is_any_residency:
                fallback_residency.append(value)

        if candidates:
            return max(candidates)
        if fallback_residency:
            return max(fallback_residency)
        return None

    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:
            env = dict(os.environ)
            env.setdefault("LC_ALL", "C")
            env.setdefault("LANG", "C")
            result = subprocess.run(
                [
                    str(self._powermetrics_path),
                    "--samplers",
                    "gpu_power",
                    "-n",
                    "1",
                    "-i",
                    "1000",
                ],
                capture_output=True,
                text=True,
                timeout=4.0,
                check=False,
                env=env,
            )
        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 = self._extract_memory_total_mb(item)

            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=[],
                    memory_model="shared",
                    memory_source="system_profiler",
                    memory_confidence="direct" if memory_total > 0 else "partial",
                )
            )

        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
        }
