# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See LICENSE for more details.
#
# Copyright (c) 2026 ScyllaDB

from unittest.mock import MagicMock, patch

import pytest
from botocore.exceptions import ClientError

from sdcm.provision.aws.az_resolver import (
    AZResolver,
    NoValidAvailabilityZoneError,
    _node_count_positive,
    run_pre_flight_capacity_probe,
)
from sdcm.provision.common.fallback import is_az_fallback_enabled
from sdcm.provision.aws.capacity_errors import ProvisioningCapacityExhausted

from unit_tests.lib.dot_dict import DotDict


def _make_params(**overrides):
    params = DotDict(
        {
            "cluster_backend": "aws",
            "region_name": "us-east-1",
            "availability_zone": "a",
            "instance_type_db": "i4i.large",
            "instance_type_loader": "t3.small",
            "instance_type_monitor": "t3.small",
            "n_db_nodes": 3,
            "n_loaders": 1,
            "n_monitor_nodes": 1,
            "pre_filter_unavailable_availability_zones": True,
            **overrides,
        }
    )
    params.region_names = params["region_name"].split()
    return params


@pytest.fixture(name="mock_aws_region_cls")
def mock_aws_region_cls_fixture():
    """Patch AwsRegion with a mock returning configurable AZ lists."""
    with patch("sdcm.provision.aws.az_resolver.AwsRegion") as mock_cls:
        instance = mock_cls.return_value
        # default: every type is offered in a, b, c, d
        instance.get_common_availability_zones.return_value = ["us-east-1a", "us-east-1b", "us-east-1c", "us-east-1d"]
        instance.region_name = "us-east-1"
        yield mock_cls, instance


def test_required_instance_types_collects_all_role_types_when_active():
    params = _make_params(
        instance_type_db="i4i.large",
        instance_type_loader="t3.small",
        instance_type_monitor="t3.medium",
        instance_type_db_oracle="i3.large",
        instance_type_db_target="m6g.large",
        nemesis_grow_shrink_instance_type="i4i.xlarge",
        zero_token_instance_type_db="t3.micro",
        instance_type_vector_store="t4g.medium",
        n_loaders=1,
        n_monitor_nodes=1,
        n_test_oracle_db_nodes=1,
        db_type="mixed_scylla",
        n_db_zero_token_nodes=1,
        use_zero_nodes=True,
        n_vector_store_nodes=2,
    )
    types = AZResolver(params).required_instance_types()
    assert set(types) == {
        "i4i.large",
        "t3.small",
        "t3.medium",
        "i3.large",
        "m6g.large",
        "i4i.xlarge",
        "t3.micro",
        "t4g.medium",
    }


def test_required_instance_types_skips_empty_and_missing():
    params = _make_params(
        instance_type_db="i4i.large",
        instance_type_loader="",
        instance_type_monitor=None,
        # instance_type_db_oracle missing entirely
        n_loaders=0,
        n_monitor_nodes=0,
    )
    assert AZResolver(params).required_instance_types() == ["i4i.large"]


def test_required_instance_types_excludes_types_with_zero_node_count():
    """Instance types whose node count is 0 or feature flag is off must not constrain AZ filtering."""
    params = _make_params(
        instance_type_db="i4i.large",
        instance_type_loader="c6i.xlarge",  # excluded: n_loaders=0
        instance_type_monitor="t3.large",  # excluded: n_monitor_nodes=0
        instance_type_vector_store="t4g.medium",  # excluded: n_vector_store_nodes=0
        instance_type_db_oracle="i8g.2xlarge",  # excluded: db_type != mixed
        zero_token_instance_type_db="i4i.large",  # excluded: use_zero_nodes=False
        n_loaders=0,
        n_monitor_nodes=0,
        n_vector_store_nodes=0,
        n_test_oracle_db_nodes=1,
        db_type="scylla",
        use_zero_nodes=False,
    )
    assert AZResolver(params).required_instance_types() == ["i4i.large"]


def test_resolve_disabled_returns_unchanged(mock_aws_region_cls):
    params = _make_params(pre_filter_unavailable_availability_zones=False, availability_zone="a,b")
    AZResolver(params).resolve()
    assert params["availability_zone"] == "a,b"


