from dataclasses import dataclass
from enum import Enum
from typing import Any, ClassVar, Dict, Optional, Tuple

# Backend slugs that name a machine already listed below. Baremetal SKUs carry the GPU
# memory in the slug, and H200/B200 shipped with the count spelled "x8" before the slugs
# were unified. Mirrors LegacyMultiCloudSlugs in grid's pkg/cluster/accelerators.go.
_SLUG_ALIASES: Dict[str, str] = {
    "lit-h200-141gb-1": "lit-h200-1",
    "lit-h200-141gb-2": "lit-h200-2",
    "lit-h200-141gb-4": "lit-h200-4",
    "lit-h200-141gb-8": "lit-h200-8",
    "lit-h200x-1": "lit-h200-1",
    "lit-h200x-2": "lit-h200-2",
    "lit-h200x-4": "lit-h200-4",
    "lit-h200x-8": "lit-h200-8",
    "lit-b200-180gb-8": "lit-b200-8",
    "lit-b200x-1": "lit-b200-1",
    "lit-b200x-8": "lit-b200-8",
}


class CloudProvider(Enum):
    """Cloud provider for running Studios and Jobs on the Lightning AI platform."""

    AWS = "AWS"
    GCP = "GCP"
    LAMBDA_LABS = "LAMBDA_LABS"
    DGX = "DGX"
    NEBIUS = "NEBIUS"
    LIGHTNING = "LIGHTNING"
    LIGHTNING_AGGREGATE = "LIGHTNING_AGGREGATE"

    def __str__(self) -> str:
        """Converts the CloudProvider to a str."""
        return self.value

    @classmethod
    def from_str(cls, provider: str) -> "CloudProvider":
        """Convert a string to a CloudProvider enum member.

        Args:
            provider: The string name of the cloud provider (case-insensitive).

        Returns:
            CloudProvider: The matching enum member.

        Raises:
            ValueError: If the string does not match any known provider.
        """
        for cp in cls:
            if cp.value.lower() == provider.lower():
                return cp
        raise ValueError(f"Unknown CloudProvider: {provider}")


