import copy
import importlib.util
import json
import io
import sys
import types
import unittest
from contextlib import redirect_stdout
from pathlib import Path


REPO_ROOT = Path(__file__).resolve().parents[1]
SCRIPT = REPO_ROOT / ".github" / "scripts" / "update_ecr_policy.py"

sys.modules.setdefault(
    "boto3",
    types.SimpleNamespace(client=lambda *args, **kwargs: None, Session=lambda *args, **kwargs: None),
)
spec = importlib.util.spec_from_file_location("update_ecr_policy", SCRIPT)
update_ecr_policy = importlib.util.module_from_spec(spec)
spec.loader.exec_module(update_ecr_policy)


class FakeECRClient:
    class exceptions:
        class RepositoryPolicyNotFoundException(Exception):
            pass

        class RepositoryNotFoundException(Exception):
            pass

    def __init__(self, policy):
        self.policy = policy
        self.updated_policies = []

    def get_repository_policy(self, repositoryName):
        return {"policyText": json.dumps(self.policy)}

    def set_repository_policy(self, repositoryName, policyText):
        self.updated_policies.append(json.loads(policyText))


class UpdateECRPolicyTest(unittest.TestCase):
    def test_upsert_preserves_unmanaged_policy_statements(self):
        unmanaged_service_statement = {
            "Sid": "AllowECRReplicationService",
            "Effect": "Allow",
            "Principal": {"Service": "replication.ecr.amazonaws.com"},
            "Action": ["ecr:ReplicateImage"],
        }
        unmanaged_invalid_aws_statement = {
            "Sid": "ManuallyOwnedStatement",
            "Effect": "Allow",
            "Principal": {"AWS": "not-an-iam-arn"},
            "Action": ["ecr:BatchGetImage"],
        }
        managed_statement = {
            "Sid": "BYOCPull",
            "Effect": "Allow",
            "Principal": {
                "AWS": [
                    "arn:aws:iam::111111111111:role/old-role",
                    "not-an-iam-arn",
                ]
            },
            "Action": ["ecr:BatchGetImage"],
        }
        policy = {
            "Version": "2012-10-17",
            "Statement": [
                copy.deepcopy(unmanaged_service_statement),
                copy.deepcopy(unmanaged_invalid_aws_statement),
                managed_statement,
            ],
        }
        client = FakeECRClient(policy)

        with redirect_stdout(io.StringIO()):
            update_ecr_policy.upsert_repo_policy(
                ecr_client=client,
                region="us-west-2",
                repo="tidbcloud/example",
                sid="BYOCPull",
                principals=["arn:aws:iam::222222222222:role/customer-node-role"],
                permissions=["ecr:BatchGetImage", "ecr:GetDownloadUrlForLayer"],
                action="sync",
            )

        self.assertEqual(len(client.updated_policies), 1)
        updated_policy = client.updated_policies[0]
        statements_by_sid = {statement["Sid"]: statement for statement in updated_policy["Statement"]}
        self.assertEqual(statements_by_sid["AllowECRReplicationService"], unmanaged_service_statement)
        self.assertEqual(statements_by_sid["ManuallyOwnedStatement"], unmanaged_invalid_aws_statement)
        self.assertEqual(
            statements_by_sid["BYOCPull"]["Principal"],
            {"AWS": ["arn:aws:iam::222222222222:role/customer-node-role"]},
        )
        self.assertEqual(
            statements_by_sid["BYOCPull"]["Action"],
            ["ecr:BatchGetImage", "ecr:GetDownloadUrlForLayer"],
        )


if __name__ == "__main__":
    unittest.main()