from typing import List, Optional, Dict

import boto3
import hashlib
import os
import subprocess
import sys
import time

from ray_release.config import RELEASE_PACKAGE_DIR
from ray_release.configs.global_config import get_global_config
from ray_release.logger import logger
from ray_release.test import (
    Test,
    DATAPLANE_ECR_REPO,
    DATAPLANE_ECR_ML_REPO,
)

DATAPLANE_S3_BUCKET = "ray-release-automation-results"
DATAPLANE_FILENAME = "dataplane_20240613.tar.gz"
DATAPLANE_DIGEST = "b7c68dd3cf9ef05b2b1518a32729fc036f06ae83bae9b9ecd241e33b37868944"
BASE_IMAGE_WAIT_TIMEOUT = 7200
BASE_IMAGE_WAIT_DURATION = 30
RELEASE_BYOD_DIR = os.path.join(RELEASE_PACKAGE_DIR, "ray_release/byod")
REQUIREMENTS_BYOD = "requirements_byod"
REQUIREMENTS_ML_BYOD = "requirements_ml_byod"


def build_champagne_image(
    ray_version: str,
    python_version: str,
    image_type: str,
) -> str:
    """
    Builds the Anyscale champagne image.
    """
    _download_dataplane_build_file()
    env = os.environ.copy()
    env["DOCKER_BUILDKIT"] = "1"
    if image_type == "cpu":
        ray_project = "ray"
        anyscale_repo = DATAPLANE_ECR_REPO
    else:
        ray_project = "ray-ml"
        anyscale_repo = DATAPLANE_ECR_ML_REPO
    ray_image = f"rayproject/{ray_project}:{ray_version}-{python_version}-{image_type}"
    anyscale_image = (
        f"{get_global_config()['byod_ecr']}/{anyscale_repo}:champagne-{ray_version}"
    )

    logger.info(f"Building champagne anyscale image from {ray_image}")
    with open(DATAPLANE_FILENAME, "rb") as build_file:
        subprocess.check_call(
            [
                "docker",
                "build",
                "--progress=plain",
                "--build-arg",
                f"BASE_IMAGE={ray_image}",
                "-t",
                anyscale_image,
                "-",
            ],
            stdin=build_file,
            stdout=sys.stderr,
            env=env,
        )
    _validate_and_push(anyscale_image)

    return anyscale_image


def build_anyscale_custom_byod_image(test: Test) -> None:
    if not test.require_custom_byod_image():
        logger.info(f"Test {test.get_name()} does not require a custom byod image")
        return
    byod_image = test.get_anyscale_byod_image()
    if _image_exist(byod_image):
        logger.info(f"Image {byod_image} already exists")
        return

    env = os.environ.copy()
    env["DOCKER_BUILDKIT"] = "1"
    subprocess.check_call(
        [
            "docker",
            "build",
            "--progress=plain",
            "--build-arg",
            f"BASE_IMAGE={test.get_anyscale_base_byod_image()}",
            "--build-arg",
            f"POST_BUILD_SCRIPT={test.get_byod_post_build_script()}",
            "-t",
            byod_image,
            "-f",
            os.path.join(RELEASE_BYOD_DIR, "byod.custom.Dockerfile"),
            RELEASE_BYOD_DIR,
        ],
        stdout=sys.stderr,
        env=env,
    )
    _validate_and_push(byod_image)


