"""Tests for arise.distributed — setup and teardown of AWS resources."""

import json
import os
import re
import sys
import tempfile
from unittest.mock import MagicMock, patch, call

sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))

from arise.config import ARISEConfig
from arise.distributed import setup_distributed, destroy_distributed, _random_suffix


# ---------------------------------------------------------------------------
# _random_suffix helper
# ---------------------------------------------------------------------------

def test_random_suffix_length():
    suffix = _random_suffix()
    assert len(suffix) == 12


def test_random_suffix_hex():
    suffix = _random_suffix()
    assert re.fullmatch(r"[0-9a-f]{12}", suffix), f"Not a 12-char hex string: {suffix!r}"


def test_random_suffix_unique():
    suffixes = {_random_suffix() for _ in range(20)}
    # Extremely unlikely to collide even once
    assert len(suffixes) > 1


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------

_FAKE_SUFFIX = "abcdef123456"
_BUCKET = f"arise-skills-{_FAKE_SUFFIX}"
_QUEUE = f"arise-trajectories-{_FAKE_SUFFIX}"
_DLQ = f"{_QUEUE}-dlq"

_QUEUE_URL = f"https://sqs.us-west-2.amazonaws.com/123456789012/{_QUEUE}"
_DLQ_URL = f"https://sqs.us-west-2.amazonaws.com/123456789012/{_DLQ}"
_QUEUE_ARN = f"arn:aws:sqs:us-west-2:123456789012:{_QUEUE}"
_DLQ_ARN = f"arn:aws:sqs:us-west-2:123456789012:{_DLQ}"


def _make_mock_session(queue_name=_QUEUE):
    """Return a mock boto3.Session with S3 and SQS clients."""
    session = MagicMock()

    s3_client = MagicMock()

    dlq_name = f"{queue_name}-dlq"
    dlq_url = f"https://sqs.us-west-2.amazonaws.com/123456789012/{dlq_name}"
    dlq_arn = f"arn:aws:sqs:us-west-2:123456789012:{dlq_name}"
    queue_url = f"https://sqs.us-west-2.amazonaws.com/123456789012/{queue_name}"
    queue_arn = f"arn:aws:sqs:us-west-2:123456789012:{queue_name}"

    sqs = MagicMock()
    sqs.create_queue.side_effect = [
        {"QueueUrl": dlq_url},
        {"QueueUrl": queue_url},
    ]
    sqs.get_queue_attributes.side_effect = [
        {"Attributes": {"QueueArn": dlq_arn}},
        {"Attributes": {"QueueArn": queue_arn}},
    ]

    def _client(service, **kwargs):
        return {"s3": s3_client, "sqs": sqs}[service]

    session.client.side_effect = _client
    return session, s3_client, sqs


# ---------------------------------------------------------------------------
# setup_distributed
# ---------------------------------------------------------------------------

@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_setup_distributed_auto_names(mock_session_cls, mock_suffix):
    session, s3, sqs = _make_mock_session()
    mock_session_cls.return_value = session

    config = setup_distributed(region="us-west-2")

    assert config.s3_bucket == _BUCKET
    assert config.aws_region == "us-west-2"
    assert _QUEUE in config.sqs_queue_url

    # Bucket created with LocationConstraint for non-us-east-1
    s3.create_bucket.assert_called_once()
    create_args = s3.create_bucket.call_args
    assert create_args[1]["CreateBucketConfiguration"] == {"LocationConstraint": "us-west-2"}

    # Versioning enabled
    s3.put_bucket_versioning.assert_called_once_with(
        Bucket=_BUCKET,
        VersioningConfiguration={"Status": "Enabled"},
    )

    # DLQ created first, then main queue
    assert sqs.create_queue.call_count == 2
    dlq_call = sqs.create_queue.call_args_list[0]
    assert dlq_call == call(QueueName=f"{_QUEUE}-dlq")

    main_call = sqs.create_queue.call_args_list[1]
    assert main_call[1]["QueueName"] == _QUEUE
    redrive = json.loads(main_call[1]["Attributes"]["RedrivePolicy"])
    assert redrive["maxReceiveCount"] == "5"
    assert _DLQ in redrive["deadLetterTargetArn"]


