"""Unit tests for EMR cluster provisioner using moto for AWS API mocking."""

from __future__ import annotations

import os
from pathlib import Path
from unittest.mock import MagicMock, patch

import boto3
import pytest

os.environ["MOTO_AMIS_PATH"] = str(Path(__file__).parent.parent / "test_data" / "mocked_ami_data.json")
from moto.server import ThreadedMotoServer

from sdcm.provision.aws.emr_provisioner import (
    EMR_EC2_INSTANCE_PROFILE_NAME,
    EMR_SERVICE_ROLE_NAME,
    EmrClusterProvisioner,
    EmrClusterTerminatedError,
    ensure_emr_roles,
    list_emr_clusters,
)

AWS_REGION = "us-east-1"


@pytest.fixture(scope="module")
def moto_server():
    """Run a mocked AWS server for testing."""
    server = ThreadedMotoServer(port=0)
    server.start()
    host, port = server.get_host_and_port()

    env_vars = {
        "AWS_ENDPOINT_URL": f"http://{host}:{port}",
        "AWS_ACCESS_KEY_ID": "testing",
        "AWS_SECRET_ACCESS_KEY": "testing",
        "AWS_DEFAULT_REGION": AWS_REGION,
    }
    saved_env = {key: os.environ.get(key) for key in env_vars}
    os.environ.update(env_vars)

    yield env_vars["AWS_ENDPOINT_URL"]

    for key, original in saved_env.items():
        if original is None:
            os.environ.pop(key, None)
        else:
            os.environ[key] = original
    server.stop()


@pytest.fixture(scope="module")
def iam_roles(moto_server):
    """Create required IAM roles for EMR."""
    iam = boto3.client("iam", region_name=AWS_REGION)
    trust_policy = '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"elasticmapreduce.amazonaws.com"},"Action":"sts:AssumeRole"}]}'
    ec2_trust_policy = '{"Version":"2012-10-17","Statement":[{"Effect":"Allow","Principal":{"Service":"ec2.amazonaws.com"},"Action":"sts:AssumeRole"}]}'

    iam.create_role(RoleName="EMR_DefaultRole", AssumeRolePolicyDocument=trust_policy)
    iam.create_role(RoleName="EMR_EC2_DefaultRole", AssumeRolePolicyDocument=ec2_trust_policy)
    iam.create_instance_profile(InstanceProfileName="EMR_EC2_DefaultRole")
    iam.add_role_to_instance_profile(InstanceProfileName="EMR_EC2_DefaultRole", RoleName="EMR_EC2_DefaultRole")


class MockParams:
    """Mock params object that behaves like SCTConfiguration for testing."""

    def __init__(self, **kwargs):
        self._data = {
            "emr_release_label": "emr-7.12.0",
            "emr_instance_type_master": "m5.xlarge",
            "emr_instance_type_core": "m5.xlarge",
            "emr_instance_count_core": 2,
            "emr_instance_type_task": "",
            "emr_instance_count_task": 0,
            "emr_spot_bid_percentage": 100,
            "emr_applications": ["Spark"],
            "emr_keep_alive": True,
            "emr_log_uri": "",
        }
        self._data.update(kwargs)

    def get(self, key, default=None):
        return self._data.get(key, default)


@pytest.fixture()
def emr_params():
    return MockParams()


@pytest.fixture()
def provisioner(moto_server, iam_roles, emr_params):
    return EmrClusterProvisioner(region_name=AWS_REGION, params=emr_params)


def test_create_emr_cluster(provisioner):
    """Test creating an EMR cluster."""
    cluster_id = provisioner.create_emr_cluster(
        test_id="test-1234-5678",
        user="test_user",
    )
    assert cluster_id is not None
    assert cluster_id.startswith("j-")
    assert provisioner.cluster_id == cluster_id


def test_create_emr_cluster_has_required_tags(provisioner):
    """Test that EMR cluster creation includes all SCP-required tags."""
    provisioner.create_emr_cluster(
        test_id="test-tags-check",
        user="test_user",
        test_name="my-test",
        version="5.4.0",
    )
    tags = {t["Key"]: t["Value"] for t in provisioner.get_emr_cluster_description().get("Tags", [])}

    expected_tags = {
        "CreatedBy": "SCT",
        "RunByUser": "test_user",
        "TestName": "my-test",
        "version": "5.4.0",
        "Name": "sct-emr-test-tag",
        "NodeType": "emr",
        "keep_action": "terminate",
    }
    assert expected_tags.items() <= tags.items()