def build_anyscale_base_byod_images(tests: List[Test]) -> None:
    """
    Builds the Anyscale BYOD images for the given tests.
    """
    _download_dataplane_build_file()
    to_be_built = {}
    built = set()
    for test in tests:
        to_be_built[test.get_anyscale_base_byod_image()] = test

    env = os.environ.copy()
    env["DOCKER_BUILDKIT"] = "1"
    start = int(time.time())
    # ray images are built on post-merge, so we can wait for them to be available
    while (
        len(built) < len(to_be_built)
        and int(time.time()) - start < BASE_IMAGE_WAIT_TIMEOUT
    ):
        for byod_image, test in to_be_built.items():
            py_version = test.get_python_version()
            if test.use_byod_ml_image():
                byod_requirements = f"{REQUIREMENTS_ML_BYOD}_{py_version}.txt"
            else:
                byod_requirements = f"{REQUIREMENTS_BYOD}_{py_version}.txt"

            if _image_exist(byod_image):
                logger.info(f"Image {byod_image} already exists")
                built.add(byod_image)
                continue
            ray_image = test.get_ray_image()
            if not _image_exist(ray_image):
                # TODO(can): instead of waiting for the base image to be built, we can
                #  build it ourselves
                timeout = BASE_IMAGE_WAIT_TIMEOUT - (int(time.time()) - start)
                logger.info(
                    f"Image {ray_image} does not exist yet. "
                    f"Wait for another {timeout}s..."
                )
                time.sleep(BASE_IMAGE_WAIT_DURATION)
                continue
            logger.info(f"Building {byod_image} from {ray_image}")
            with open(DATAPLANE_FILENAME, "rb") as build_file:
                subprocess.check_call(
                    [
                        "docker",
                        "build",
                        "--progress=plain",
                        "--build-arg",
                        f"BASE_IMAGE={ray_image}",
                        "-t",
                        byod_image,
                        "-",
                    ],
                    stdin=build_file,
                    stdout=sys.stderr,
                    env=env,
                )
                subprocess.check_call(
                    [
                        "docker",
                        "build",
                        "--progress=plain",
                        "--build-arg",
                        f"BASE_IMAGE={byod_image}",
                        "--build-arg",
                        f"PIP_REQUIREMENTS={byod_requirements}",
                        "--build-arg",
                        "DEBIAN_REQUIREMENTS=requirements_debian_byod.txt",
                        "-t",
                        byod_image,
                        "-f",
                        os.path.join(RELEASE_BYOD_DIR, "byod.Dockerfile"),
                        RELEASE_BYOD_DIR,
                    ],
                    stdout=sys.stderr,
                    env=env,
                )
                _validate_and_push(byod_image)
                built.add(byod_image)


def _validate_and_push(byod_image: str) -> None:
    """
    Validates the given image and pushes it to ECR.
    """
    docker_ray_commit = (
        subprocess.check_output(
            [
                "docker",
                "run",
                "-ti",
                "--entrypoint",
                "python",
                byod_image,
                "-c",
                "import ray; print(ray.__commit__)",
            ],
        )
        .decode("utf-8")
        .strip()
    )
    if os.environ.get("RAY_IMAGE_TAG"):
        logger.info(f"Ray commit from image: {docker_ray_commit}")
    else:
        expected_ray_commit = _get_ray_commit()
        assert (
            docker_ray_commit == expected_ray_commit
        ), f"Expected ray commit {expected_ray_commit}, found {docker_ray_commit}"
    logger.info(f"Pushing image to registry: {byod_image}")
    subprocess.check_call(
        ["docker", "push", byod_image],
        stdout=sys.stderr,
    )


def _get_ray_commit(envs: Optional[Dict[str, str]] = None) -> str:
    if envs is None:
        envs = os.environ
    for key in [
        "RAY_WANT_COMMIT_IN_IMAGE",
        "COMMIT_TO_TEST",
        "BUILDKITE_COMMIT",
    ]:
        commit = envs.get(key, "")
        if commit:
            return commit
    return ""


def _download_dataplane_build_file() -> None:
    """
    Downloads the dataplane build file from S3.
    """
    s3 = boto3.client("s3")
    s3.download_file(
        Bucket=DATAPLANE_S3_BUCKET,
        Key=DATAPLANE_FILENAME,
        Filename=DATAPLANE_FILENAME,
    )
    with open(DATAPLANE_FILENAME, "rb") as build_context:
        digest = hashlib.sha256(build_context.read()).hexdigest()
        assert digest == DATAPLANE_DIGEST, "Mismatched dataplane digest found!"


def _image_exist(image: str) -> bool:
    """
    Checks if the given image exists in Docker
    """
    p = subprocess.run(
        ["docker", "manifest", "inspect", image],
        stdout=sys.stderr,
        stderr=sys.stderr,
    )
    return p.returncode == 0