@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_setup_distributed_custom_names(mock_session_cls, mock_suffix):
    session, s3, sqs = _make_mock_session(queue_name="my-queue")
    mock_session_cls.return_value = session

    config = setup_distributed(
        region="us-east-1",
        bucket_name="my-bucket",
        queue_name="my-queue",
    )

    assert config.s3_bucket == "my-bucket"
    s3.create_bucket.assert_called_once_with(Bucket="my-bucket")
    # us-east-1 should NOT include LocationConstraint
    create_args = s3.create_bucket.call_args
    assert "CreateBucketConfiguration" not in create_args[1]


@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_setup_distributed_passes_profile(mock_session_cls, mock_suffix):
    session, _, _ = _make_mock_session()
    mock_session_cls.return_value = session

    setup_distributed(region="eu-west-1", profile="staging")

    mock_session_cls.assert_called_once_with(region_name="eu-west-1", profile_name="staging")


@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_setup_distributed_returns_config(mock_session_cls, mock_suffix):
    session, s3, sqs = _make_mock_session()
    mock_session_cls.return_value = session

    config = setup_distributed(region="us-west-2")

    assert isinstance(config, ARISEConfig)
    assert config.s3_bucket is not None
    assert config.sqs_queue_url is not None
    assert config.aws_region == "us-west-2"


@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_setup_distributed_auto_name_format(mock_session_cls, mock_suffix):
    """Auto-generated names follow the arise-skills-{12-char} / arise-trajectories-{12-char} pattern."""
    session, s3, sqs = _make_mock_session()
    mock_session_cls.return_value = session

    config = setup_distributed(region="us-west-2")

    assert re.fullmatch(r"arise-skills-[0-9a-f]{12}", config.s3_bucket)
    # Extract queue name from URL
    queue_name = config.sqs_queue_url.rstrip("/").split("/")[-1]
    assert re.fullmatch(r"arise-trajectories-[0-9a-f]{12}", queue_name)


# ---------------------------------------------------------------------------
# destroy_distributed
# ---------------------------------------------------------------------------

@patch("arise.distributed.boto3.Session")
def test_destroy_distributed(mock_session_cls):
    session = MagicMock()
    s3_resource = MagicMock()
    bucket_obj = MagicMock()
    s3_resource.Bucket.return_value = bucket_obj

    sqs = MagicMock()
    sqs.get_queue_attributes.return_value = {
        "Attributes": {
            "RedrivePolicy": json.dumps({
                "deadLetterTargetArn": "arn:aws:sqs:us-west-2:123456789012:my-dlq",
                "maxReceiveCount": "5",
            })
        }
    }
    sqs.get_queue_url.return_value = {"QueueUrl": "https://sqs.us-west-2.amazonaws.com/123456789012/my-dlq"}

    def _client(service, **kwargs):
        return {"sqs": sqs}[service]

    session.client.side_effect = _client
    session.resource.return_value = s3_resource
    mock_session_cls.return_value = session

    config = ARISEConfig(
        s3_bucket="my-bucket",
        sqs_queue_url="https://sqs.us-west-2.amazonaws.com/123456789012/my-queue",
        aws_region="us-west-2",
    )

    destroy_distributed(config)

    # Bucket emptied (all versions deleted) and then bucket deleted
    bucket_obj.object_versions.all().delete.assert_called_once()
    bucket_obj.delete.assert_called_once()

    # DLQ and main queue both deleted
    assert sqs.delete_queue.call_count == 2


@patch("arise.distributed.boto3.Session")
def test_destroy_distributed_no_bucket(mock_session_cls):
    """destroy_distributed skips S3 when s3_bucket is None."""
    session = MagicMock()
    s3_resource = MagicMock()
    sqs = MagicMock()
    sqs.get_queue_attributes.return_value = {
        "Attributes": {
            "RedrivePolicy": json.dumps({
                "deadLetterTargetArn": "arn:aws:sqs:us-west-2:123456789012:dlq",
            })
        }
    }
    sqs.get_queue_url.return_value = {"QueueUrl": "https://sqs.us-west-2.amazonaws.com/123456789012/dlq"}

    def _client(service, **kwargs):
        return {"sqs": sqs}[service]

    session.client.side_effect = _client
    session.resource.return_value = s3_resource
    mock_session_cls.return_value = session

    config = ARISEConfig(
        s3_bucket=None,
        sqs_queue_url="https://sqs.us-west-2.amazonaws.com/123456789012/my-queue",
        aws_region="us-west-2",
    )

    destroy_distributed(config)

    s3_resource.Bucket.assert_not_called()
    assert sqs.delete_queue.call_count >= 1