def test_resolve_replaces_invalid_az_with_valid_alternative(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    # configured "a" not in the supported list - must be replaced
    region_instance.get_common_availability_zones.return_value = ["us-east-1b", "us-east-1c"]
    params = _make_params(availability_zone="a")
    AZResolver(params).resolve()
    assert params["availability_zone"] in {"b", "c"}


def test_resolve_multi_az_drops_unsupported_and_fills_alternatives(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = ["us-east-1a", "us-east-1c", "us-east-1d"]
    # configured a,b,e — only "a" supported; must replace b + e with valid alternatives
    params = _make_params(availability_zone="a,b,e")
    AZResolver(params).resolve()
    result = params["availability_zone"].split(",")
    assert len(result) == 3
    assert set(result) <= {"a", "c", "d"}  # only valid AZs, no duplicates implied by len check
    assert "a" in result
    assert "b" not in result
    assert "e" not in result
    assert len(set(result)) == 3


def test_resolve_raises_when_no_valid_az_in_region(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = []
    params = _make_params(availability_zone="a")
    with pytest.raises(NoValidAvailabilityZoneError):
        AZResolver(params).resolve()


def test_resolve_target_intersection_when_target_differs_from_db(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    # only "c" supports both db (x86) and target (arm)
    region_instance.get_common_availability_zones.return_value = ["us-east-1c"]
    params = _make_params(
        instance_type_db="i4i.large",
        instance_type_db_target="m6g.large",
        availability_zone="a",
    )
    AZResolver(params).resolve()
    assert params["availability_zone"] == "c"


@pytest.fixture(name="mock_multi_region")
def mock_multi_region_fixture():
    """Patch AwsRegion to return configurable per-region AZ lists."""
    region_returns: dict[str, list[str]] = {}

    class _RegionStub:
        def __init__(self, region_name):
            self.region_name = region_name

        def get_common_availability_zones(self, instance_types, preferred_azs):  # noqa: ARG002
            return region_returns.get(self.region_name, [])

    with patch("sdcm.provision.aws.az_resolver.AwsRegion", _RegionStub):
        yield region_returns


def test_resolve_multi_region_intersects_supported_az_letters(mock_multi_region):
    mock_multi_region["us-east-1"] = ["us-east-1a", "us-east-1b"]
    mock_multi_region["eu-west-1"] = ["eu-west-1b", "eu-west-1c"]
    params = _make_params(region_name="us-east-1 eu-west-1", availability_zone="a")
    AZResolver(params).resolve()
    assert params["availability_zone"] == "b"


def test_resolve_multi_region_raises_when_no_common_letter(mock_multi_region):
    mock_multi_region["us-east-1"] = ["us-east-1a", "us-east-1b"]
    mock_multi_region["eu-west-1"] = ["eu-west-1c", "eu-west-1d"]
    params = _make_params(region_name="us-east-1 eu-west-1", availability_zone="a")
    with pytest.raises(NoValidAvailabilityZoneError):
        AZResolver(params).resolve()


def test_resolve_multi_region_multi_az_drops_unsupported_and_fills(mock_multi_region):
    """Multi-AZ × multi-region: letter must be supported in EVERY region; drop+fill applies."""
    mock_multi_region["us-east-1"] = ["us-east-1a", "us-east-1b", "us-east-1c", "us-east-1d"]
    mock_multi_region["eu-west-1"] = ["eu-west-1b", "eu-west-1c", "eu-west-1d"]
    params = _make_params(region_name="us-east-1 eu-west-1", availability_zone="a,b,c")
    AZResolver(params).resolve()
    result = params["availability_zone"].split(",")
    assert set(result) == {"b", "c", "d"}


@pytest.mark.parametrize(
    "value, expected",
    [
        # scalars
        (None, False),
        (True, True),
        (False, False),
        (0, False),
        (1, True),
        (3, True),
        (-1, False),
        (0.0, False),
        (1.5, True),
        # lists
        ([], False),
        ([0, 0], False),
        ([2, 0], True),
        ([-1, 0], False),
        (["2", "0"], True),
        (["0"], False),
        # strings
        ("", False),
        ("0", False),
        ("3", True),
        ("3 4", True),
        ("0 0", False),
        ("-1", False),
        ("abc", False),
        # unsupported types
        ({"a": 1}, False),
    ],
)
def test_node_count_positive(value, expected):
    assert _node_count_positive(value) is expected


def test_required_instance_types_skips_target_when_none():
    """`instance_type_db_target` unset / None must not constrain AZ filtering."""
    params = _make_params(
        instance_type_db="i4i.large",
        instance_type_db_target=None,
        nemesis_grow_shrink_instance_type="",
        instance_type_loader="",
        n_loaders=0,
        instance_type_monitor="",
        n_monitor_nodes=0,
    )
    assert AZResolver(params).required_instance_types() == ["i4i.large"]


@pytest.mark.parametrize("az_value", ["", None])
def test_resolve_with_unset_availability_zone_stays_unset(mock_aws_region_cls, az_value):
    """Empty / None `availability_zone` means no explicit user choice; resolver leaves it effectively unset."""
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = ["us-east-1b", "us-east-1c"]
    params = _make_params(availability_zone=az_value)
    AZResolver(params).resolve()
    assert not params["availability_zone"]


@pytest.mark.parametrize(
    "new_param, alias, expected",
    [
        (None, True, True),
        (False, True, False),
    ],
)
def test_is_az_fallback_enabled(new_param, alias, expected):
    params = _make_params(fallback_to_next_availability_zone=new_param, aws_fallback_to_next_availability_zone=alias)
    assert is_az_fallback_enabled(params) is expected


def test_get_fallback_candidates_single_az(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = ["us-east-1a", "us-east-1b", "us-east-1c"]
    params = _make_params(availability_zone="a")
    assert AZResolver(params).get_fallback_candidates() == [["a"], ["b"], ["c"]]


def test_get_fallback_candidates_multi_az(mock_aws_region_cls):
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = [
        "us-east-1a",
        "us-east-1b",
        "us-east-1c",
        "us-east-1d",
    ]
    params = _make_params(availability_zone="a,b,c")
    candidates = AZResolver(params).get_fallback_candidates()
    # first candidate is the configured set; subsequent candidates replace one slot at a time with "d"
    assert candidates[0] == ["a", "b", "c"]
    assert ["d", "b", "c"] in candidates
    assert ["a", "d", "c"] in candidates
    assert ["a", "b", "d"] in candidates


def test_get_fallback_candidates_enumerates_alternatives_when_upfront_filter_disabled(mock_aws_region_cls):
    """Fallback enumeration must NOT depend on `pre_filter_unavailable_availability_zones`."""
    _, region_instance = mock_aws_region_cls
    region_instance.get_common_availability_zones.return_value = ["us-east-1a", "us-east-1b"]
    params = _make_params(availability_zone="c", pre_filter_unavailable_availability_zones=False)
    assert AZResolver(params).get_fallback_candidates() == [["a"], ["b"]]


def _capacity_error() -> ClientError:
    return ClientError({"Error": {"Code": "InsufficientInstanceCapacity", "Message": "no capacity"}}, "RunInstances")


def _non_capacity_error() -> ClientError:
    return ClientError({"Error": {"Code": "UnauthorizedOperation", "Message": "not allowed"}}, "RunInstances")


@pytest.fixture(name="mock_probe_aws")
def mock_probe_aws_fixture():
    """Patch boto3 clients and AwsRegion used by `probe_capacity_for_types`.
    Yields the ec2 mock so tests can configure run_instances / terminate_instances.
    """
    ec2 = MagicMock()
    ec2.describe_instance_types.return_value = {
        "InstanceTypes": [{"ProcessorInfo": {"SupportedArchitectures": ["x86_64"]}}]
    }
    ssm = MagicMock()
    ssm.get_parameter.return_value = {"Parameter": {"Value": "ami-deadbeef"}}

    clients = {"ec2": ec2, "ssm": ssm}

    region_obj = MagicMock()
    region_obj.sct_subnet.return_value = MagicMock(subnet_id="subnet-abc")
    region_obj.sct_security_group = MagicMock(group_id="sg-xyz")

    with (
        patch("sdcm.provision.aws.az_resolver.boto3.client", side_effect=lambda svc, **_: clients[svc]),
        patch("sdcm.provision.aws.az_resolver.AwsRegion", return_value=region_obj),
        patch("sdcm.provision.aws.az_resolver.TestConfig.common_tags", return_value={"RunByUser": "tester"}),
    ):
        yield ec2


def test_probe_capacity_returns_true_on_successful_launch(mock_probe_aws):
    mock_probe_aws.run_instances.return_value = {"Instances": [{"InstanceId": "i-1"}]}
    params = _make_params(test_id="t-1", instance_type_db_target="m6g.large")

    results = AZResolver(params).probe_capacity_for_types(types=["m6g.large"], az="a")

    assert results == {"m6g.large": True}
    mock_probe_aws.terminate_instances.assert_called_once_with(InstanceIds=["i-1"])


def test_probe_capacity_returns_false_on_capacity_error(mock_probe_aws):
    mock_probe_aws.run_instances.side_effect = _capacity_error()
    resolver = AZResolver(_make_params(test_id="t-1"))

    assert resolver.probe_capacity_for_types(types=["m6g.large"], az="a") == {"m6g.large": False}

    mock_probe_aws.terminate_instances.assert_not_called()


def test_probe_capacity_passes_through_non_capacity_error(mock_probe_aws):
    """SCP denies, IAM gaps, etc. must not fail the test — the probe degrades to a no-op pass."""
    mock_probe_aws.run_instances.side_effect = _non_capacity_error()
    resolver = AZResolver(_make_params(test_id="t-1"))

    assert resolver.probe_capacity_for_types(types=["m6g.large"], az="a") == {"m6g.large": True}

    mock_probe_aws.terminate_instances.assert_not_called()


@pytest.mark.parametrize(
    "overrides",
    [
        {"pre_flight_capacity_probe": False, "instance_type_db_target": "m6g.large"},
        {
            "pre_flight_capacity_probe": True,
            "instance_type_db_target": None,
            "nemesis_grow_shrink_instance_type": "",
        },
    ],
    ids=["flag_off", "no_dynamic_types"],
)
def test_run_pre_flight_capacity_probe_noop(mock_probe_aws, overrides):
    """The orchestrator skips all AWS calls when the flag is off or no dynamic-add instance types are configured."""
    run_pre_flight_capacity_probe(_make_params(**overrides))
    mock_probe_aws.run_instances.assert_not_called()


def test_run_pre_flight_capacity_probe_succeeds_when_all_types_available(mock_probe_aws):
    mock_probe_aws.run_instances.return_value = {"Instances": [{"InstanceId": "i-1"}]}
    run_pre_flight_capacity_probe(
        _make_params(
            pre_flight_capacity_probe=True,
            instance_type_db_target="m6g.large",
            nemesis_grow_shrink_instance_type="i4i.xlarge",
            availability_zone="a,b",
            test_id="t-1",
        )
    )

    assert mock_probe_aws.run_instances.call_count == 4
    assert mock_probe_aws.terminate_instances.call_count == 4


def test_run_pre_flight_capacity_probe_raises_on_capacity_failure(mock_probe_aws):
    mock_probe_aws.run_instances.side_effect = _capacity_error()
    with pytest.raises(ProvisioningCapacityExhausted, match="m6g.large@a"):
        run_pre_flight_capacity_probe(
            _make_params(
                pre_flight_capacity_probe=True,
                instance_type_db_target="m6g.large",
                availability_zone="a",
                test_id="t-1",
            )
        )


def test_get_region_fallback_candidates_excludes_current_and_non_peered():
    """Only actively-peered regions (other than the current one) become candidates."""
    params = _make_params(region_name="us-east-1", availability_zone="a")
    resolver = AZResolver(params)
    peered = {"eu-west-1", "us-west-2"}

    with (
        patch.object(AZResolver, "_is_region_peered", side_effect=lambda cur, cand: cand in peered),
        patch.object(AZResolver, "_common_supported_letters", return_value=["a", "b", "c"]),
    ):
        candidates = resolver.get_region_fallback_candidates()

    regions = [region for region, _ in candidates]
    assert "us-east-1" not in regions
    assert "us-east-2" not in regions
    assert set(regions) == peered


def test_get_dc_fallback_candidates_requires_peering_with_every_staying_dc():
    """A candidate must be peered with ALL staying DCs, not just one."""
    params = _make_params(region_name="us-east-1 eu-west-1 us-west-2", availability_zone="a")
    resolver = AZResolver(params)

    peering = {
        ("us-east-1", "eu-west-2"): True,
        ("us-west-2", "eu-west-2"): True,
        ("us-east-1", "eu-west-3"): True,
        ("us-west-2", "eu-west-3"): False,
    }

    def is_region_peered(staying: str, candidate: str) -> bool:
        return peering.get((staying, candidate), False)

    with (
        patch.object(AZResolver, "_is_region_peered", side_effect=is_region_peered),
        patch.object(AZResolver, "_common_supported_letters", return_value=["a"]),
    ):
        regions = [region for region, _ in resolver.get_dc_fallback_candidates(dc_index=1)]

    assert "eu-west-2" in regions
    assert "eu-west-3" not in regions


def test_is_region_peered_true_only_for_active_connection():
    """_is_region_peered must accept only active peering, not pending-acceptance or missing."""
    region_pair = ("us-east-1", "eu-west-1")
    peer_name = "peer-name"

    with (
        patch("sdcm.provision.aws.az_resolver.AwsRegion"),
        patch("sdcm.utils.aws_peering.AwsVpcPeering._find_existing_peering") as mock_find,
    ):
        mock_find.return_value = ({"Status": {"Code": "active"}}, peer_name)
        assert AZResolver._is_region_peered(*region_pair) is True

        mock_find.return_value = ({"Status": {"Code": "pending-acceptance"}}, peer_name)
        assert AZResolver._is_region_peered(*region_pair) is False

        mock_find.return_value = (None, peer_name)
        assert AZResolver._is_region_peered(*region_pair) is False
