from __future__ import annotations

import re
from pathlib import Path

ROOT = Path(__file__).resolve().parents[3]

# One entry per browser-direct upload purpose. Kept in lock-step with
# ``app.core.storage_uploads.PURPOSE_MAX_BYTES`` by
# ``test_every_upload_purpose_has_an_iam_prefix`` below — a purpose without a
# prefix here is a 403 that only happens in production.
READ_WRITE_PREFIXES = {
    "inventory_evidence",
    "feed_post_image",
    "feed_announcement",
    "approval_supplement",
    "evidence_documents",
    "vendor_assessment_evidence",
    "self_check_evidence",
    "inbound_questionnaire",
    "bounty_attachment",
    "asset_disposal_evidence",
    "awareness_material",
    "ai_labs_image",
    "ai_labs_document",
    "ai_labs_audio",
    # Not an upload purpose — the backend writes generated artifacts here
    # with the EC2 role's own credentials and presigns the employee's
    # download from it. Listed so a `terraform apply` cannot revoke it.
    "ai_labs_output",
}
READ_ONLY_PREFIXES = {"cli-dist", "plugin-dist"}
CLEANUP_PREFIXES = READ_WRITE_PREFIXES - {"approval_supplement"}


def _read(path: str) -> str:
    return (ROOT / path).read_text()


def _braced_block(text: str, marker: str) -> str:
    marker_start = text.index(marker)
    start = text.index("{", marker_start)
    depth = 0
    for index in range(start, len(text)):
        if text[index] == "{":
            depth += 1
        elif text[index] == "}":
            depth -= 1
            if depth == 0:
                return text[start : index + 1]
    raise AssertionError(f"unterminated Terraform block: {marker}")


def _statement(policy: str, sid: str) -> str:
    match = re.search(rf'sid\s*=\s*"{re.escape(sid)}"', policy)
    assert match
    sid_start = match.start()
    statement_start = policy.rfind("statement {", 0, sid_start)
    assert statement_start >= 0
    return _braced_block(policy[statement_start:], "statement")


def _actions(statement: str) -> set[str]:
    match = re.search(r"actions\s*=\s*\[(.*?)\]", statement, re.DOTALL)
    assert match
    return set(re.findall(r'"([^"]+)"', match.group(1)))


def _prefixes(terraform: str, local_name: str) -> set[str]:
    match = re.search(
        rf"{local_name}\s*=\s*toset\(\[(.*?)\]\)", terraform, re.DOTALL
    )
    assert match
    return set(re.findall(r'"([^"]+)"', match.group(1)))


def test_uploads_policy_has_exact_object_actions_and_prefixes() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    policy = _braced_block(
        terraform, 'data "aws_iam_policy_document" "uploads_access"'
    )

    assert _prefixes(terraform, "uploads_read_write_prefixes") == READ_WRITE_PREFIXES
    assert _prefixes(terraform, "uploads_read_only_prefixes") == READ_ONLY_PREFIXES

    read_write = _statement(policy, "ApplicationObjectReadWrite")
    assert _actions(read_write) == {"s3:GetObject", "s3:PutObject"}
    assert "for prefix in local.uploads_read_write_prefixes" in read_write
    assert '"${var.uploads_bucket_arn}/${prefix}/*"' in read_write

    read_only = _statement(policy, "SelfServiceArtifactRead")
    assert _actions(read_only) == {"s3:GetObject"}
    assert "for prefix in local.uploads_read_only_prefixes" in read_only
    assert '"${var.uploads_bucket_arn}/${prefix}/*"' in read_only

    cleanup = _statement(policy, "OrphanUploadCleanup")
    assert _actions(cleanup) == {"s3:DeleteObject"}
    assert "for prefix in local.uploads_cleanup_prefixes" in cleanup
    assert '"${var.uploads_bucket_arn}/${prefix}/20??/*"' in cleanup
    assert "uploads_cleanup_prefixes = setsubtract(" in terraform
    assert 'toset(["approval_supplement"])' in terraform
    assert CLEANUP_PREFIXES == {
        "inventory_evidence",
        "feed_post_image",
        "feed_announcement",
        "evidence_documents",
        "vendor_assessment_evidence",
        "self_check_evidence",
        "inbound_questionnaire",
        "bounty_attachment",
        "asset_disposal_evidence",
        "awareness_material",
        "ai_labs_image",
        "ai_labs_document",
        "ai_labs_audio",
        # AI Labs deletes its own objects when an employee deletes a
        # conversation or a generation — that is a product promise, not an
        # orphan sweep, so DeleteObject on the output namespace is load-
        # bearing rather than incidental.
        "ai_labs_output",
    }


