#!/usr/bin/env python3
"""Intel-specific GPU monitoring implementation."""

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

from .base import BaseGPUMonitor, GPUInfo


class IntelGPUMonitor(BaseGPUMonitor):
    """Intel GPU monitoring implementation."""

    def __init__(self) -> None:
        """Initialize Intel GPU monitor with immutable driver-version cache."""
        super().__init__()
        self._cached_driver_version = ""
        self._driver_version_initialized = False

    def _find_smi(self) -> Optional[Path]:
        """Find Intel GPU monitoring tool.

        Avoid probing with runtime flags at discovery time because support for
        flags such as ``--version`` differs across intel-gpu-tools releases.
        """
        self.logger.debug("Searching for intel_gpu_top")
        candidate_paths: list[Path] = []

        resolved = shutil.which("intel_gpu_top")
        if resolved:
            candidate_paths.append(Path(resolved))

        for raw_path in ("/usr/bin/intel_gpu_top", "/usr/local/bin/intel_gpu_top"):
            path = Path(raw_path)
            if path not in candidate_paths:
                candidate_paths.append(path)

        for path in candidate_paths:
            if path.exists() and os.access(path, os.X_OK):
                self.logger.info("Found intel_gpu_top at %s", path)
                return path

        self.logger.debug("No executable intel_gpu_top found")
        return None

    def _get_driver_version(self) -> str:
        """Get Intel GPU driver version (cached forever as it doesn't change)."""
        if self._driver_version_initialized:
            return self._cached_driver_version

        version = ""

        # Method 1: sysfs module versions (newer kernels may expose xe)
        for path in (Path("/sys/module/xe/version"), Path("/sys/module/i915/version")):
            if version:
                break
            try:
                if path.exists():
                    value = path.read_text(encoding="utf-8", errors="ignore").strip()
                    if value:
                        version = value
            except Exception as e:
                self.logger.debug("sysfs version probe failed for %s: %s", path, e)

        # Method 2: modinfo fallback (supports built module metadata)
        if not version:
            for module in ("xe", "i915"):
                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:"):
                            maybe_version = line.split(":", 1)[1].strip()
                            if maybe_version:
                                version = maybe_version
                                break
                except Exception as e:
                    self.logger.debug("modinfo %s probe failed: %s", module, e)
                if version:
                    break

        self._cached_driver_version = version
        self._driver_version_initialized = True
        if version:
            self.logger.info(f"Detected Intel driver version: {version}")
        return version

    @staticmethod
    def _coerce_float(value: Any) -> Optional[float]:
        """Convert mixed CLI/JSON numeric payloads to float."""
        if value is None:
            return None
        if isinstance(value, (int, float)):
            return float(value)
        if isinstance(value, dict):
            if "value" in value:
                return IntelGPUMonitor._coerce_float(value.get("value"))
            return None
        if isinstance(value, str):
            match = re.search(r"-?\d+(?:\.\d+)?", value)
            if match:
                try:
                    return float(match.group(0))
                except ValueError:
                    return None
        return None

    @staticmethod
    def _iter_nodes(node: Any, path: str = "") -> List[tuple[str, Any]]:
        """Flatten nested JSON objects/lists into path/value tuples."""
        nodes: list[tuple[str, Any]] = [(path, node)]
        if isinstance(node, dict):
            for key, child in node.items():
                child_path = f"{path}.{key}" if path else str(key)
                nodes.extend(IntelGPUMonitor._iter_nodes(child, child_path))
        elif isinstance(node, list):
            for idx, child in enumerate(node):
                child_path = f"{path}[{idx}]" if path else f"[{idx}]"
                nodes.extend(IntelGPUMonitor._iter_nodes(child, child_path))
        return nodes

    @staticmethod
    def _normalize_memory_mb(value: Any) -> float:
        """Convert process memory values to MB when units are available."""
        numeric = IntelGPUMonitor._coerce_float(value)
        if numeric is None:
            return 0.0

        unit = ""
        if isinstance(value, dict):
            unit = str(value.get("unit") or value.get("units") or "").strip().lower()

        if unit in {"b", "byte", "bytes"}:
            return max(0.0, numeric / (1024.0 * 1024.0))
        if unit in {"kb", "kib"}:
            return max(0.0, numeric / 1024.0)
        if unit in {"gb", "gib"}:
            return max(0.0, numeric * 1024.0)
        if unit in {"mb", "mib"}:
            return max(0.0, numeric)

        # Heuristic fallback when unit metadata is absent.
        if numeric > 1_000_000:
            return max(0.0, numeric / (1024.0 * 1024.0))
        return max(0.0, numeric)

    def _parse_device_list(self, output: str) -> List[Dict[str, Any]]:
        """Parse ``intel_gpu_top -L`` output into index/name records."""
        devices: list[Dict[str, Any]] = []
        used_indices: set[int] = set()

        for line in output.splitlines():
            stripped = line.strip()
            if not stripped:
                continue
            if stripped.lower().startswith("list"):
                continue

            index: Optional[int] = None
            for pattern in (
                r"/dev/dri/card(\d+)",
                r"/dev/dri/renderd(\d+)",
                r"\bcard(\d+)\b",
                r"^(\d+)\s*[:\-]",
            ):
                match = re.search(pattern, stripped, flags=re.IGNORECASE)
                if match:
                    parsed = int(match.group(1))
                    if "renderd" in pattern.lower() and parsed >= 128:
                        index = parsed - 128
                    else:
                        index = parsed
                    break

            name = stripped
            if re.match(
                r"^(?:\d+|card\d+|/dev/dri/card\d+|drm:/dev/dri/card\d+)\s*:",
                stripped,
                flags=re.IGNORECASE,
            ):
                name = stripped.split(":", 1)[1].strip() or stripped

            if index is None:
                # Ignore lines without a stable card/render identifier.
                continue

            if index in used_indices:
                continue

            used_indices.add(index)
            devices.append({"index": index, "name": name})

        return devices

    @staticmethod
    def _read_sysfs_int(path: Path) -> Optional[int]:
        """Read integer value from sysfs."""
        try:
            raw = path.read_text(encoding="utf-8", errors="ignore").strip()
        except (OSError, UnicodeError):
            return None
        if not raw:
            return None
        try:
            return int(raw, 0)
        except ValueError:
            return None

    def _resolve_memory_metrics(
        self, index: int, memory_used_mb: float, memory_total_mb: float
    ) -> tuple[float, float, str, str, str]:
        """Attach memory provenance and shared-memory fallbacks."""
        used_mb = max(0.0, float(memory_used_mb))
        total_mb = max(0.0, float(memory_total_mb))
        if total_mb > 0:
            return used_mb, total_mb, "dedicated", "intel_gpu_top", "direct"

        device_root = Path(f"/sys/class/drm/card{index}/device")
        shared_total_bytes = self._read_sysfs_int(device_root / "mem_info_system_total")
        shared_used_bytes = self._read_sysfs_int(device_root / "mem_info_system_used")

        if shared_total_bytes and shared_total_bytes > 0:
            total_mb = float(shared_total_bytes) / (1024.0 * 1024.0)
            if shared_used_bytes is not None and shared_used_bytes >= 0:
                used_mb = float(shared_used_bytes) / (1024.0 * 1024.0)
                confidence = "direct"
            else:
                confidence = "partial"
            return used_mb, total_mb, "shared", "sysfs:mem_info_system", confidence

        return used_mb, total_mb, "unknown", "intel_gpu_top", "unknown"

    def _parse_json_samples(self, output: str) -> List[Dict[str, Any]]:
        """Parse intel_gpu_top JSON stream output.

        ``intel_gpu_top -J`` emits a stream of JSON objects that may end with
        an incomplete trailing object when the process is interrupted.
        """
        text = output.strip()
        if not text:
            return []

        if text.startswith("["):
            try:
                payload = json.loads(text)
                if isinstance(payload, list):
                    return [entry for entry in payload if isinstance(entry, dict)]
            except json.JSONDecodeError:
                pass

        decoder = json.JSONDecoder()
        samples: list[Dict[str, Any]] = []
        cursor = 0
        text_len = len(text)

        while cursor < text_len:
            while cursor < text_len and text[cursor] in {" ", "\n", "\r", "\t", ","}:
                cursor += 1
            if cursor >= text_len:
                break

            try:
                payload, next_cursor = decoder.raw_decode(text, cursor)
            except json.JSONDecodeError:
                break

            if isinstance(payload, dict):
                samples.append(payload)
            cursor = next_cursor

        return samples

    def _extract_engine_utilization(self, sample: Dict[str, Any]) -> float:
        """Best-effort extraction of per-GPU utilization from JSON sample."""
        busy_entries: list[tuple[str, float]] = []

        engines = sample.get("engines", sample)
        for path, node in self._iter_nodes(engines):
            if not isinstance(node, dict) or "busy" not in node:
                continue
            busy = self._coerce_float(node.get("busy"))
            if busy is None:
                continue
            busy_entries.append((path.lower(), max(0.0, min(100.0, busy))))

        if not busy_entries:
            return 0.0

        preferred = [
            busy
            for path, busy in busy_entries
            if any(tag in path for tag in ("render", "3d", "compute", "rcs", "ccs"))
        ]
        values = preferred or [busy for _path, busy in busy_entries]
        return max(values) if values else 0.0

    def _extract_gpu_index(
        self,
        sample: Dict[str, Any],
        fallback_index: int,
        known_indices: Optional[List[int]] = None,
    ) -> int:
        """Extract GPU index from sample metadata."""
        for key in ("card", "gpu", "drm_minor", "device"):
            value = sample.get(key)
            numeric = self._coerce_float(value)
            if numeric is not None:
                return int(numeric)

            if isinstance(value, str):
                match = re.search(r"card(\d+)", value, flags=re.IGNORECASE)
                if match:
                    return int(match.group(1))

        for key in ("drm-device", "drm_device", "path"):
            value = sample.get(key)
            if isinstance(value, str):
                match = re.search(r"card(\d+)", value, flags=re.IGNORECASE)
                if match:
                    return int(match.group(1))

        if known_indices:
            return known_indices[fallback_index % len(known_indices)]

        # If intel_gpu_top sample metadata omits card IDs, keep everything
        # collapsed under a single logical device to avoid duplicate phantom
        # GPUs from multi-sample streams.
        return 0

    def _sample_has_ai_ml_engines(self, sample: Dict[str, Any]) -> bool:
        """Return whether sample exposes render/compute engines."""
        engines = sample.get("engines", sample)
        for path, node in self._iter_nodes(engines):
            if not isinstance(node, dict):
                continue
            lowered_path = path.lower()
            if any(
                tag in lowered_path for tag in ("render", "3d", "compute", "rcs", "ccs")
            ):
                return True
        return False

    def _extract_processes(self, sample: Dict[str, Any]) -> List[Dict[str, Any]]:
        """Extract process payloads from JSON samples."""
        processes_by_pid: dict[int, Dict[str, Any]] = {}
        clients = sample.get("clients")
        if clients is None:
            return []

        for _path, node in self._iter_nodes(clients):
            if not isinstance(node, dict):
                continue

            pid_raw = node.get("pid") or node.get("client_pid")
            pid = self._coerce_float(pid_raw)
            if pid is None:
                continue
            pid_int = int(pid)
            if pid_int <= 0:
                continue

            name = (
                node.get("name")
                or node.get("process")
                or node.get("comm")
                or node.get("cmd")
                or "unknown"
            )
            memory_raw = (
                node.get("memory")
                or node.get("memory_used")
                or node.get("memory-used")
                or node.get("rss")
                or node.get("vram")
            )
            memory_mb = self._normalize_memory_mb(memory_raw)

            current = processes_by_pid.get(pid_int)
            if current is None or memory_mb > current["memory"]:
                processes_by_pid[pid_int] = {
                    "pid": pid_int,
                    "name": str(name).strip() or "unknown",
                    "memory": memory_mb,
                }

        return sorted(processes_by_pid.values(), key=lambda item: item["pid"])

    def _build_gpu_records(
        self, samples: List[Dict[str, Any]], devices: List[Dict[str, Any]]
    ) -> List[Dict[str, Any]]:
        """Build normalized GPU records from JSON samples and device metadata."""
        records: dict[int, Dict[str, Any]] = {}
        known_indices = sorted(entry["index"] for entry in devices)
        device_names = {entry["index"]: entry["name"] for entry in devices}

        for entry in devices:
            records[entry["index"]] = {
                "index": entry["index"],
                "name": entry["name"],
                "utilization": 0.0,
                "memory_used": 0.0,
                "memory_total": 0.0,
                "temperature": 0.0,
                "power_draw": 0.0,
                "power_limit": 0.0,
                "processes": [],
                "_ai_ml_capable": False,
            }

        for fallback_index, sample in enumerate(samples):
            index = self._extract_gpu_index(
                sample,
                fallback_index=fallback_index,
                known_indices=known_indices,
            )
            record = records.setdefault(
                index,
                {
                    "index": index,
                    "name": device_names.get(index, f"Intel GPU {index}"),
                    "utilization": 0.0,
                    "memory_used": 0.0,
                    "memory_total": 0.0,
                    "temperature": 0.0,
                    "power_draw": 0.0,
                    "power_limit": 0.0,
                    "processes": [],
                    "_ai_ml_capable": False,
                },
            )

            record["utilization"] = self._extract_engine_utilization(sample)
            sample_processes = self._extract_processes(sample)
            if self._sample_has_ai_ml_engines(sample) or sample_processes:
                record["_ai_ml_capable"] = True
            if sample_processes:
                record["processes"] = sample_processes

        ordered_records = [records[idx] for idx in sorted(records.keys())]
        if any(record.get("_ai_ml_capable") for record in ordered_records):
            ordered_records = [
                record for record in ordered_records if record.get("_ai_ml_capable")
            ]
        for record in ordered_records:
            record.pop("_ai_ml_capable", None)
        return ordered_records

    def _parse_legacy_text_output(self, output: str) -> List[Dict[str, Any]]:
        """Fallback parser for older text output formats."""
        if not output:
            return []

        gpus: list[Dict[str, Any]] = []
        gpu_sections = output.split("Intel GPU device")
        for section in gpu_sections[1:]:
            try:
                device_match = re.search(r"(\d+):\s*([^\n]+)", section)
                if not device_match:
                    continue

                index = int(device_match.group(1))
                name = device_match.group(2).strip()
                util_match = re.search(r"Render/3D/0\s+([0-9.]+)%", section)
                utilization = float(util_match.group(1)) if util_match else 0.0

                gpus.append(
                    {
                        "index": index,
                        "name": name,
                        "utilization": utilization,
                        "memory_used": 0.0,
                        "memory_total": 0.0,
                        "temperature": 0.0,
                        "power_draw": 0.0,
                        "power_limit": 0.0,
                        "processes": [],
                    }
                )
            except Exception:
                continue
        return gpus

    def _parse_gpu_info(
        self, output: str, devices: Optional[List[Dict[str, Any]]] = None
    ) -> List[Dict[str, Any]]:
        """Parse intel_gpu_top output for all GPUs."""
        if not output:
            return []

        samples = self._parse_json_samples(output)
        if samples:
            return self._build_gpu_records(samples, devices or [])
        return self._parse_legacy_text_output(output)

    def get_gpu_info(self) -> List[GPUInfo]:
        """Get Intel GPU information for all devices."""
        self.logger.debug("Starting Intel GPU info collection")
        if not self.smi_path:
            self.logger.debug("Skipping Intel GPU info collection: no intel_gpu_top")
            return []

        try:
            # Get driver version once (cached forever)
            driver_version = self._get_driver_version()

            device_output = self._run_smi_command([str(self.smi_path), "-L"]) or ""
            devices = self._parse_device_list(device_output)

            # `intel_gpu_top -J` is a stream; capture one sampled window.
            result = self._run_smi_command([str(self.smi_path), "-J", "-s", "1000"])
            if not result:
                if devices:
                    fallback_gpus: List[GPUInfo] = []
                    for entry in devices:
                        (
                            memory_used_mb,
                            memory_total_mb,
                            memory_model,
                            memory_source,
                            memory_confidence,
                        ) = self._resolve_memory_metrics(
                            entry["index"], memory_used_mb=0.0, memory_total_mb=0.0
                        )
                        fallback_gpus.append(
                            GPUInfo(
                                index=entry["index"],
                                name=entry["name"],
                                utilization=0.0,
                                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=memory_model,
                                memory_source=memory_source,
                                memory_confidence=memory_confidence,
                            )
                        )
                    return fallback_gpus
                self.logger.debug("Intel detailed GPU sampling returned no data")
                return []

            gpu_info = self._parse_gpu_info(result, devices=devices)

            gpus: List[GPUInfo] = []
            for info in gpu_info:
                (
                    memory_used_mb,
                    memory_total_mb,
                    memory_model,
                    memory_source,
                    memory_confidence,
                ) = self._resolve_memory_metrics(
                    index=int(info["index"]),
                    memory_used_mb=float(info.get("memory_used", 0.0)),
                    memory_total_mb=float(info.get("memory_total", 0.0)),
                )
                gpus.append(
                    GPUInfo(
                        index=info["index"],
                        name=info["name"],
                        utilization=info["utilization"],
                        memory_used=memory_used_mb,
                        memory_total=memory_total_mb,
                        temperature=info["temperature"],
                        power_draw=info["power_draw"],
                        power_limit=info.get("power_limit", 0.0),
                        processes=info["processes"],
                        driver_version=driver_version,
                        memory_model=memory_model,
                        memory_source=memory_source,
                        memory_confidence=memory_confidence,
                    )
                )

            self.logger.debug(
                "Successfully created " + str(len(gpus)) + " GPUInfo objects"
            )
            return gpus

        except Exception as e:
            self.logger.error("Failed to get GPU information: " + str(e))
            return []

    def get_quick_metrics(self) -> Dict[int, Dict[str, float]]:
        """Get only essential metrics (utilization, memory) for quick updates.

        Returns:
            Dictionary mapping GPU indices to metric dictionaries
        """
        metrics: Dict[int, Dict[str, float]] = {}

        if not self.smi_path:
            self.logger.debug("Skipping Intel quick metrics: no intel_gpu_top")
            return metrics

        self.logger.debug("Getting quick GPU metrics for Intel GPUs")

        try:
            device_output = self._run_smi_command([str(self.smi_path), "-L"]) or ""
            devices = self._parse_device_list(device_output)

            result = self._run_smi_command([str(self.smi_path), "-J", "-s", "1000"])
            if not result:
                return metrics

            gpu_info = self._parse_gpu_info(result, devices=devices)

            for info in gpu_info:
                index = info.get("index", 0)
                (
                    memory_used_mb,
                    memory_total_mb,
                    _memory_model,
                    _memory_source,
                    _memory_confidence,
                ) = self._resolve_memory_metrics(
                    index=int(index),
                    memory_used_mb=float(info.get("memory_used", 0.0)),
                    memory_total_mb=float(info.get("memory_total", 0.0)),
                )
                # Memory values already in MB
                metrics[index] = {
                    "utilization": info.get("utilization", 0.0),
                    "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,
                }
                self.logger.debug(
                    f"Quick metrics for Intel GPU {index}: {metrics[index]}"
                )

            return metrics

        except Exception as e:
            self.logger.error(f"Error collecting Intel GPU quick metrics: {e}")
            return metrics

    def _run_smi_command(self, cmd: List[str]) -> Optional[str]:
        """Run an Intel GPU monitoring command."""
        try:
            self.logger.debug("Executing command: " + " ".join(cmd))
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=False,
                timeout=5,
            )

            if result.returncode != 0 and not result.stdout.strip():
                self.logger.debug(
                    "Failed command (%s): %s", result.returncode, result.stderr.strip()
                )
                return None

            if result.returncode != 0:
                self.logger.debug(
                    "Non-zero return code %s with partial output preserved",
                    result.returncode,
                )
            return result.stdout

        except subprocess.TimeoutExpired as e:
            raw_partial = e.stdout
            partial: Optional[str]
            if isinstance(raw_partial, bytes):
                partial = raw_partial.decode("utf-8", errors="replace")
            else:
                partial = raw_partial
            if partial and partial.strip():
                self.logger.debug("Command timed out; using partial output")
                return partial
            self.logger.debug(
                "Failed to run command: %s - Error: %s", " ".join(cmd), str(e)
            )
            return None
        except (
            subprocess.SubprocessError,
            FileNotFoundError,
        ) as e:
            self.logger.debug(
                "Failed to run command: " + " ".join(cmd) + " - Error: " + str(e)
            )
            return None