def test_create_emr_cluster_has_ssh_key(provisioner):
    """Test that EMR cluster creation includes the SSH key pair."""
    provisioner.create_emr_cluster(test_id="test-ssh-key", user="test_user")
    ec2_attrs = provisioner.get_emr_cluster_description().get("Ec2InstanceAttributes", {})
    assert ec2_attrs.get("Ec2KeyName") == "scylla_test_id_ed25519"


def test_create_emr_cluster_no_sct_security_groups(provisioner):
    """Test that EMR cluster creation does not pass SCT security groups."""
    provisioner.create_emr_cluster(test_id="test-no-sct-sg", user="test_user")
    ec2_attrs = provisioner.get_emr_cluster_description().get("Ec2InstanceAttributes", {})
    # EMR should use its own managed SGs, not SCT's
    assert ec2_attrs.get("EmrManagedMasterSecurityGroup", "") != "sg-0691c61126d99dd41"  # SCT-2-sg


def test_get_emr_cluster_status(provisioner):
    """Test getting EMR cluster status after creation."""
    provisioner.create_emr_cluster(
        test_id="test-status-check",
        user="test_user",
    )
    status = provisioner.get_emr_cluster_status()
    assert "State" in status
    assert status["State"] in ("STARTING", "BOOTSTRAPPING", "RUNNING", "WAITING")


def test_get_emr_cluster_description(provisioner):
    """Test getting full EMR cluster description."""
    provisioner.create_emr_cluster(
        test_id="test-describe",
        user="test_user",
    )
    desc = provisioner.get_emr_cluster_description()
    assert desc["Id"] == provisioner.cluster_id
    assert "Status" in desc


def test_create_emr_cluster_default_log_uri_falls_back_to_shared_bucket(provisioner):
    """Empty/unset emr_log_uri defaults to s3://sct-emr-spark-migrator-<region>/emr-logs/
    so EMR step stdout/stderr are uploaded and SparkMigratorRunner.get_step_stdout works."""
    provisioner.create_emr_cluster(test_id="test-default-loguri", user="test_user")
    desc = provisioner.get_emr_cluster_description()
    assert desc.get("LogUri") == f"s3://sct-emr-spark-migrator-{AWS_REGION}/emr-logs/"


def test_create_emr_cluster_explicit_log_uri_overrides_default():
    """Explicit emr_log_uri param wins over the default."""
    prov = EmrClusterProvisioner(region_name=AWS_REGION, params=MockParams(emr_log_uri="s3://my-custom-emr-logs/path/"))
    prov.create_emr_cluster(test_id="test-explicit-loguri", user="test_user")
    assert prov.get_emr_cluster_description().get("LogUri") == "s3://my-custom-emr-logs/path/"


def test_terminate_emr_cluster(provisioner):
    """Test terminating an EMR cluster."""
    provisioner.create_emr_cluster(
        test_id="test-terminate",
        user="test_user",
    )
    cluster_id = provisioner.cluster_id
    provisioner.terminate_emr_cluster()

    status = provisioner.get_emr_cluster_status(cluster_id)
    assert status["State"] in ("TERMINATING", "TERMINATED")


def test_list_emr_clusters_by_tags(moto_server, iam_roles):
    """Test listing EMR clusters filtered by tags."""
    params = MockParams()
    prov = EmrClusterProvisioner(region_name=AWS_REGION, params=params)

    prov.create_emr_cluster(test_id="list-test-001", user="user1")
    prov.create_emr_cluster(test_id="list-test-002", user="user1")

    clusters = list_emr_clusters({"RunByUser": "user1"}, region_name=AWS_REGION)
    assert len(clusters) >= 2

    clusters_by_test = list_emr_clusters({"TestId": "list-test-001"}, region_name=AWS_REGION)
    assert len(clusters_by_test) >= 1
    assert all(c["Tags"].get("TestId") == "list-test-001" for c in clusters_by_test)


def test_list_emr_clusters_uses_fail_fast_client_config():
    """The EMR client must use fail-fast network settings for unreachable regions."""
    emr_client = MagicMock()
    emr_client.get_paginator.return_value.paginate.return_value = []

    with patch("sdcm.provision.aws.emr_provisioner.boto3.client", return_value=emr_client) as mock_client:
        list_emr_clusters({"RunByUser": "user1"}, region_name="me-south-1")

    config = mock_client.call_args.kwargs["config"]
    assert config.connect_timeout == 10
    assert config.retries == {"total_max_attempts": 1}


def test_build_instance_groups_default(provisioner):
    """Test default instance groups configuration."""
    groups = provisioner._build_instance_groups()
    assert len(groups) == 2  # Master + Core

    master = groups[0]
    assert master["Name"] == "Master"
    assert master["InstanceRole"] == "MASTER"
    assert master["Market"] == "ON_DEMAND"
    assert master["InstanceCount"] == 1

    core = groups[1]
    assert core["Name"] == "Core"
    assert core["InstanceRole"] == "CORE"
    assert core["InstanceCount"] == 2


