import sys
from pathlib import Path

import pytest
from botocore.stub import Stubber

from conductor.cli.cmd.clean import BranchManager, clean
from tests.cli.cmd.test_files import FIXTURE_DIR
from tests.conftest import IMAGE_ECR_RESPONSE, S3_BUCKET_RESPONSES

IMAGE_ECR_RESPONSE_BAD = {
    "imageIds": [
        {"imageDigest": "a1", "imageTag": "main-123"},
        {"imageDigest": "a2", "imageTag": "test1-123"},
        {"imageDigest": "a3", "imageTag": "closed-test1-123"},
        {"imageDigest": "a4", "imageTag": "closed-test2-123"},
        {"imageDigest": "a5", "imageTag": "closed_test3_123"},
    ]
}


def test_get_s3_bucket_response(mocker, monkeypatch):
    mocker.patch("conductor.cli.cmd.clean.BranchManager._prune_remote_branch")
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager._get_live_branches",
        return_value=["main", "test1", "test2", "test3"],
    )
    monkeypatch.setenv(
        "AWS_DEFAULT_REGION", "us-west-2"
    )  # Allows BranchManager init to work before stubbing.
    branch_manager = BranchManager(
        project_root=".",
        s3_bucket="test-project.test-env",
        s3_prefix="",
        airflow_bucket="test-airflow-bucket",
        airflow_prefix="dags/",
        delimiter="/",
        dry_run=False,
    )
    stubber = Stubber(branch_manager.s3_client)
    bucket = "test-project.test-env"
    prefix = ""
    delimiter = "/"
    expected_params = {"Bucket": bucket, "Prefix": prefix, "Delimiter": delimiter}
    stubber.add_response("list_objects_v2", S3_BUCKET_RESPONSES, expected_params)
    stubber.activate()
    resulted_s3_response = branch_manager._get_s3_bucket_response(
        bucket, prefix, delimiter
    )
    assert resulted_s3_response == S3_BUCKET_RESPONSES


def test_get_ecr_response(mocker, monkeypatch):
    mocker.patch("conductor.cli.cmd.clean.BranchManager._prune_remote_branch")
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager._get_live_branches",
        return_value=["main", "test1", "test2", "test3"],
    )
    monkeypatch.setenv(
        "AWS_DEFAULT_REGION", "us-west-2"
    )  # Allows BranchManager init to work before stubbing.
    branch_manager = BranchManager(
        project_root=".",
        s3_bucket="test-project.test-env",
        s3_prefix="",
        airflow_bucket="test-airflow-bucket",
        airflow_prefix="dags/",
        delimiter="/",
        dry_run=False,
    )
    stubber = Stubber(branch_manager.ecr_client)
    expected_params = {"repositoryName": "test-project.test-env"}
    stubber.add_response("list_images", IMAGE_ECR_RESPONSE, expected_params)
    stubber.activate()
    resulted_s3_response = branch_manager._get_ecr_response()
    assert resulted_s3_response == IMAGE_ECR_RESPONSE


def test_prompt_for_branch(branch_manager, mocker):
    mocker.patch(
        "conductor.cli.cmd.clean.input", side_effect=["Y", "YES", "no", "YESSS"]
    )
    _, _, manager = branch_manager
    assert manager.prompt_for_branch("test_branch_1") is True
    assert manager.prompt_for_branch("test_branch_1") is True
    assert manager.prompt_for_branch("test_branch_1") is False
    assert manager.prompt_for_branch("test_branch_1") is False


def test_stage_for_deletion(branch_manager, mocker):
    _, _, branch_manager_mocks = branch_manager
    to_be_deleted = ["closed_1", "closed_2", "closed_3"]
    for i in to_be_deleted:
        branch_manager_mocks.stage_for_deletion(i)
    assert to_be_deleted == branch_manager_mocks.to_be_deleted_branch


