import os
from datetime import datetime
from enum import Enum
from typing import Dict, List

from ci.ray_ci.configs import DEFAULT_ARCHITECTURE, DEFAULT_PYTHON_TAG_VERSION
from ci.ray_ci.linux_container import LinuxContainer
from ci.ray_ci.supported_images import (
    get_architectures,
    get_default,
    get_platforms,
    get_python_versions,
)

PLATFORMS_RAY = get_platforms("ray")
PLATFORMS_RAY_ML = get_platforms("ray-ml")
PLATFORMS_RAY_LLM = get_platforms("ray-llm")
GPU_PLATFORM = get_default("ray", "gpu_platform")

PYTHON_VERSIONS_RAY = get_python_versions("ray")
PYTHON_VERSIONS_RAY_ML = get_python_versions("ray-ml")
PYTHON_VERSIONS_RAY_LLM = get_python_versions("ray-llm")
ARCHITECTURES_RAY = get_architectures("ray")
ARCHITECTURES_RAY_ML = get_architectures("ray-ml")
ARCHITECTURES_RAY_LLM = get_architectures("ray-llm")


class RayType(str, Enum):
    RAY = "ray"
    RAY_EXTRA = "ray-extra"
    RAY_ML = "ray-ml"
    RAY_ML_EXTRA = "ray-ml-extra"
    RAY_LLM = "ray-llm"
    RAY_LLM_EXTRA = "ray-llm-extra"


RAY_REPO_MAP: Dict[str, str] = {
    RayType.RAY.value: RayType.RAY.value,
    RayType.RAY_ML.value: RayType.RAY_ML.value,
    RayType.RAY_LLM.value: RayType.RAY_LLM.value,
    RayType.RAY_EXTRA.value: RayType.RAY.value,
    RayType.RAY_ML_EXTRA.value: RayType.RAY_ML.value,
    RayType.RAY_LLM_EXTRA.value: RayType.RAY_LLM.value,
}


class DockerContainer(LinuxContainer):
    """
    Container for building and publishing ray docker images
    """

    def __init__(
        self,
        python_version: str,
        platform: str,
        image_type: str,
        architecture: str = DEFAULT_ARCHITECTURE,
        canonical_tag: str = None,
        upload: bool = False,
    ) -> None:
        assert "RAYCI_CHECKOUT_DIR" in os.environ, "RAYCI_CHECKOUT_DIR not set"
        rayci_checkout_dir = os.environ["RAYCI_CHECKOUT_DIR"]

        super().__init__(
            "forge" if architecture == "x86_64" else "forge-aarch64",
            python_version=python_version,
            volumes=[
                f"{rayci_checkout_dir}:/rayci",
                "/var/run/docker.sock:/var/run/docker.sock",
            ],
        )

        if image_type in [RayType.RAY_ML, RayType.RAY_ML_EXTRA]:
            assert python_version in PYTHON_VERSIONS_RAY_ML
            assert platform in PLATFORMS_RAY_ML
            assert architecture in ARCHITECTURES_RAY_ML
        elif image_type in [RayType.RAY_LLM, RayType.RAY_LLM_EXTRA]:
            assert python_version in PYTHON_VERSIONS_RAY_LLM
            assert platform in PLATFORMS_RAY_LLM
            assert architecture in ARCHITECTURES_RAY_LLM
        else:
            # ray or ray-extra
            assert python_version in PYTHON_VERSIONS_RAY
            assert platform in PLATFORMS_RAY
            assert architecture in ARCHITECTURES_RAY

        self.platform = platform
        self.image_type = image_type
        self.architecture = architecture
        self.canonical_tag = canonical_tag
        self.upload = upload

    def _get_image_version_tags(self, external: bool) -> List[str]:
        """
        Get version tags.

        Args:
            external: If True, return the external image tags. If False, return the
                internal image tags.
        """
        branch = os.environ.get("BUILDKITE_BRANCH", "")
        sha_tag = os.environ["BUILDKITE_COMMIT"][:6]
        rayci_build_id = os.environ["RAYCI_BUILD_ID"]
        pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false")
        formatted_date = datetime.now().strftime("%y%m%d")

        if branch == "master":
            if external and os.environ.get("RAYCI_SCHEDULE") == "nightly":
                return [f"nightly.{formatted_date}.{sha_tag}", "nightly"]
            return [sha_tag, rayci_build_id]

        if branch and branch.startswith("releases/"):
            release_name = branch[len("releases/") :]
            release_tag = f"{release_name}.{sha_tag}"
            if external:
                # Avoid saving build ID ones when saving it on public registries.
                return [release_tag]
            return [release_tag, rayci_build_id]

        if pr != "false":
            return [f"pr-{pr}.{sha_tag}", rayci_build_id]

        return [sha_tag, rayci_build_id]

    def _get_canonical_tag(self) -> str:
        # The canonical tag is the first tag in the list of tags. The list of tag is
        # never empty because the image is always tagged with at least the sha tag.
        #
        # The canonical tag is the most complete tag with no abbreviation,
        # e.g. sha-pyversion-platform
        return self.canonical_tag if self.canonical_tag else self._get_image_tags()[0]

    def _get_python_version_tag(self) -> str:
        return f"-py{self.python_version.replace('.', '')}"  # 3.x -> py3x

    def _get_platform_tag(self) -> str:
        if self.platform == "cpu":
            return "-cpu"
        if self.platform == "tpu":
            return "-tpu"
        versions = self.platform.split(".")
        return f"-{versions[0]}{versions[1]}"  # cu11.8.0 -> cu118

    def _get_image_tags(self, external: bool = False) -> List[str]:
        """
        Get image tags & alias for the container image.

        An image tag is composed by ray version tag, python version and platform.
        See https://docs.ray.io/en/latest/ray-overview/installation.html for
        more information on the image tags.

        Args:
            external: If True, return the external image tags. If False, return the
                internal image tags.
        """

        versions = self._get_image_version_tags(external)

        platforms = [self._get_platform_tag()]
        if self.platform == "cpu" and self.image_type in [
            RayType.RAY,
            RayType.RAY_EXTRA,
        ]:
            # no tag is alias to cpu for ray image
            platforms.append("")
        elif self.platform == GPU_PLATFORM:
            # gpu is alias to cu118 for ray image
            platforms.append("-gpu")
            if self.image_type in [RayType.RAY_ML, RayType.RAY_ML_EXTRA]:
                # no tag is alias to gpu for ray-ml image
                platforms.append("")

        py_versions = [self._get_python_version_tag()]
        if self.python_version == DEFAULT_PYTHON_TAG_VERSION:
            py_versions.append("")

        variation = ""
        if self.image_type in [
            RayType.RAY_EXTRA,
            RayType.RAY_ML_EXTRA,
            RayType.RAY_LLM_EXTRA,
        ]:
            variation = "-extra"

        tags = []
        for version in versions:
            for platform in platforms:
                for py_version in py_versions:
                    tag = f"{version}{variation}{py_version}{platform}"
                    if self.architecture != DEFAULT_ARCHITECTURE:
                        tag += f"-{self.architecture}"
                    tags.append(tag)
        return tags