@dataclass(frozen=True)
class Machine:
    """Hardware configuration for running Studios and Jobs on the Lightning AI platform.

    Use the pre-defined class attributes (e.g. ``Machine.A10G``, ``Machine.H100``) rather than
    constructing instances directly. Pass a ``Machine`` to :meth:`~lightning_sdk.Studio.start`,
    :meth:`~lightning_sdk.Job.run`, or :meth:`~lightning_sdk.MMT.run` to select the hardware.
    """

    # supported CPU variations
    CPU_X_2: ClassVar["Machine"]
    CPU_X_4: ClassVar["Machine"]
    CPU_X_8: ClassVar["Machine"]
    CPU_X_16: ClassVar["Machine"]
    # default CPU machines
    CPU_SMALL: ClassVar["Machine"]
    CPU: ClassVar["Machine"]
    # supported data-prep variations (big disk)
    DATA_PREP: ClassVar["Machine"]
    DATA_PREP_MAX: ClassVar["Machine"]
    DATA_PREP_ULTRA: ClassVar["Machine"]

    # supported GPU types
    # supported T4 variations
    T4_SMALL: ClassVar["Machine"]
    T4: ClassVar["Machine"]
    T4_X_2: ClassVar["Machine"]
    T4_X_4: ClassVar["Machine"]
    T4_X_8: ClassVar["Machine"]
    # supported L4 variations
    L4: ClassVar["Machine"]
    L4_X_2: ClassVar["Machine"]
    L4_X_4: ClassVar["Machine"]
    L4_X_8: ClassVar["Machine"]
    # supported L40S variations
    L40S: ClassVar["Machine"]
    L40S_X_2: ClassVar["Machine"]
    L40S_X_4: ClassVar["Machine"]
    L40S_X_8: ClassVar["Machine"]
    # supported RTXP 6000 variations
    RTXP_6000: ClassVar["Machine"]
    RTXP_6000_X_2: ClassVar["Machine"]
    RTXP_6000_X_4: ClassVar["Machine"]
    RTXP_6000_X_8: ClassVar["Machine"]
    # supported A100 variations
    # defaults, can be either A100 type depending on cloud provider availability
    A100: ClassVar["Machine"]
    A100_X_2: ClassVar["Machine"]
    A100_X_4: ClassVar["Machine"]
    A100_X_8: ClassVar["Machine"]
    # A100 40GB versions
    A100_40GB: ClassVar["Machine"]
    A100_40GB_X_2: ClassVar["Machine"]
    A100_40GB_X_4: ClassVar["Machine"]
    A100_40GB_X_8: ClassVar["Machine"]
    # A100 80GB versions
    A100_80GB: ClassVar["Machine"]
    A100_80GB_X_2: ClassVar["Machine"]
    A100_80GB_X_4: ClassVar["Machine"]
    A100_80GB_X_8: ClassVar["Machine"]

    H100: ClassVar["Machine"]
    H100_X_2: ClassVar["Machine"]
    H100_X_4: ClassVar["Machine"]
    H100_X_8: ClassVar["Machine"]

    H200: ClassVar["Machine"]
    H200_X_2: ClassVar["Machine"]
    H200_X_4: ClassVar["Machine"]
    H200_X_8: ClassVar["Machine"]

    B200: ClassVar["Machine"]
    B200_X_8: ClassVar["Machine"]

    # Specialized Machines

    name: str
    slug: str
    instance_type: Optional[str] = None
    family: Optional[str] = None
    accelerator_count: Optional[int] = None
    cost: Optional[float] = None
    interruptible_cost: Optional[float] = None
    provider: Optional[str] = None
    wait_time: Optional[float] = None
    interruptible_wait_time: Optional[float] = None
    _include_in_cli: bool = True

    def __str__(self) -> str:
        """String representation of the Machine."""
        return str(self.name) if self.name else (self.slug if self.slug else str(self.instance_type))

    def __eq__(self, other: object) -> bool:
        """Machines are equal if the instance type is equal."""
        if isinstance(other, Machine):
            return (
                # equality based on raw instance type (provider specific)
                (self.instance_type and self.instance_type == other.instance_type)
                # equality based on slug (provider agnostic)
                or self.slug == other.slug
                # equality based on machine specs (e.g. A100_80GB_X_8 == A100_X_8)
                or (self.family == other.family and self.accelerator_count == other.accelerator_count)
            )
        return False

    def is_cpu(self) -> bool:
        """Whether the machine is a CPU."""
        return self.family in ("CPU", "DATA_PREP")

    @classmethod
    def _predefined_from_str(cls, machine: Optional[str], *additional_machine_ids: Any) -> Optional["Machine"]:
        """Return the predefined machine matching any of the identifiers, or ``None`` if unknown.

        Args:
            machine: Primary machine identifier (name, slug, or instance type).
            *additional_machine_ids: Extra identifiers to try when the primary does not match.

        Returns:
            Machine: The matching predefined machine, or ``None`` when no attribute matches.
        """
        possible_values: Tuple["Machine", ...] = tuple(
            [machine for machine in cls.__dict__.values() if isinstance(machine, cls)]
        )
        for m in possible_values:
            for machine_id in [machine, *additional_machine_ids]:
                if machine_id and _SLUG_ALIASES.get(machine_id, machine_id) in (
                    getattr(m, "name", None),
                    getattr(m, "instance_type", None),
                    getattr(m, "slug", None),
                ):
                    return m
        return None

    @classmethod
    def from_str(cls, machine: str, *additional_machine_ids: Any) -> "Machine":
        """Resolve a machine from one or more string identifiers.

        Tries to match against known machine names, slugs, and instance types in order.
        Falls back to constructing an ad-hoc ``Machine`` if no match is found.

        Args:
            machine: Primary machine identifier (name, slug, or instance type).
            *additional_machine_ids: Extra identifiers to try when the primary does not match.

        Returns:
            Machine: The resolved or constructed machine instance.
        """
        predefined = cls._predefined_from_str(machine, *additional_machine_ids)
        if predefined is not None:
            return predefined

        if additional_machine_ids:
            return cls(machine, *additional_machine_ids)
        return cls(machine, machine, machine)

    @classmethod
    def _from_accelerator(cls, accelerator: Any) -> "Machine":
        """Construct a Machine from a raw accelerator object returned by the cluster API.

        Args:
            accelerator: An accelerator descriptor with ``accelerator_type``, ``resources``,
                ``slug_multi_cloud``, ``slug``, ``instance_id``, ``secondary_instance_id``,
                and ``family`` attributes.

        Returns:
            Machine: The resolved or ad-hoc ``Machine`` instance.
        """
        if accelerator.accelerator_type == "GPU":
            accelerator_resources_count = accelerator.resources.gpu
        else:
            accelerator_resources_count = accelerator.resources.cpu

        return Machine.from_str(
            accelerator.slug_multi_cloud,
            accelerator.slug,
            accelerator.instance_id,
            accelerator.secondary_instance_id,
            f"lit-{accelerator.family.lower()}-{accelerator_resources_count}",
        )


