"""Seeweb Cloud

History:
@ Aug 6, 2025: Initial version of the integration.
- Francesco Massa
- Marco Cristofanilli (marco.cATseeweb.it)

"""

from __future__ import annotations

import typing
from typing import Any, Dict, Iterator, List, Optional, Tuple, Union

from sky import catalog
from sky import clouds
from sky import sky_logging
from sky.adaptors import seeweb as seeweb_adaptor
from sky.provision import seeweb as seeweb_provision
from sky.utils import registry
from sky.utils import resources_utils
from sky.utils import ux_utils

if typing.TYPE_CHECKING:
    from sky import resources as resources_lib
    from sky.utils import status_lib
    from sky.utils import volume as volume_lib

# ---------- key file path -----------------
_SEEWEB_KEY_FILE = '~/.seeweb_cloud/seeweb_keys'

logger = sky_logging.init_logger(__name__)
# (content: ini-like)
#   api_key = <TOKEN>


@registry.CLOUD_REGISTRY.register
class Seeweb(clouds.Cloud):
    """Seeweb GPU Cloud."""

    _REPR = 'Seeweb'
    # Define unsupported features to provide clear error messages
    # This helps users understand what Seeweb can and cannot do
    _CLOUD_UNSUPPORTED_FEATURES = {
        clouds.CloudImplementationFeatures.MULTI_NODE:
            ('Multi-node not supported. '
             'Seeweb does not support multi-node clusters.'),
        clouds.CloudImplementationFeatures.CUSTOM_DISK_TIER:
            ('Custom disk tiers not supported. '
             'Seeweb does not support custom disk tiers.'),
        clouds.CloudImplementationFeatures.STORAGE_MOUNTING:
            ('Storage mounting not supported. '
             'Seeweb does not support storage mounting.'),
        clouds.CloudImplementationFeatures.HIGH_AVAILABILITY_CONTROLLERS:
            ('High availability controllers not supported. '
             'Seeweb does not support high availability controllers.'),
        clouds.CloudImplementationFeatures.SPOT_INSTANCE:
            ('Spot instances not supported. '
             'Seeweb does not support spot instances.'),
        clouds.CloudImplementationFeatures.CLONE_DISK_FROM_CLUSTER:
            ('Disk cloning not supported. '
             'Seeweb does not support disk cloning.'),
        clouds.CloudImplementationFeatures.IMAGE_ID:
            ('Custom image IDs not supported. '
             'Seeweb does not support custom image IDs.'),
        clouds.CloudImplementationFeatures.CUSTOM_NETWORK_TIER:
            ('Custom network tiers not supported. '
             'Seeweb does not support custom network tiers.'),
        clouds.CloudImplementationFeatures.HOST_CONTROLLERS:
            ('Host controllers not supported. '
             'Seeweb does not support host controllers.'),
        clouds.CloudImplementationFeatures.CUSTOM_MULTI_NETWORK:
            ('Custom multi-network not supported. '
             'Seeweb does not support custom multi-network.'),
    }
    _MAX_CLUSTER_NAME_LEN_LIMIT = 120
    _regions: List[clouds.Region] = []

    PROVISIONER_VERSION = clouds.ProvisionerVersion.SKYPILOT
    STATUS_VERSION = clouds.StatusVersion.SKYPILOT

    # Enable port support with updatable version
    OPEN_PORTS_VERSION = clouds.OpenPortsVersion.UPDATABLE

    @classmethod
    def _unsupported_features_for_resources(
        cls,
        resources: 'resources_lib.Resources',
        region: Optional[str] = None,
    ) -> Dict[clouds.CloudImplementationFeatures, str]:
        return cls._CLOUD_UNSUPPORTED_FEATURES

    @classmethod
    def max_cluster_name_length(cls) -> Optional[int]:
        return cls._MAX_CLUSTER_NAME_LEN_LIMIT

    @classmethod
    def regions(cls) -> List['clouds.Region']:
        """Return available regions for Seeweb."""
        # Get regions from the catalog system
        # This reads from the CSV files generated by fetch_seeweb.py
        regions = catalog.regions(clouds='seeweb')
        return regions

    @classmethod
    def regions_with_offering(
        cls,
        instance_type: str,
        accelerators: Optional[Dict[str, int]],
        use_spot: bool,
        region: Optional[str],
        zone: Optional[str],
        resources: Optional['resources_lib.Resources'] = None,
    ) -> List[clouds.Region]:
        assert zone is None, 'Seeweb does not support zones.'
        del zone
        if use_spot:
            return []

        # Get regions from catalog based on instance type
        # This will read the CSV and return only regions
        # where the instance type exists
        regions = catalog.get_region_zones_for_instance_type(
            instance_type, use_spot, 'seeweb')

        if region is not None:
            regions = [r for r in regions if r.name == region]

        return regions

    @classmethod
    def zones_provision_loop(
        cls,
        *,
        region: str,
        num_nodes: int,
        instance_type: str,
        accelerators: Optional[Dict[str, int]] = None,
        use_spot: bool = False,
    ) -> Iterator[None]:
        del num_nodes
        regions = cls.regions_with_offering(instance_type,
                                            accelerators,
                                            use_spot,
                                            region=region,
                                            zone=None)
        for r in regions:
            assert r.zones is None, r
            yield r.zones

    @classmethod
    def get_zone_shell_cmd(cls) -> Optional[str]:
        """Seeweb doesn't support zones."""
        return None

    def instance_type_to_hourly_cost(
        self,
        instance_type: str,
        use_spot: bool,
        region: Optional[str],
        zone: Optional[str],
    ) -> float:
        cost = catalog.get_hourly_cost(instance_type,
                                       use_spot=use_spot,
                                       region=region,
                                       zone=zone,
                                       clouds='seeweb')
        return cost

    def accelerators_to_hourly_cost(
        self,
        accelerators: Dict[str, int],
        use_spot: bool,
        region: Optional[str],
        zone: Optional[str],
    ) -> float:

        return 0.0

    def get_egress_cost(self, num_gigabytes: float):
        return 0.0

    def make_deploy_resources_variables(
        self,
        resources: 'resources_lib.Resources',
        cluster_name: resources_utils.ClusterName,
        region: 'clouds.Region',
        zones: Optional[List['clouds.Zone']],
        num_nodes: int,
        dryrun: bool = False,
        volume_mounts: Optional[List['volume_lib.VolumeMount']] = None,
    ) -> Dict[str, Any]:
        """Create deployment variables for Seeweb."""

        # Note: Spot instances and multi-node are automatically handled by
        # the framework via _CLOUD_UNSUPPORTED_FEATURES

        resources = resources.assert_launchable()

        acc_dict = self.get_accelerators_from_instance_type(
            resources.instance_type)
        docker_image = resources.extract_docker_image()
        docker_run_options: List[str] = []
        if docker_image is not None:
            if acc_dict:
                docker_run_options.append('--gpus all')
            logger.info(
                'Launching Seeweb cluster with docker image %s. Ensure the '
                'image is Debian-based and allows passwordless sudo.',
                docker_image)

        # Standard custom_resources string for Ray
        custom_resources = resources_utils.make_ray_custom_resources_str(
            acc_dict)

        # Seeweb-specific GPU configuration for the provisioner
        # This tells the provisioner how to configure GPU resources
        seeweb_gpu_config = None
        if resources.accelerators:
            # If the instance has accelerators, prepare GPU configuration
            accelerator_name = list(resources.accelerators.keys())[0]
            accelerator_count = resources.accelerators[accelerator_name]
            seeweb_gpu_config = {
                'gpu': accelerator_count,
                'gpu_label': accelerator_name,
            }

        # Seeweb uses pre-configured images based on instance type
        # Determine image based on whether the instance type name contains "GPU"
        if resources.instance_type and 'GPU' in resources.instance_type.upper():
            # GPU instance - use image with NVIDIA drivers
            if resources.instance_type in ['ECS1GPU10', 'ECS2GPU10']:
                # H200 GPU instance - use UEFI image with NVIDIA drivers
                image_id = 'ubuntu-2204-uefi-nvidia-driver'
            else:
                # Other GPU instance - use standard image with NVIDIA drivers
                image_id = 'ubuntu-2204-nvidia-driver'
        else:
            # CPU-only instance - use standard Ubuntu image
            image_id = 'ubuntu-2204'

        result = {
            'instance_type': resources.instance_type,
            'region': region.name,
            'cluster_name': cluster_name,
            'custom_resources': custom_resources,
            'seeweb_gpu_config': seeweb_gpu_config,
            'image_id': image_id,
        }
        if docker_run_options:
            result['docker_run_options'] = docker_run_options
        return result

    @classmethod
    def get_vcpus_mem_from_instance_type(
            cls, instance_type: str) -> Tuple[Optional[float], Optional[float]]:
        result = catalog.get_vcpus_mem_from_instance_type(instance_type,
                                                          clouds='seeweb')
        return result

    @classmethod
    def get_accelerators_from_instance_type(
        cls,
        instance_type: str,
    ) -> Optional[Dict[str, Union[int, float]]]:
        result = catalog.get_accelerators_from_instance_type(instance_type,
                                                             clouds='seeweb')
        return result

    @classmethod
    def get_default_instance_type(
        cls,
        cpus: Optional[str] = None,
        memory: Optional[str] = None,
        disk_tier: Optional[resources_utils.DiskTier] = None,
        region: Optional[str] = None,
        zone: Optional[str] = None,
    ) -> Optional[str]:
        result = catalog.get_default_instance_type(cpus=cpus,
                                                   memory=memory,
                                                   disk_tier=disk_tier,
                                                   clouds='seeweb')
        return result

    def _get_feasible_launchable_resources(
        self, resources: 'resources_lib.Resources'
    ) -> 'resources_utils.FeasibleResources':
        """Get feasible resources for Seeweb."""
        if resources.use_spot:
            return resources_utils.FeasibleResources(
                [], [], 'Spot instances not supported on Seeweb')

        if resources.accelerators and len(resources.accelerators) > 1:
            return resources_utils.FeasibleResources(
                [], [], 'Multiple accelerator types not supported on Seeweb')

        # If no instance_type is specified, try to get a default one
        if not resources.instance_type:
            # If accelerators are specified, try to find instance
            # type forthat accelerator
            if resources.accelerators:
                # Get the first accelerator
                # (we already checked there's only one)
                acc_name, acc_count = list(resources.accelerators.items())[0]

                # Use catalog to find instance type for this accelerator
                # This leverages the catalog system to find suitable instances
                (
                    instance_types,
                    fuzzy_candidates,
                ) = catalog.get_instance_type_for_accelerator(
                    acc_name=acc_name,
                    acc_count=acc_count,
                    cpus=resources.cpus,
                    memory=resources.memory,
                    use_spot=resources.use_spot,
                    region=resources.region,
                    zone=resources.zone,
                    clouds='seeweb',
                )

                if instance_types and len(instance_types) > 0:
                    # Use the first (cheapest) instance type
                    selected_instance_type = instance_types[0]
                    resources = resources.copy(
                        instance_type=selected_instance_type)
                else:
                    return resources_utils.FeasibleResources(
                        [],
                        fuzzy_candidates,
                        f'No instance type found for accelerator'
                        f'{acc_name}:{acc_count} on Seeweb',
                    )
            else:
                # No accelerators specified, use default instance type
                default_instance_type = self.get_default_instance_type(
                    cpus=resources.cpus,
                    memory=resources.memory,
                    region=resources.region,
                    zone=resources.zone,
                )

                if default_instance_type:
                    # Create new resources with the default instance type
                    resources = resources.copy(
                        instance_type=default_instance_type)
                else:
                    return resources_utils.FeasibleResources(
                        [],
                        [],
                        f'No suitable instance type found for'
                        f'cpus={resources.cpus}, memory={resources.memory}',
                    )

        # Check if instance type exists
        if resources.instance_type:
            exists = catalog.instance_type_exists(resources.instance_type,
                                                  clouds='seeweb')
            if not exists:
                return resources_utils.FeasibleResources(
                    [],
                    [],
                    f'Instance type {resources.instance_type}'
                    f' not available on Seeweb',
                )

        # Set the cloud if not already set
        if not resources.cloud:
            resources = resources.copy(cloud=self)

        # Return the resources as feasible
        return resources_utils.FeasibleResources([resources], [], None)

    @classmethod
    def _check_compute_credentials(cls) -> Tuple[bool, Optional[str]]:
        """Check Seeweb compute credentials."""
        try:
            result = seeweb_adaptor.check_compute_credentials()
            return result, None
        except Exception as e:  # pylint: disable=broad-except
            return False, str(e)

    @classmethod
    def _check_storage_credentials(cls) -> Tuple[bool, Optional[str]]:
        """Check Seeweb storage credentials."""
        try:
            result = seeweb_adaptor.check_storage_credentials()
            return result, None
        except Exception as e:  # pylint: disable=broad-except
            return False, str(e)

    @classmethod
    def get_user_identities(cls) -> Optional[List[List[str]]]:
        # Seeweb doesn't have user identity concept
        return None

    @classmethod
    def query_status(
        cls,
        name: str,
        tag_filters: Dict[str, str],
        region: Optional[str],
        zone: Optional[str],
        **kwargs,
    ) -> List['status_lib.ClusterStatus']:
        """Query the status of Seeweb cluster instances."""
        cluster_name_on_cloud = name

        result = seeweb_provision.instance.query_instances(
            cluster_name=name,
            cluster_name_on_cloud=cluster_name_on_cloud,
            provider_config={},
            non_terminated_only=True)
        # Convert Dict[str, Tuple[Optional[ClusterStatus],
        # Optional[str]]] to List[ClusterStatus]
        return [status for status, _ in result.values() if status is not None]

    def get_credential_file_mounts(self) -> Dict[str, str]:
        """Returns the credential files to mount."""
        # Mount the Seeweb API key file to the remote instance
        # This allows the provisioner to authenticate with Seeweb API
        result = {
            _SEEWEB_KEY_FILE: _SEEWEB_KEY_FILE,
        }
        return result

    def instance_type_exists(self, instance_type: str) -> bool:
        """Returns whether the instance type exists for Seeweb."""
        result = catalog.instance_type_exists(instance_type, clouds='seeweb')
        return result

    @classmethod
    def get_image_size(cls, image_id: str, region: Optional[str]) -> float:
        """Seeweb doesn't support custom images."""
        del image_id, region
        with ux_utils.print_exception_no_traceback():
            raise ValueError(f'Custom images are not supported on {cls._REPR}. '
                             'Seeweb clusters use pre-configured images only.')

    # Image-related methods (not supported)
    @classmethod
    def create_image_from_cluster(
        cls,
        cluster_name: resources_utils.ClusterName,
        region: Optional[str],
        zone: Optional[str],
    ) -> str:
        del cluster_name, region, zone  # unused
        with ux_utils.print_exception_no_traceback():
            raise ValueError(
                f'Creating images from clusters is not supported on'
                f' {cls._REPR}. Seeweb does not support custom'
                f' image creation.')

    @classmethod
    def maybe_move_image(
        cls,
        image_id: str,
        source_region: str,
        target_region: str,
        source_zone: Optional[str],
        target_zone: Optional[str],
    ) -> str:
        del image_id, source_region, target_region, source_zone, target_zone
        with ux_utils.print_exception_no_traceback():
            raise ValueError(
                f'Moving images between regions is not supported on'
                f' {cls._REPR}. '
                'Seeweb does not support custom images.')

    @classmethod
    def delete_image(cls, image_id: str, region: Optional[str]) -> None:
        del image_id, region
        with ux_utils.print_exception_no_traceback():
            raise ValueError(
                f'Deleting images is not supported on {cls._REPR}. '
                'Seeweb does not support custom image management.')
