import logging
import sys
import time
from dataclasses import dataclass
from enum import Enum
from logging import Logger
from multiprocessing import Pipe, Process
from multiprocessing.connection import Connection

from transformers.utils.import_utils import is_cuda_platform, is_rocm_platform


if is_cuda_platform():
    import pynvml

if is_rocm_platform():
    import amdsmi

import psutil
import torch

from transformers.utils import is_torch_accelerator_available


_logger = logging.getLogger(__name__)


# Data class to hold the hardware information
def get_device_name_and_memory_total() -> tuple[str, float]:
    """Returns the name and memory total of GPU 0."""
    device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
    torch_accelerator_module = getattr(torch, device_type, torch.cuda)
    device_name = torch_accelerator_module.get_device_properties(0).name
    device_memory_total = torch_accelerator_module.get_device_properties(0).total_memory / 1024**3
    return device_name, device_memory_total


class HardwareInfo:
    """A class to hold information about the hardware."""

    def __init__(self) -> None:
        # Retrieve GPU stats
        try:
            self.gpu_name, self.gpu_memory_total_gb = get_device_name_and_memory_total()
        except Exception:
            self.gpu_name, self.gpu_memory_total_gb = None, None
        # Retrieve python, torch and CUDA version
        self.python_version = f"{sys.version.split()[0]}"
        self.torch_version = torch.__version__
        if hasattr(torch, "cuda") and torch.cuda.is_available():
            self.cuda_version = torch.version.cuda
        else:
            self.cuda_version = None
        # Retrieve general hardware information
        self.cpu_count = psutil.cpu_count()
        self.memory_total_mb = int(psutil.virtual_memory().total / (1024 * 1024))

    def to_dict(self) -> dict[str, None | int | float | str]:
        return {
            "gpu_name": self.gpu_name,
            "gpu_memory_total_gb": self.gpu_memory_total_gb,
            "python_version": self.python_version,
            "torch_version": self.torch_version,
        }


# Functions to get information about the GPU
def get_amd_gpu_stats(device_handle) -> tuple[int, float]:
    """Get AMD GPU stats using amdsmi library."""
    utilization = amdsmi.amdsmi_get_gpu_activity(device_handle)["gfx_activity"]
    memory_used = amdsmi.amdsmi_get_gpu_vram_usage(device_handle)["vram_used"]
    return int(utilization), float(memory_used) / 1024**3  # Convert bytes to GB


def get_intel_xpu_stats() -> tuple[int, float]:
    """Get Intel XPU stats using torch's Level Zero Sysman bindings."""
    utilization = torch.xpu.utilization(0)
    memory_used = torch.xpu.device_memory_used(0)
    return int(utilization), float(memory_used) / 1024**3  # Convert bytes to GB


def get_nvidia_gpu_stats(device_handle) -> tuple[int, float]:
    """Returns the utilization and memory used of an NVIDIA GPU using pynvml."""
    utilization = pynvml.nvmlDeviceGetUtilizationRates(device_handle).gpu
    memory_info = pynvml.nvmlDeviceGetMemoryInfo(device_handle)
    memory_used_gb = memory_info.used / 1024**3
    return int(utilization), float(memory_used_gb)


# Simple data classes to hold the raw GPU metrics
class GPUMonitoringStatus(Enum):
    """Status of GPU monitoring."""

    SUCCESS = "success"
    FAILED = "failed"
    NO_GPUS_AVAILABLE = "no_gpus_available"
    NO_SAMPLES_COLLECTED = "no_samples_collected"


@dataclass
class GPURawMetrics:
    """Raw values for GPU utilization and memory used."""

    utilization: list[float]  # in percent
    memory_used: list[float]  # in GB
    timestamps: list[float]  # in seconds
    timestamp_0: float  # in seconds
    monitoring_status: GPUMonitoringStatus

    def to_dict(self) -> dict[str, None | int | float | str]:
        return {
            "utilization": self.utilization,
            "memory_used": self.memory_used,
            "timestamps": self.timestamps,
            "timestamp_0": self.timestamp_0,
            "monitoring_status": self.monitoring_status.value,
        }

    @classmethod
    def from_dict(cls, data: dict[str, None | int | float | str]) -> "GPURawMetrics":
        """Create a GPURawMetrics instance from a dictionary."""
        return cls(
            utilization=data["utilization"],
            memory_used=data["memory_used"],
            timestamps=data["timestamps"],
            timestamp_0=data["timestamp_0"],
            monitoring_status=GPUMonitoringStatus(data["monitoring_status"]),
        )