# CPU machines
# default CPU machines
Machine.CPU_SMALL = Machine(name="CPU_SMALL", slug="cpu-2", family="CPU", accelerator_count=2)
Machine.CPU = Machine(name="CPU", slug="cpu-4", family="CPU", accelerator_count=4)
# available CPU variations
Machine.CPU_X_2 = Machine(name="CPU_X_2", slug="cpu-2", family="CPU", accelerator_count=2)
Machine.CPU_X_4 = Machine(name="CPU_X_4", slug="cpu-4", family="CPU", accelerator_count=4)
Machine.CPU_X_8 = Machine(name="CPU_X_8", slug="cpu-8", family="CPU", accelerator_count=8)
Machine.CPU_X_16 = Machine(name="CPU_X_16", slug="cpu-16", family="CPU", accelerator_count=16)
# available data-prep (big disk) machines
Machine.DATA_PREP = Machine(name="DATA_PREP", slug="data-prep-mid", family="DATA-PREP", accelerator_count=32)
Machine.DATA_PREP_MAX = Machine(
    name="DATA_PREP_MAX", slug="data-prep-max-large", family="DATA-PREP", accelerator_count=64
)
Machine.DATA_PREP_ULTRA = Machine(
    name="DATA_PREP_ULTRA", slug="data-prep-ultra-extra-large", family="DATA-PREP", accelerator_count=96
)