def test_build_instance_groups_with_task():
    """Test instance groups with task nodes (Spot)."""
    params = MockParams(
        emr_instance_type_task="m5.xlarge",
        emr_instance_count_task=3,
        emr_spot_bid_percentage=80,
    )
    prov = EmrClusterProvisioner(region_name=AWS_REGION, params=params)
    groups = prov._build_instance_groups()
    assert len(groups) == 3

    task = groups[2]
    assert task["Name"] == "Task"
    assert task["InstanceRole"] == "TASK"
    assert task["Market"] == "SPOT"
    assert task["InstanceCount"] == 3
    assert task["BidPrice"] == "80"


def test_add_step(provisioner):
    """Test that add_step adds a generic step to the EMR cluster."""
    provisioner.create_emr_cluster(test_id="test-step", user="test_user")
    step_id = provisioner.add_step(
        cluster_id=provisioner.cluster_id,
        step_name="test-step",
        jar="command-runner.jar",
        args=["spark-submit", "--class", "com.example.Main", "s3://bucket/app.jar"],
    )
    assert step_id is not None


def test_get_step_status(provisioner):
    """Test getting step status."""
    provisioner.create_emr_cluster(
        test_id="test-step-status",
        user="test_user",
    )
    step_id = provisioner.add_step(
        cluster_id=provisioner.cluster_id,
        step_name="test-step-status",
        jar="command-runner.jar",
        args=["spark-submit", "s3://bucket/test.jar"],
    )
    status = provisioner.get_step_status(provisioner.cluster_id, step_id)
    assert "State" in status


def test_ensure_emr_roles_creates_missing_roles(moto_server):
    """Test that ensure_emr_roles creates both roles when they don't exist."""
    iam = boto3.client("iam", region_name=AWS_REGION)
    iam.attach_role_policy = lambda **kw: None

    with patch("sdcm.provision.aws.emr_provisioner.boto3.client", return_value=iam):
        ensure_emr_roles(AWS_REGION)

    role = iam.get_role(RoleName=EMR_SERVICE_ROLE_NAME)
    assert role["Role"]["RoleName"] == EMR_SERVICE_ROLE_NAME
    profile = iam.get_instance_profile(InstanceProfileName=EMR_EC2_INSTANCE_PROFILE_NAME)
    assert profile["InstanceProfile"]["InstanceProfileName"] == EMR_EC2_INSTANCE_PROFILE_NAME


def test_ensure_emr_roles_idempotent(moto_server, iam_roles):
    """Test that ensure_emr_roles is safe to call when roles already exist."""
    iam = boto3.client("iam", region_name=AWS_REGION)
    iam.attach_role_policy = lambda **kw: None

    with patch("sdcm.provision.aws.emr_provisioner.boto3.client", return_value=iam):
        ensure_emr_roles(AWS_REGION)

    assert iam.get_role(RoleName=EMR_SERVICE_ROLE_NAME)["Role"]["RoleName"] == EMR_SERVICE_ROLE_NAME


def test_wait_for_emr_cluster_ready_fails_fast_on_terminal_state(provisioner):
    """A terminal TERMINATED_WITH_ERRORS state raises immediately — no retry, no ~30 min wait in SetUp."""
    cluster_id = "j-CLUSTER"
    terminal_message = "bootstrap action 1 returned a non-zero return code"
    terminal_status = {
        "State": "TERMINATED_WITH_ERRORS",
        "StateChangeReason": {"Message": terminal_message},
    }

    with (
        patch.object(provisioner, "get_emr_cluster_status", return_value=terminal_status) as mock_status,
        patch("sdcm.utils.decorators.time.sleep") as mock_sleep,
        pytest.raises(EmrClusterTerminatedError, match=terminal_message),
    ):
        provisioner.wait_for_emr_cluster_ready(cluster_id)

    assert mock_status.call_count == 1
    mock_sleep.assert_not_called()


def test_wait_for_emr_cluster_ready_retries_while_provisioning(provisioner):
    """A still-provisioning state is retried until the cluster reaches WAITING."""
    cluster_id = "j-CLUSTER"
    waiting_status = {"State": "WAITING"}
    states = [{"State": "STARTING"}, {"State": "BOOTSTRAPPING"}, waiting_status]

    with (
        patch.object(provisioner, "get_emr_cluster_status", side_effect=states) as mock_status,
        patch("sdcm.utils.decorators.time.sleep"),
    ):
        result = provisioner.wait_for_emr_cluster_ready(cluster_id)

    assert result == waiting_status
    assert mock_status.call_count == 3