# Main class, used to monitor the GPU utilization during benchmark execution
class GPUMonitor:
    """Monitor GPU utilization during benchmark execution using a separate process."""

    def __init__(self, sample_interval_sec: float = 0.05, logger: Logger | None = None):
        self.sample_interval_sec = sample_interval_sec
        self.logger = logger if logger is not None else _logger
        self.gpu_type = None
        self.process = None

        device_type = torch.accelerator.current_accelerator().type if is_torch_accelerator_available() else "cuda"
        torch_accelerator_module = getattr(torch, device_type, torch.cuda)
        self.num_available_gpus = torch_accelerator_module.device_count()
        if self.num_available_gpus == 0:
            self.logger.warning(f"No GPUs detected by torch.{device_type}.device_count().")
            return

        # Determine GPU type
        device_name, _ = get_device_name_and_memory_total()
        if "amd" in device_name.lower():
            self.gpu_type = "amd"
        elif "nvidia" in device_name.lower():
            self.gpu_type = "nvidia"
        elif "intel" in device_name.lower() or device_type == "xpu":
            self.gpu_type = "intel"
        else:
            self.logger.warning(f"Unsupported GPU for monitoring: {device_name}")

    @staticmethod
    def _monitor_worker(gpu_type: str, sample_interval_sec: float, connection: Connection):
        """Worker process for GPU monitoring."""
        gpu_utilization = []
        gpu_memory_used = []
        timestamps = []
        device_handle = None

        # Initialize GPU-specific monitoring
        if gpu_type == "amd":
            amdsmi.amdsmi_init()
            device_handle = amdsmi.amdsmi_get_processor_handles()[0]
        elif gpu_type == "nvidia":
            pynvml.nvmlInit()
            device_handle = pynvml.nvmlDeviceGetHandleByIndex(0)

        # Signal ready
        try:
            connection.send(0)
        except Exception:
            return

        # Monitoring loop
        stop = False
        while not stop:
            try:
                if gpu_type == "amd":
                    utilization, memory_used = get_amd_gpu_stats(device_handle)
                elif gpu_type == "nvidia":
                    utilization, memory_used = get_nvidia_gpu_stats(device_handle)
                elif gpu_type == "intel":
                    utilization, memory_used = get_intel_xpu_stats()
                else:
                    break

                gpu_utilization.append(utilization)
                gpu_memory_used.append(memory_used)
                timestamps.append(time.time())
            except Exception as e:
                # Skips failed measurements
                _logger.debug(f"Failed to collect GPU metrics sample: {e}")

            stop = connection.poll(sample_interval_sec)

        # Cleanup
        if gpu_type == "amd":
            try:
                amdsmi.amdsmi_shut_down()
            except Exception as e:
                _logger.debug(f"Failed to shutdown AMD GPU monitoring: {e}")
        elif gpu_type == "nvidia":
            try:
                pynvml.nvmlShutdown()
            except Exception as e:
                _logger.debug(f"Failed to shutdown NVIDIA GPU monitoring: {e}")

        # Send results back
        try:
            connection.send((gpu_utilization, gpu_memory_used, timestamps))
        except Exception as e:
            _logger.error(f"Failed to send GPU monitoring results: {e}")

        connection.close()

    def start(self):
        """Start monitoring GPU metrics in a separate process."""
        if self.gpu_type is None:
            self.logger.debug("GPU monitoring skipped (no supported GPU)")
            return

        self.child_connection, self.parent_connection = Pipe()
        self.process = Process(
            target=GPUMonitor._monitor_worker,
            args=(self.gpu_type, self.sample_interval_sec, self.child_connection),
            daemon=True,
        )
        self.process.start()

        # Wait for worker to signal ready
        if self.process.is_alive():
            self.parent_connection.recv()
        self.logger.debug("GPU monitoring started (multiprocessing)")

    def stop_and_collect(self) -> GPURawMetrics:
        """Stop monitoring and return collected metrics."""
        # No GPU available or unsupported GPU
        if self.process is None:
            return GPURawMetrics(
                utilization=[],
                memory_used=[],
                timestamps=[],
                timestamp_0=0.0,
                monitoring_status=GPUMonitoringStatus.NO_GPUS_AVAILABLE,
            )

        # Process crashed before we could collect results
        process_failed = False
        if not self.process.is_alive():
            process_failed = True
            gpu_utilization, gpu_memory_used, timestamps = [], [], []
        else:
            # Signal stop
            self.parent_connection.send(0)
            # Get results
            try:
                gpu_utilization, gpu_memory_used, timestamps = self.parent_connection.recv()
            except Exception:
                process_failed = True
                gpu_utilization, gpu_memory_used, timestamps = [], [], []

        self.parent_connection.close()
        self.process.join(timeout=2.0)
        if self.process.is_alive():
            self.process.terminate()

        if gpu_utilization:
            timestamp_0 = timestamps[0]
            metrics = GPURawMetrics(
                utilization=gpu_utilization,
                memory_used=gpu_memory_used,
                timestamps=[t - timestamp_0 for t in timestamps],
                timestamp_0=timestamp_0,
                monitoring_status=GPUMonitoringStatus.SUCCESS,
            )
            self.logger.debug(f"GPU monitoring completed: {len(gpu_utilization)} samples collected")
        elif process_failed:
            metrics = GPURawMetrics(
                utilization=[],
                memory_used=[],
                timestamps=[],
                timestamp_0=0.0,
                monitoring_status=GPUMonitoringStatus.FAILED,
            )
            self.logger.warning("GPU monitoring failed (process crashed or timed out)")
        else:
            metrics = GPURawMetrics(
                utilization=[],
                memory_used=[],
                timestamps=[],
                timestamp_0=0.0,
                monitoring_status=GPUMonitoringStatus.NO_SAMPLES_COLLECTED,
            )
        return metrics