import os
from pathlib import Path
import random
import shutil
import tempfile
from typing import Any, Dict, Union
from unittest.mock import patch
import uuid
from zipfile import ZipFile

import pytest
import smart_open
import yaml

from anyscale.sdk.anyscale_client.models import CreateProductionJob
from anyscale.sdk.anyscale_client.models.create_production_job_config import (
    CreateProductionJobConfig,
)
from anyscale.sdk.anyscale_client.sdk import (
    _upload_and_rewrite_working_dir_in_create_production_job,
    _upload_and_rewrite_working_dir_ray_serve_config,
)
from anyscale.utils.ray_utils import (  # type: ignore
    _dir_travel,
    _get_excludes,
    zip_directory,
)
from anyscale.utils.runtime_env import (
    _get_remote_storage_object_name,
    upload_and_rewrite_working_dir,
)


CONDA_DICT = {"dependencies": ["pip", {"pip": ["pip-install-test==0.5"]}]}
PIP_LIST = ["requests==1.0.0", "pip-install-test"]
GS_BUCKET_NAME = "anyscale-bk-e2e-test-job"


@pytest.fixture()
def test_directory():
    with tempfile.TemporaryDirectory() as tmp_dir:
        path = Path(tmp_dir)
        subdir = path / "subdir"
        subdir.mkdir(parents=True)
        requirements_file = subdir / "requirements.txt"
        with requirements_file.open(mode="w") as f:
            print("\n".join(PIP_LIST), file=f)

        good_conda_file = subdir / "good_conda_env.yaml"
        with good_conda_file.open(mode="w") as f:
            yaml.dump(CONDA_DICT, f)

        bad_conda_file = subdir / "bad_conda_env.yaml"
        with bad_conda_file.open(mode="w") as f:
            print("% this is not a YAML file %", file=f)

        gitignore_file = subdir / ".gitignore"
        with gitignore_file.open(mode="w") as f:
            print("file_in_gitignore.txt", file=f)

        file_in_gitignore = subdir / "file_in_gitignore.txt"
        with file_in_gitignore.open(mode="w") as f:
            print("file_in_gitignore", file=f)

        yield subdir, requirements_file, good_conda_file, bad_conda_file


@pytest.mark.parametrize(
    ("upload_path", "expected_output"),
    [
        ("gs://bucket", "file.zip"),
        ("gs://bucket/path", "path/file.zip"),
        ("gs://bucket/path/", "path/file.zip"),
    ],
)
def test_get_remote_storage_object_name(upload_path, expected_output):
    assert (
        _get_remote_storage_object_name(upload_path, upload_filename="file.zip")
        == expected_output
    )


