#!/usr/bin/env python3
"""Base GPU monitoring functionality."""

import logging
import os
from abc import ABC, abstractmethod
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, List, Optional, Tuple


@dataclass
class GPUInfo:
    """Container for GPU information.

    All memory values (memory_used, memory_total, and processes memory) are
    in MB (megabytes). Framework version fields store vendor-specific driver
    and toolkit versions (e.g., CUDA, ROCm).
    """

    index: int
    name: str
    utilization: float  # GPU utilization percentage (0-100)
    memory_used: float  # In MB (megabytes)
    memory_total: float  # In MB (megabytes)
    temperature: float  # In degrees Celsius
    power_draw: float  # In watts
    power_limit: float  # In watts
    processes: List[Dict[str, Any]]  # Each process dict includes 'memory' in MB
    # Framework version information (optional, populated per vendor)
    driver_version: str = ""
    cuda_version: str = ""  # NVIDIA only
    rocm_version: str = ""  # AMD only
    # Memory provenance metadata:
    # - memory_model: dedicated | shared | hybrid | unknown
    # - memory_source: backend source string (sysfs/smi/etc)
    # - memory_confidence: direct | estimated | partial | unknown
    memory_model: str = "dedicated"
    memory_source: str = ""
    memory_confidence: str = "direct"
    memory_dedicated_used: float = 0.0
    memory_dedicated_total: float = 0.0
    memory_shared_used: float = 0.0
    memory_shared_total: float = 0.0

    @staticmethod
    def _safe_float(value: Any) -> float:
        """Convert values to non-negative floats for UI math."""
        try:
            return max(0.0, float(value))
        except (TypeError, ValueError):
            return 0.0

    def _shared_memory_is_estimated(self) -> bool:
        """Return True when total memory is a shared-memory estimate."""
        model = str(self.memory_model).strip().lower()
        confidence = str(self.memory_confidence).strip().lower()
        return model == "shared" and confidence not in {"", "direct"}

    def memory_process_used_mb(self) -> float:
        """Return the best-effort aggregate process memory footprint in MB."""
        total = 0.0
        for process in self.processes:
            if not isinstance(process, dict):
                continue
            total += self._safe_float(process.get("memory", 0.0))
        return total

    def memory_display_used_mb(self) -> float:
        """Return the memory value that should drive display and bar math."""
        memory_used_mb = self._safe_float(self.memory_used)
        if memory_used_mb > 0:
            return memory_used_mb

        if self._shared_memory_is_estimated():
            return self.memory_process_used_mb()

        return memory_used_mb

    def memory_hard_total_mb(self) -> float:
        """Return the absolute total/budget when the backend reports one."""
        return self._safe_float(self.memory_total)

    def memory_breakdown_totals_mb(self) -> Tuple[float, float]:
        """Return known dedicated/shared totals in MB."""
        return (
            self._safe_float(self.memory_dedicated_total),
            self._safe_float(self.memory_shared_total),
        )

    def has_memory_breakdown(self) -> bool:
        """Return True when dedicated and shared memory totals are both known."""
        dedicated_total_mb, shared_total_mb = self.memory_breakdown_totals_mb()
        return dedicated_total_mb > 0.0 and shared_total_mb > 0.0

    def memory_visual_summary(self) -> str:
        """Return a short memory-budget summary for shared/hybrid devices."""
        hard_total_mb = self.memory_hard_total_mb()
        if self.uses_soft_memory_limit():
            soft_total_gb = self.memory_soft_total_mb() / 1024.0
            hard_total_gb = hard_total_mb / 1024.0
            return f"{soft_total_gb:.1f}G soft / {hard_total_gb:.1f}G host"

        dedicated_total_mb, shared_total_mb = self.memory_breakdown_totals_mb()
        if dedicated_total_mb > 0.0 and shared_total_mb > 0.0:
            return (
                f"{dedicated_total_mb / 1024.0:.1f}G dedicated + "
                f"{shared_total_mb / 1024.0:.1f}G dynamic"
            )

        if str(self.memory_model).strip().lower() == "shared" and hard_total_mb > 0.0:
            return f"{hard_total_mb / 1024.0:.1f}G shared"

        return ""

    def memory_soft_total_mb(self) -> float:
        """Return a softer visualization budget for shared estimated memory.

        Dedicated accelerators and shared-memory devices with direct totals use
        the reported total directly. For shared-memory NPUs without a direct cap,
        the bar uses a coarse bucketed budget while the hard ceiling remains the
        host-RAM estimate reported in ``memory_total``.
        """
        hard_total_mb = self.memory_hard_total_mb()
        if hard_total_mb <= 0:
            return 0.0

        if not self._shared_memory_is_estimated():
            return hard_total_mb

        minimum_soft_mb = min(hard_total_mb, 4096.0)
        display_used_mb = self.memory_display_used_mb()
        if display_used_mb <= 0:
            return minimum_soft_mb

        target_mb = max(minimum_soft_mb, display_used_mb * 1.25)
        bucket_mb = 1024.0
        while bucket_mb < target_mb and bucket_mb < hard_total_mb:
            bucket_mb *= 2.0

        return min(hard_total_mb, max(minimum_soft_mb, bucket_mb))

    def uses_soft_memory_limit(self) -> bool:
        """Return True when the memory bar should show a soft budget."""
        soft_total_mb = self.memory_soft_total_mb()
        hard_total_mb = self.memory_hard_total_mb()
        return (
            self._shared_memory_is_estimated()
            and soft_total_mb > 0
            and hard_total_mb > soft_total_mb
        )

    def memory_bar_total_mb(self) -> float:
        """Return the denominator used by usage bars and percentages."""
        soft_total_mb = self.memory_soft_total_mb()
        if soft_total_mb > 0:
            return soft_total_mb
        return self.memory_hard_total_mb()