# GPU machines
# available T4 machines
Machine.T4_SMALL = Machine(name="T4_SMALL", slug="lit-t4-1-small", family="T4", accelerator_count=1)
Machine.T4 = Machine(name="T4", slug="lit-t4-1", family="T4", accelerator_count=1)
Machine.T4_X_2 = Machine(name="T4_X_2", slug="lit-t4-2", family="T4", accelerator_count=2)
Machine.T4_X_4 = Machine(name="T4_X_4", slug="lit-t4-4", family="T4", accelerator_count=4)
Machine.T4_X_8 = Machine(name="T4_X_8", slug="lit-t4-8", family="T4", accelerator_count=8)
# available L4 machines
Machine.L4 = Machine(name="L4", slug="lit-l4-1", family="L4", accelerator_count=1)
Machine.L4_X_2 = Machine(name="L4_X_2", slug="lit-l4-2", family="L4", accelerator_count=2)
Machine.L4_X_4 = Machine(name="L4_X_4", slug="lit-l4-4", family="L4", accelerator_count=4)
Machine.L4_X_8 = Machine(name="L4_X_8", slug="lit-l4-8", family="L4", accelerator_count=8)
# available L40S machines
Machine.L40S = Machine(name="L40S", slug="lit-l40s-1", family="L40S", accelerator_count=1)
Machine.L40S_X_2 = Machine(name="L40S_X_2", slug="lit-l40s-2", family="L40S", accelerator_count=2)
Machine.L40S_X_4 = Machine(name="L40S_X_4", slug="lit-l40s-4", family="L40S", accelerator_count=4)
Machine.L40S_X_8 = Machine(name="L40S_X_8", slug="lit-l40s-8", family="L40S", accelerator_count=8)
# available RTXP 6000 machines
Machine.RTXP_6000 = Machine(name="RTXP_6000", slug="lit-rtx-6000-pro-1", family="RTX PRO", accelerator_count=1)
Machine.RTXP_6000_X_2 = Machine(name="RTXP_6000_X_2", slug="lit-rtx-6000-pro-2", family="RTX PRO", accelerator_count=2)
Machine.RTXP_6000_X_4 = Machine(name="RTXP_6000_X_4", slug="lit-rtx-6000-pro-4", family="RTX PRO", accelerator_count=4)
Machine.RTXP_6000_X_8 = Machine(name="RTXP_6000_X_8", slug="lit-rtx-6000-pro-8", family="RTX PRO", accelerator_count=8)
# available A100 Machines
Machine.A100 = Machine(name="A100", slug="lit-a100-1", family="A100", accelerator_count=1)
Machine.A100_X_2 = Machine(name="A100_X_2", slug="lit-a100-2", family="A100", accelerator_count=2)
Machine.A100_X_4 = Machine(name="A100_X_4", slug="lit-a100-4", family="A100", accelerator_count=4)
Machine.A100_X_8 = Machine(name="A100_X_8", slug="lit-a100-8", family="A100", accelerator_count=8)
# don't include variants in cli, only default types that can match for all variants
Machine.A100_40GB = Machine(
    name="A100_40GB", slug="lit-a100-40gb-1", family="A100", accelerator_count=1, _include_in_cli=False
)
Machine.A100_40GB_X_2 = Machine(
    name="A100_40GB_X_2", slug="lit-a100-40gb-2", family="A100", accelerator_count=2, _include_in_cli=False
)
Machine.A100_40GB_X_4 = Machine(
    name="A100_40GB_X_4", slug="lit-a100-40gb-4", family="A100", accelerator_count=4, _include_in_cli=False
)
Machine.A100_40GB_X_8 = Machine(
    name="A100_40GB_X_8", slug="lit-a100-40gb-8", family="A100", accelerator_count=8, _include_in_cli=False
)
Machine.A100_80GB = Machine(
    name="A100_80GB", slug="lit-a100-80gb-1", family="A100", accelerator_count=1, _include_in_cli=False
)
Machine.A100_80GB_X_2 = Machine(
    name="A100_80GB_X_2", slug="lit-a100-80gb-2", family="A100", accelerator_count=2, _include_in_cli=False
)
Machine.A100_80GB_X_4 = Machine(
    name="A100_80GB_X_4", slug="lit-a100-80gb-4", family="A100", accelerator_count=4, _include_in_cli=False
)
Machine.A100_80GB_X_8 = Machine(
    name="A100_80GB_X_8", slug="lit-a100-80gb-8", family="A100", accelerator_count=8, _include_in_cli=False
)
# available H100 machines
Machine.H100 = Machine(name="H100", slug="lit-h100-1", family="H100", accelerator_count=1)
Machine.H100_X_2 = Machine(name="H100_X_2", slug="lit-h100-2", family="H100", accelerator_count=2)
Machine.H100_X_4 = Machine(name="H100_X_4", slug="lit-h100-4", family="H100", accelerator_count=4)
Machine.H100_X_8 = Machine(name="H100_X_8", slug="lit-h100-8", family="H100", accelerator_count=8)
# available H200 machines
Machine.H200 = Machine(name="H200", slug="lit-h200-1", family="H200", accelerator_count=1)
Machine.H200_X_2 = Machine(name="H200_X_2", slug="lit-h200-2", family="H200", accelerator_count=2)
Machine.H200_X_4 = Machine(name="H200_X_4", slug="lit-h200-4", family="H200", accelerator_count=4)
Machine.H200_X_8 = Machine(name="H200_X_8", slug="lit-h200-8", family="H200", accelerator_count=8)
# available B200 machines
Machine.B200 = Machine(name="B200", slug="lit-b200-1", family="B200", accelerator_count=1)
Machine.B200_X_8 = Machine(name="B200_X_8", slug="lit-b200-8", family="B200", accelerator_count=8)


DEFAULT_MACHINE = Machine.CPU.name