@patch("arise.distributed.boto3.Session")
def test_destroy_distributed_no_queue(mock_session_cls):
    """destroy_distributed skips SQS when sqs_queue_url is None."""
    session = MagicMock()
    s3_resource = MagicMock()
    bucket_obj = MagicMock()
    s3_resource.Bucket.return_value = bucket_obj
    sqs = MagicMock()

    def _client(service, **kwargs):
        return {"sqs": sqs}[service]

    session.client.side_effect = _client
    session.resource.return_value = s3_resource
    mock_session_cls.return_value = session

    config = ARISEConfig(
        s3_bucket="my-bucket",
        sqs_queue_url=None,
        aws_region="us-west-2",
    )

    destroy_distributed(config)

    bucket_obj.delete.assert_called_once()
    sqs.delete_queue.assert_not_called()


# ---------------------------------------------------------------------------
# CLI integration
# ---------------------------------------------------------------------------

@patch("arise.distributed._random_suffix", return_value=_FAKE_SUFFIX)
@patch("arise.distributed.boto3.Session")
def test_cli_setup_distributed(mock_session_cls, mock_suffix):
    session, s3, sqs = _make_mock_session()
    mock_session_cls.return_value = session

    old_cwd = os.getcwd()
    tmpdir = tempfile.mkdtemp()
    try:
        os.chdir(tmpdir)

        from arise.cli import main
        sys.argv = ["arise", "setup-distributed", "--region", "us-west-2"]
        main()

        config_path = os.path.join(tmpdir, ".arise.json")
        assert os.path.exists(config_path)
        with open(config_path) as f:
            data = json.load(f)
        assert re.fullmatch(r"arise-skills-[0-9a-f]{12}", data["s3_bucket"])
        assert data["aws_region"] == "us-west-2"
        assert "sqs_queue_url" in data
    finally:
        os.chdir(old_cwd)


@patch("arise.distributed.boto3.Session")
def test_cli_setup_distributed_destroy(mock_session_cls):
    session = MagicMock()
    s3_resource = MagicMock()
    bucket_obj = MagicMock()
    s3_resource.Bucket.return_value = bucket_obj

    sqs = MagicMock()
    sqs.get_queue_attributes.return_value = {
        "Attributes": {
            "RedrivePolicy": json.dumps({
                "deadLetterTargetArn": "arn:aws:sqs:us-west-2:123456789012:dlq",
                "maxReceiveCount": "5",
            })
        }
    }
    sqs.get_queue_url.return_value = {"QueueUrl": "https://sqs.us-west-2.amazonaws.com/123456789012/dlq"}

    def _client(service, **kwargs):
        return {"sqs": sqs}[service]

    session.client.side_effect = _client
    session.resource.return_value = s3_resource
    mock_session_cls.return_value = session

    old_cwd = os.getcwd()
    tmpdir = tempfile.mkdtemp()
    try:
        os.chdir(tmpdir)

        # Write a fake .arise.json
        with open(".arise.json", "w") as f:
            json.dump({
                "s3_bucket": "my-bucket",
                "sqs_queue_url": "https://sqs.us-west-2.amazonaws.com/123456789012/my-queue",
                "aws_region": "us-west-2",
            }, f)

        from arise.cli import main
        sys.argv = ["arise", "setup-distributed", "--destroy"]
        main()

        assert not os.path.exists(".arise.json")
        bucket_obj.delete.assert_called_once()
    finally:
        os.chdir(old_cwd)


@patch("arise.distributed.boto3.Session")
def test_cli_setup_distributed_destroy_missing_config(mock_session_cls):
    """--destroy with no .arise.json exits with error."""
    old_cwd = os.getcwd()
    tmpdir = tempfile.mkdtemp()
    try:
        os.chdir(tmpdir)

        from arise.cli import main
        import pytest
        sys.argv = ["arise", "setup-distributed", "--destroy"]
        with pytest.raises(SystemExit) as exc:
            main()
        assert exc.value.code != 0
    finally:
        os.chdir(old_cwd)
