"""AMD GPU configuration handler.

This module provides AMD-specific GPU configuration for different
container runtimes using ROCm.
"""

from __future__ import annotations

from pathlib import Path

from gaol.gpu.models import DeviceMapping, GPUConfig


class AmdGPU:
    """AMD GPU configuration handler.

    Provides methods to generate runtime-specific GPU configurations
    for AMD GPUs using ROCm.

    Example:
        >>> config = GPUConfig(enabled=True, device_ids=[0])
        >>> docker_config = AmdGPU.get_docker_config(config)
        >>> "/dev/kfd" in docker_config["devices"]
        True
    """

    # Standard AMD ROCm device paths
    KFD_DEVICE = "/dev/kfd"
    DRI_PATH = Path("/dev/dri")

    @staticmethod
    def get_docker_config(gpu_config: GPUConfig) -> dict:
        """Get Docker-specific GPU configuration for AMD.

        AMD GPUs use device mappings rather than device_requests.
        Requires ROCm-compatible container images.

        Args:
            gpu_config: GPU configuration.

        Returns:
            Dictionary with devices and environment variables.
        """
        devices = AmdGPU.get_required_devices(gpu_config.device_ids)

        # Environment variables for ROCm
        environment = {
            "HSA_OVERRIDE_GFX_VERSION": "",  # May need to be set for specific GPUs
        }

        # Set visible devices if specific ones requested
        if gpu_config.device_ids:
            environment["HIP_VISIBLE_DEVICES"] = ",".join(
                str(d) for d in gpu_config.device_ids
            )
            environment["ROCR_VISIBLE_DEVICES"] = ",".join(
                str(d) for d in gpu_config.device_ids
            )

        return {
            "devices": devices,
            "environment": environment,
        }

    @staticmethod
    def get_nspawn_config(gpu_config: GPUConfig) -> dict:
        """Get nspawn-specific GPU configuration for AMD.

        Args:
            gpu_config: GPU configuration.

        Returns:
            Dictionary with bind mounts and environment variables.
        """
        binds = AmdGPU.get_required_devices(gpu_config.device_ids)

        environment: dict[str, str] = {}
        if gpu_config.device_ids:
            environment["HIP_VISIBLE_DEVICES"] = ",".join(
                str(d) for d in gpu_config.device_ids
            )
            environment["ROCR_VISIBLE_DEVICES"] = ",".join(
                str(d) for d in gpu_config.device_ids
            )

        return {
            "binds": binds,
            "environment": environment,
        }

    @staticmethod
    def get_libvirt_config(
        gpu_config: GPUConfig,  # noqa: ARG004
        pci_addresses: list[str],
    ) -> dict:
        """Get libvirt-specific GPU configuration for AMD.

        Args:
            gpu_config: GPU configuration (reserved for future use).
            pci_addresses: List of PCI addresses.

        Returns:
            Dictionary with XML snippets.
        """
        xml_parts: list[str] = []

        for pci_addr in pci_addresses:
            # Parse PCI address
            parts = pci_addr.replace(":", ".").split(".")
            if len(parts) >= 4:
                domain = parts[0]
                bus = parts[1]
                slot = parts[2]
                function = parts[3]

                xml_parts.append(
                    f"""    <hostdev mode='subsystem' type='pci' managed='yes'>
      <source>
        <address domain='0x{domain}' bus='0x{bus}' slot='0x{slot}' function='0x{function}'/>
      </source>
    </hostdev>"""
                )

        return {
            "hostdev_xml": "\n".join(xml_parts),
            "requires_iommu": True,
        }

    @staticmethod
    def get_required_devices(device_ids: list[int] | None = None) -> list[str]:
        """Get list of required AMD device paths.

        Args:
            device_ids: Specific GPU device IDs, or None for all.

        Returns:
            List of device paths.
        """
        devices: list[str] = []

        # Add KFD device (required for ROCm)
        if Path(AmdGPU.KFD_DEVICE).exists():
            devices.append(AmdGPU.KFD_DEVICE)

        # Add DRI render nodes
        if AmdGPU.DRI_PATH.exists():
            if device_ids:
                # Map device IDs to render nodes
                for device_id in device_ids:
                    # renderD128 is typically GPU 0, renderD129 is GPU 1, etc.
                    render_node = AmdGPU.DRI_PATH / f"renderD{128 + device_id}"
                    if render_node.exists():
                        devices.append(str(render_node))

                    # Also add corresponding card device
                    card_node = AmdGPU.DRI_PATH / f"card{device_id}"
                    if card_node.exists():
                        devices.append(str(card_node))
            else:
                # Add all render and card devices
                for path in sorted(AmdGPU.DRI_PATH.iterdir()):
                    if path.name.startswith("renderD") or path.name.startswith("card"):
                        devices.append(str(path))

        return devices

    @staticmethod
    def get_device_mappings(
        device_ids: list[int] | None = None,
    ) -> list[DeviceMapping]:
        """Get device mappings for AMD GPUs.

        Args:
            device_ids: Specific GPU device IDs, or None for all.

        Returns:
            List of DeviceMapping objects.
        """
        mappings: list[DeviceMapping] = []

        for device_path in AmdGPU.get_required_devices(device_ids):
            mappings.append(
                DeviceMapping(
                    host_path=device_path,
                    container_path=device_path,
                    permissions="rwm",
                )
            )

        return mappings

    @staticmethod
    def get_required_capabilities() -> list[str]:
        """Get Linux capabilities required for AMD GPU access.

        Returns:
            List of capability names.
        """
        return ["SYS_ADMIN"]

    @staticmethod
    def get_render_node_for_gpu(gpu_index: int) -> str | None:
        """Get the render node path for a specific GPU index.

        Args:
            gpu_index: GPU device index (0, 1, 2, etc.)

        Returns:
            Path to render node or None if not found.
        """
        render_node = AmdGPU.DRI_PATH / f"renderD{128 + gpu_index}"
        if render_node.exists():
            return str(render_node)
        return None
