"""Volume synchronization between Mithril and SkyPilot.

This module ensures that volumes referenced in a SkyPilot task are registered
with SkyPilot before launch. It checks volumes against the Mithril API and
automatically registers any that exist in Mithril but aren't yet known to SkyPilot.
"""

from __future__ import annotations

from typing import TYPE_CHECKING

from mithril.api.bindings.models.volume_model_interface import VolumeModelInterface

# SkyPilot volume type constants from sky.utils.volume.VolumeType enum
MITHRIL_FILE_SHARE_TYPE = "mithril-file-share"
MITHRIL_BLOCK_TYPE = "mithril-block"

if TYPE_CHECKING:
    from mithril.api.bindings.models import VolumeModel
    from mithril.api.client import MithrilClient
    from mithril.sky import SkyClient


def _get_skypilot_volumes(sky: SkyClient) -> set[str]:
    """Get volume names registered with SkyPilot.

    Args:
        sky: SkyClient instance.

    Returns:
        Set of volume names currently registered in SkyPilot.
    """
    request_id = sky.volumes.ls()
    volumes = sky.get(request_id)
    return {vol.name for vol in volumes if vol.name}


def _register_volume_with_skypilot(vol: VolumeModel, sky: SkyClient) -> None:
    """Register a Mithril volume with SkyPilot.

    Args:
        vol: VolumeModel from Mithril API.
        sky: SkyClient instance.

    Raises:
        Exception: If registration fails. sky.get() will raise RuntimeError,
            RequestCancelled, or the original exception from the request.
    """
    # Map Mithril volume to SkyPilot volume spec
    match vol.interface:
        case VolumeModelInterface.FILE:
            volume_type = MITHRIL_FILE_SHARE_TYPE
        case VolumeModelInterface.BLOCK:
            volume_type = MITHRIL_BLOCK_TYPE

    # Note: SkyPilot treats GB and Gi identically (both map to 2**30 bytes)
    sky_volume = sky.volumes.Volume(
        name=vol.name,
        type=volume_type,
        infra=f"mithril/{vol.region}",
        size=f"{vol.capacity_gb}GB",
    )
    request_id = sky.volumes.apply(sky_volume)
    sky.get(request_id)


def ensure_volumes_registered(
    volume_names: list[str],
    *,
    mithril: MithrilClient,
    sky: SkyClient,
) -> list[str]:
    """Ensure task volumes are registered with SkyPilot.

    For each volume name:
    1. Check if already registered with SkyPilot
    2. If not, check if it exists in Mithril
    3. If found in Mithril, register it with SkyPilot

    Args:
        volume_names: List of volume names to ensure are registered.
        mithril: MithrilClient instance for Mithril API operations.
        sky: SkyClient instance for SkyPilot operations.

    Returns:
        List of volume names that were newly registered.
    """
    if not volume_names:
        return []

    volume_names_set = set(volume_names)

    # Check what's already in SkyPilot
    sky_volumes = _get_skypilot_volumes(sky)
    missing = volume_names_set - sky_volumes

    if not missing:
        return []

    mithril_volumes = mithril.list_volumes()
    mithril_by_name = {v.name: v for v in mithril_volumes}

    registered: list[str] = []
    for name in missing:
        if name in mithril_by_name:
            vol = mithril_by_name[name]
            _register_volume_with_skypilot(vol, sky)
            registered.append(name)

    return registered