def test_list_bucket_is_bucket_scoped_and_limited_to_the_registered_prefixes() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    policy = _braced_block(
        terraform, 'data "aws_iam_policy_document" "uploads_access"'
    )
    existence = _statement(policy, "PermittedPrefixExistenceCheck")

    assert _actions(existence) == {"s3:ListBucket"}
    assert "resources = [var.uploads_bucket_arn]" in existence
    assert 'test     = "StringLike"' in existence
    assert 'variable = "s3:prefix"' in existence
    assert "for prefix in local.uploads_all_prefixes" in existence
    assert '"${prefix}/*"' in existence
    assert "setunion(" in terraform
    assert READ_WRITE_PREFIXES | READ_ONLY_PREFIXES == {
        "inventory_evidence",
        "feed_post_image",
        "feed_announcement",
        "approval_supplement",
        "evidence_documents",
        "vendor_assessment_evidence",
        "self_check_evidence",
        "inbound_questionnaire",
        "bounty_attachment",
        "asset_disposal_evidence",
        "awareness_material",
        "ai_labs_image",
        "ai_labs_document",
        "ai_labs_audio",
        "ai_labs_output",
        "cli-dist",
        "plugin-dist",
    }


def test_uploads_policy_has_no_broad_delete_multipart_kms_or_bucket_wide_grant() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    policy = _braced_block(
        terraform, 'data "aws_iam_policy_document" "uploads_access"'
    )

    forbidden = {
        "s3:AbortMultipartUpload",
        "s3:ListMultipartUploadParts",
        "s3:ListBucketMultipartUploads",
        "s3:*",
        "kms:Decrypt",
        "kms:GenerateDataKey",
        '"${var.uploads_bucket_arn}/*"',
    }
    for value in forbidden:
        assert value not in policy

    cleanup = _statement(policy, "OrphanUploadCleanup")
    assert '"${var.uploads_bucket_arn}/${prefix}/*"' not in cleanup
    assert "approval_supplement" not in cleanup


def test_ec2_role_policy_inventory_cannot_hide_an_extra_inline_policy() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    inline_policy_blocks = re.findall(
        r'resource\s+"aws_iam_role_policy"\s+"([^"]+)"\s*\{(.*?)\n\}',
        terraform,
        re.DOTALL,
    )
    policies_for_ec2 = {
        name
        for name, block in inline_policy_blocks
        if "role   = aws_iam_role.ec2.id" in block
    }
    assert policies_for_ec2 == {"ecr_secrets", "uploads_access"}
    assert terraform.count("role       = aws_iam_role.ec2.name") == 1
    assert "AmazonSSMManagedInstanceCore" in terraform


def test_live_inline_policy_name_and_prod_bucket_arn_are_stable() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    resource = _braced_block(
        terraform, 'resource "aws_iam_role_policy" "uploads_access"'
    )
    variables = _read("infra/aws/modules/ec2/variables.tf")
    prod = _read("infra/aws/environments/prod/main.tf")

    assert 'name   = "HelixProdUploadsAccess"' in resource
    assert "policy = data.aws_iam_policy_document.uploads_access.json" in resource
    assert 'variable "uploads_bucket_arn"' in variables
    assert "uploads_bucket_arn must be an S3 bucket ARN without an object suffix" in variables
    assert 'uploads_bucket_arn = "arn:aws:s3:::helix-prod-uploads-ap-southeast-1"' in prod


def test_head_object_403_404_risk_is_part_of_policy_and_ops_contract() -> None:
    terraform = _read("infra/aws/modules/ec2/main.tf")
    docs = _read("docs/ops/storage-direct-upload.md")

    assert "S3 returns 403 for a missing HeadObject when ListBucket is absent" in terraform
    assert "HeadObject missing-object contract" in docs
    assert "404" in docs
    assert "403" in docs
    assert "multipart-specific list/abort permissions" in docs


def test_every_upload_purpose_has_an_iam_prefix() -> None:
    """Every browser-direct upload purpose must have an S3 prefix grant.

    This is the test that would have caught the 2026-09-02 gap. The presigned
    PUT is signed with the prod EC2 role's credentials, so S3 checks THIS
    policy when the browser uploads. A purpose registered in
    ``PURPOSE_MAX_BYTES`` without a matching prefix here therefore fails with
    ``403 AccessDenied`` — **in production only**. Nothing local reproduces
    it: the presign endpoint happily signs a URL for a prefix the role cannot
    write, the failure lands in the browser's PUT, and every backend test
    passes because they never touch real S3.

    ``asset_disposal_evidence`` shipped that way (PR #3401): code deployed,
    migration applied, feature dead on arrival.
    """
    from app.core.storage_uploads import ALLOWED_PURPOSES

    missing = set(ALLOWED_PURPOSES) - READ_WRITE_PREFIXES
    assert missing == set(), (
        f"upload purpose(s) with no IAM prefix grant: {sorted(missing)} — "
        "add them to local.uploads_read_write_prefixes in "
        "infra/aws/modules/ec2/main.tf AND to the live inline policy "
        "helix-prod-ec2-role/HelixProdUploadsAccess, or the browser PUT 403s "
        "in production only."
    )


def test_terraform_prefixes_match_the_contract_sets() -> None:
    """The Terraform locals and this file's sets are the same list.

    Two hand-maintained copies is how ``bounty_attachment`` ended up in the
    LIVE policy but not in Terraform — leaving a ``terraform apply`` of this
    module primed to silently revoke it. Reconciled 2026-09-02; this test
    keeps them from drifting again.
    """
    terraform = _read("infra/aws/modules/ec2/main.tf")
    assert _prefixes(terraform, "uploads_read_write_prefixes") == READ_WRITE_PREFIXES
    assert _prefixes(terraform, "uploads_read_only_prefixes") == READ_ONLY_PREFIXES