class TestUploadLocalWorkingDir:
    def test_upload_and_rewrite_working_dir_with_fake_upload(self, test_directory):
        def fake_upload_file_to_remote_storage(
            source_file: str, upload_path: str, upload_filename: str
        ) -> str:
            shutil.copyfile(source_file, upload_filename)
            return os.path.join(upload_path, upload_filename)

        upload_path = "s3://fake/dir"
        runtime_env = {
            "pip": ["blah"],
            "working_dir": str(test_directory),
            "upload_path": upload_path,
        }
        new_runtime_env = upload_and_rewrite_working_dir(
            runtime_env,
            upload_file_to_remote_storage_fn=fake_upload_file_to_remote_storage,
        )

        assert new_runtime_env["working_dir"].startswith(upload_path)
        assert "upload_path" not in new_runtime_env
        assert Path(new_runtime_env["working_dir"]).suffix == ".zip"

    @pytest.mark.skipif(
        os.environ.get("CI") != "true",
        reason="Skip test if not running in CI because S3/GS credentials are not available",
    )
    @pytest.mark.parametrize(
        "upload_path",
        [
            "s3://bk-premerge-first-jawfish-artifacts/e2e_tests/job",
            f"gs://{GS_BUCKET_NAME}",
        ],
    )
    @pytest.mark.parametrize("mode", ["sdk", "cli"])
    def test_upload_and_rewrite_working_dir(
        self, test_directory, upload_path: str, mode
    ):
        """Test that the working_dir correctly gets uploaded to upload_path in the runtime_env."""
        subdir, _, _, _ = test_directory
        runtime_env = {
            "pip": ["blah"],
            "working_dir": str(subdir),
            "upload_path": upload_path,
        }
        if mode == "sdk":
            job_config = CreateProductionJobConfig(
                compute_config_id="cpt_U8RCfD7Wr1vCD4iqGi4cBbj1",
                build_id="bld_1277XIinoJmiM8Z3gNdcHN",
                runtime_env=runtime_env,
                entrypoint="python my_job_script.py --option1=value1",
                max_retries=3,
            )

            create_production_job = CreateProductionJob(
                name="my-production-job",
                description="A production job running on Anyscale.",
                project_id="prj_7S7Os7XBvO6vdiVC1J0lgj",
                config=job_config,
            )

            _upload_and_rewrite_working_dir_in_create_production_job(
                create_production_job
            )
            new_runtime_env = create_production_job.config.runtime_env
        elif mode == "cli":
            # Upload the working_dir and rewrite the working_dir field to point to the new remote URI
            new_runtime_env = upload_and_rewrite_working_dir(runtime_env)
        else:
            raise ValueError(f"Invalid mode: {mode}")

        # Check all non-working_dir fields in the runtime_env are unchanged
        runtime_env_copy = runtime_env.copy()
        new_runtime_env_copy = new_runtime_env.copy()
        runtime_env_copy.pop("working_dir")
        runtime_env_copy.pop("upload_path")
        new_runtime_env_copy.pop("working_dir")
        assert runtime_env_copy == new_runtime_env_copy

        # Check the format of the remote URI
        remote_uri = new_runtime_env["working_dir"]
        assert remote_uri.startswith(upload_path)
        assert Path(remote_uri).suffix == ".zip"

        # Download the zip file from the remote bucket
        local_zip_file = "test_job_controller_working_dir.zip"
        with smart_open.open(remote_uri, "rb") as package_zip, open(
            local_zip_file, "wb"
        ) as fin:
            fin.write(package_zip.read())

        # Check the downloaded zip is the same directory that was uploaded
        unzip_target_dir = os.path.join("unzip_target_dir")
        with ZipFile(local_zip_file, "r") as zip_ref:
            zip_ref.extractall(unzip_target_dir)
        assert set(os.listdir(os.path.join(unzip_target_dir, "subdir"))) == {
            "requirements.txt",
            "good_conda_env.yaml",
            "bad_conda_env.yaml",
            ".gitignore",
        }

    @pytest.mark.skipif(
        os.environ.get("CI") != "true",
        reason="Skip test if not running in CI because S3/GS credentials are not available",
    )
    @pytest.mark.parametrize(
        "upload_path",
        [
            "s3://bk-premerge-first-jawfish-artifacts/e2e_tests/job",
            f"gs://{GS_BUCKET_NAME}",
        ],
    )
    def test_upload_and_rewrite_working_dir_service_v2(
        self, test_directory, upload_path: str
    ):
        """Test that the working_dir correctly gets uploaded to upload_path in the runtime_env."""
        subdir, _, _, _ = test_directory
        runtime_env = {
            "pip": ["blah"],
            "working_dir": str(subdir),
            "upload_path": upload_path,
        }
        ray_serve_config = {
            "runtime_env": runtime_env,
            "import_path": "serve_deploy_update:bound",
            "deployments": [
                {
                    "name": "Updatable",
                    "num_replicas": 2,
                    "user_config": {
                        "response": "Hello from e2e test!",
                        "should_fail": False,
                    },
                },
            ],
        }
        new_ray_serve_config = _upload_and_rewrite_working_dir_ray_serve_config(
            ray_serve_config
        )
        new_runtime_env = new_ray_serve_config["runtime_env"]

        # Check all non-working_dir fields in the runtime_env are unchanged
        runtime_env_copy = runtime_env.copy()
        new_runtime_env_copy = new_runtime_env.copy()
        runtime_env_copy.pop("working_dir")
        runtime_env_copy.pop("upload_path")
        new_runtime_env_copy.pop("working_dir")
        assert runtime_env_copy == new_runtime_env_copy

        # Check the format of the remote URI
        remote_uri = new_runtime_env["working_dir"]
        assert remote_uri.startswith(upload_path)
        assert Path(remote_uri).suffix == ".zip"

        # Download the zip file from the remote bucket
        local_zip_file = "test_job_controller_working_dir.zip"
        with smart_open.open(remote_uri, "rb") as package_zip, open(
            local_zip_file, "wb"
        ) as fin:
            fin.write(package_zip.read())

        # Check the downloaded zip is the same directory that was uploaded
        unzip_target_dir = os.path.join("unzip_target_dir")
        with ZipFile(local_zip_file, "r") as zip_ref:
            zip_ref.extractall(unzip_target_dir)
        assert set(os.listdir(os.path.join(unzip_target_dir, "subdir"))) == {
            "requirements.txt",
            "good_conda_env.yaml",
            "bad_conda_env.yaml",
            ".gitignore",
        }

    @pytest.mark.parametrize(
        "job_config",
        [
            CreateProductionJobConfig(
                compute_config_id="cpt_U8RCfD7Wr1vCD4iqGi4cBbj1",
                build_id="bld_1277XIinoJmiM8Z3gNdcHN",
                entrypoint="python my_job_script.py --option1=value1",
                max_retries=3,
            ),
            CreateProductionJobConfig(
                compute_config_id="cpt_U8RCfD7Wr1vCD4iqGi4cBbj1",
                build_id="bld_1277XIinoJmiM8Z3gNdcHN",
                entrypoint="python my_job_script.py --option1=value1",
                runtime_env={"pip": ["blah"],},
                max_retries=3,
            ),
            {
                "compute_config_id": "cpt_U8RCfD7Wr1vCD4iqGi4cBbj1",
                "build_id": "bld_1277XIinoJmiM8Z3gNdcHN",
                "entrypoint": "python my_job_script.py --option1=value1",
                "max_retries": 3,
            },
            {
                "compute_config_id": "cpt_U8RCfD7Wr1vCD4iqGi4cBbj1",
                "build_id": "bld_1277XIinoJmiM8Z3gNdcHN",
                "entrypoint": "python my_job_script.py --option1=value1",
                "runtime_env": {"pip": ["blah"],},
                "max_retries": 3,
            },
        ],
    )
    def test_empty_runtime_env(
        self, job_config: Union[Dict[str, Any], CreateProductionJobConfig]
    ):
        original_runtime_env = None
        if isinstance(job_config, dict):
            original_runtime_env = job_config.get("runtime_env")
        else:
            original_runtime_env = job_config.runtime_env
        create_production_job = CreateProductionJob(
            name="my-production-job",
            description="A production job running on Anyscale.",
            project_id="prj_7S7Os7XBvO6vdiVC1J0lgj",
            config=job_config,
        )

        _upload_and_rewrite_working_dir_in_create_production_job(create_production_job)
        if isinstance(create_production_job.config, dict):
            assert (
                create_production_job.config.get("runtime_env") == original_runtime_env
            )
        else:
            assert create_production_job.config.runtime_env == original_runtime_env

    def test_remote_uri_working_dir_noop(self):
        runtime_env = {"working_dir": "s3://already-a-uri"}
        assert runtime_env == upload_and_rewrite_working_dir(runtime_env)

    @pytest.mark.parametrize("include_parent_dir", [True, False])
    @pytest.mark.parametrize("diable_gitignore_exclusion", [True, False])
    def test_zip_directory(
        self, test_directory, include_parent_dir, diable_gitignore_exclusion
    ):
        subdir, _, _, _ = test_directory

        with tempfile.TemporaryDirectory() as tmpdir, patch(
            "anyscale.utils.ray_utils.DISABLE_GITIGNORE_EXCLUSION",
            diable_gitignore_exclusion,
        ):
            zip_file = os.path.join(tmpdir, "test_zip.zip")
            # NOTE(architkulkarni): Must use "requirements.txt" in excludes instead of requirements_file
            # until https://github.com/ray-project/ray/issues/23473 is fixed.
            zip_directory(
                subdir,
                excludes=["requirements.txt"],
                output_path=zip_file,
                include_parent_dir=include_parent_dir,
            )
            unzip_target_dir = os.path.join(tmpdir, "unzip_target_dir")
            with ZipFile(zip_file, "r") as zip_ref:
                zip_ref.extractall(unzip_target_dir)
            if include_parent_dir:
                assert os.listdir(unzip_target_dir) == ["subdir"]
                unzip_target_dir = os.path.join(unzip_target_dir, "subdir")

            expected_files = {
                "good_conda_env.yaml",
                "bad_conda_env.yaml",
                ".gitignore",
            }
            if diable_gitignore_exclusion:
                expected_files.add("file_in_gitignore.txt")

            assert set(os.listdir(unzip_target_dir)) == expected_files

    def test_travel(self, tmp_path):
        """Copied from OSS Ray at python/ray/tests/test_runtime_env_packaging.py"""
        dir_paths = set()
        file_paths = set()
        item_num = 0
        excludes = []
        root = tmp_path / "test"

        def construct(path, excluded=False, depth=0):
            nonlocal item_num
            path.mkdir(parents=True)
            if not excluded:
                dir_paths.add(str(path))
            if depth > 8:
                return
            if item_num > 500:
                return
            dir_num = random.randint(0, 10)
            file_num = random.randint(0, 10)
            for _ in range(dir_num):
                uid = str(uuid.uuid4()).split("-")[0]
                dir_path = path / uid
                exclud_sub = random.randint(0, 5) == 0
                if not excluded and exclud_sub:
                    excludes.append(str(dir_path.relative_to(root)))
                if not excluded:
                    construct(dir_path, exclud_sub or excluded, depth + 1)
                item_num += 1
            if item_num > 1000:
                return

            for _ in range(file_num):
                uid = str(uuid.uuid4()).split("-")[0]
                v = random.randint(0, 1000)
                with (path / uid).open("w") as f:
                    f.write(str(v))
                if not excluded:
                    if random.randint(0, 5) == 0:
                        excludes.append(str((path / uid).relative_to(root)))
                    else:
                        file_paths.add((str(path / uid), str(v)))
                item_num += 1

        construct(root)
        exclude_spec = _get_excludes(root, excludes)
        visited_dir_paths = set()
        visited_file_paths = set()

        def handler(path):
            if path.is_dir():
                visited_dir_paths.add(str(path))
            else:
                with open(path) as f:
                    visited_file_paths.add((str(path), f.read()))

        _dir_travel(root, [exclude_spec], handler)
        assert file_paths == visited_file_paths
        assert dir_paths == visited_dir_paths


@pytest.mark.skipif(
    os.environ.get("CI") != "true",
    reason="Skip test if not running in CI because credentials are not available.",
)
def test_gs_upload_download():
    """Test that the CI machine has access to the Google Storage bucket used for testing."""
    random_int = random.randint(0, 100000)
    filename = f"test_gs_upload_download_{random_int}.txt"

    from google.cloud import storage

    storage_client = storage.Client()
    bucket_obj = storage_client.bucket(GS_BUCKET_NAME)
    blob = bucket_obj.blob(filename)
    try:
        random_str = str(random_int)
        blob.upload_from_string(random_str)
        assert blob.exists()

        from smart_open import open

        uri = f"gs://{GS_BUCKET_NAME}/{filename}"

        with open(uri, "r") as f:
            assert random_str == f.read()
    finally:
        # Delete the file from the Google Storage bucket
        blob.delete()