@pytest.mark.datafiles(Path(FIXTURE_DIR, "test_project"))
def test_clean_delete_all(datafiles, branch_manager, mocker):
    s3_deletion_mock, ecr_deletion_mock, branch_manager_mocks = branch_manager
    sys.path.append(str(datafiles))
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager.prompt_for_branch", return_value=True
    )
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager",
        return_value=branch_manager_mocks,
    )
    mocker.patch("conductor.cli.cmd.clean.get_dags_s3_path", return_value=[None, None])
    clean(datafiles)
    s3_deletion_mock.assert_any_call("test-project.test-env", "closed_test1")
    s3_deletion_mock.assert_any_call("test-project.test-env", "closed-test2")
    s3_deletion_mock.assert_any_call(
        "test-airflow-bucket", "dags/test-project.test-env.closed_test1.py"
    )
    s3_deletion_mock.assert_any_call(
        "test-airflow-bucket", "dags/test-project.test-env.closed-test2.py"
    )
    ecr_deletion_mock.assert_any_call("a3")
    ecr_deletion_mock.assert_any_call("a4")
    ecr_deletion_mock.assert_any_call("a5")
    assert sorted(branch_manager_mocks.to_be_deleted_branch) == sorted(
        ["closed_test1", "closed-test2", "closed_test3"]
    )


@pytest.mark.datafiles(Path(FIXTURE_DIR, "test_project"))
def test_clean_delete_none(datafiles, branch_manager, mocker):
    s3_deletion_mock, ecr_deletion_mock, branch_manager_mocks = branch_manager
    sys.path.append(str(datafiles))
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager.prompt_for_branch", return_value=False
    )
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager",
        return_value=branch_manager_mocks,
    )
    mocker.patch("conductor.cli.cmd.clean.get_dags_s3_path", return_value=[None, None])
    clean(datafiles)
    s3_deletion_mock.assert_not_called()
    ecr_deletion_mock.assert_not_called()
    assert branch_manager_mocks.to_be_deleted_branch == []


@pytest.mark.datafiles(Path(FIXTURE_DIR, "test_project"))
@pytest.mark.dry_run
def test_clean_delete_all_dry_run(datafiles, branch_manager, mocker):
    s3_deletion_mock, ecr_deletion_mock, branch_manager_mocks = branch_manager
    sys.path.append(str(datafiles))
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager.prompt_for_branch", return_value=True
    )
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager",
        return_value=branch_manager_mocks,
    )
    mocker.patch("conductor.cli.cmd.clean.get_dags_s3_path", return_value=[None, None])
    clean(datafiles, dry_run=True)
    s3_deletion_mock.assert_not_called()
    ecr_deletion_mock.assert_not_called()
    assert sorted(branch_manager_mocks.to_be_deleted_branch) == sorted(
        ["closed_test1", "closed-test2", "closed_test3"]
    )


def mocked_prompt_for_branch(self, branch):
    if branch == "closed-test2":
        return True
    else:
        False


@pytest.mark.datafiles(Path(FIXTURE_DIR, "test_project"))
def test_clean_delete_closed_test2(datafiles, branch_manager, mocker):
    s3_deletion_mock, ecr_deletion_mock, branch_manager_mocks = branch_manager
    sys.path.append(str(datafiles))
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager.prompt_for_branch",
        new=mocked_prompt_for_branch,
    )
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager",
        return_value=branch_manager_mocks,
    )
    mocker.patch("conductor.cli.cmd.clean.get_dags_s3_path", return_value=[None, None])
    clean(datafiles)
    s3_deletion_mock.assert_any_call(
        "test-airflow-bucket", "dags/test-project.test-env.closed-test2.py"
    )
    ecr_deletion_mock.assert_any_call("a4")
    assert branch_manager_mocks.to_be_deleted_branch == ["closed-test2"]


def test_collect_ecr_resources_exception(branch_manager, mocker):
    _, _, branch_manager_mocks = branch_manager
    mocker.patch(
        "conductor.cli.cmd.clean.BranchManager._get_ecr_response",
        return_value=IMAGE_ECR_RESPONSE_BAD,
    )
    with pytest.raises(Exception) as e_info:
        branch_manager_mocks._collect_ecr_resources()
        assert (
            e_info
            == "closed_test_3 could not find the appropriate branch name. \
                    imageTag has to be in the format of <branch_name>-<commit-hash>"
        )