class BaseGPUMonitor(ABC):
    """Abstract base class for GPU monitoring."""

    def __init__(self) -> None:
        """Initialize the GPU monitor with common setup."""
        self.logger = logging.getLogger(self.__class__.__name__)
        self.logger.debug("Initializing GPU monitor")
        # Initialize with None and defer actual finding until first use
        self._smi_path: Optional[Path] = None
        self._smi_path_initialized = False

    @property
    def smi_path(self) -> Optional[Path]:
        """Get the SMI path lazily."""
        # Only find SMI path when first accessed, not during initialization
        if not self._smi_path_initialized:
            self._smi_path = self._find_smi()
            self._smi_path_initialized = True

            if self._smi_path:
                self.logger.info(f"Found SMI tool at: {self._smi_path}")
            else:
                self.logger.warning("SMI tool not found")

        return self._smi_path

    @abstractmethod
    def _find_smi(self) -> Optional[Path]:
        """Find the vendor-specific SMI tool."""
        pass

    @abstractmethod
    def get_gpu_info(self) -> List[GPUInfo]:
        """Get GPU information for this vendor."""
        pass

    @abstractmethod
    def get_quick_metrics(self) -> Dict[int, Dict[str, float]]:
        """Get lightweight per-GPU metrics for frequent refresh."""
        pass

    def _run_smi_command(self, cmd: List[str]) -> Optional[str]:
        """Run an SMI command and return its output."""
        try:
            import subprocess

            # First check if we have a valid command
            if not self.smi_path:
                self.logger.debug("No SMI path available")
                return None

            # If it's a full path, check if it exists and is executable
            if (
                str(self.smi_path) != self.smi_path.name
            ):  # If path contains directory components
                if not self.smi_path.exists():
                    self.logger.debug(f"SMI path does not exist: {self.smi_path}")
                    return None
                if not os.access(self.smi_path, os.X_OK):
                    self.logger.debug(f"SMI path is not executable: {self.smi_path}")
                    return None

            self.logger.debug(f"Running SMI command: {' '.join(cmd)}")

            # Run the command with a timeout to prevent hanging
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                check=True,
                timeout=5,  # 5 second timeout
            )

            self.logger.debug(f"SMI command output: {result.stdout[:200]}...")
            return result.stdout

        except subprocess.CalledProcessError as e:
            self.logger.error(
                f"SMI command failed with return code {e.returncode}: {e.stderr}"
            )
            return None
        except subprocess.TimeoutExpired:
            self.logger.error("SMI command timed out after 5 seconds")
            return None
        except FileNotFoundError:
            self.logger.error(f"SMI command not found: {cmd[0]}")
            return None
        except PermissionError:
            self.logger.error(f"Permission denied running SMI command: {cmd[0]}")
            return None
        except Exception as e:
            self.logger.error(f"Unexpected error running SMI command: {str(e)}")
            return None
