#!/usr/bin/env python3
# /// script
# requires-python = ">=3.10"
# dependencies = ["pandas", "pyarrow"]
# ///

import re
import pandas as pd


# Type aliases: map AWS doc type -> list of known aliases (including Resource Explorer names)
TYPE_ALIASES = {
    # elasticloadbalancing: trailing slash differences
    ("elasticloadbalancing", "loadbalancer/app/"): ["loadbalancer/app"],
    ("elasticloadbalancing", "loadbalancer/net/"): ["loadbalancer/net"],
    ("elasticloadbalancing", "loadbalancer/gwy/"): ["loadbalancer/gwy"],
    # events: more specific name in docs
    ("events", "rule-on-default-event-bus"): ["rule"],
    # secretsmanager: capitalization
    ("secretsmanager", "Secret"): ["secret"],
    # mq: plural vs singular
    ("mq", "configurations"): ["configuration"],
    # inspector: slash vs hyphen
    ("inspector", "target-template"): ["target/template"],
    # backup: camelCase vs kebab-case
    ("backup", "backupPlan"): ["backup-plan"],
    ("backup", "backupVault"): ["backup-vault"],
    # ssm: hyphen removal
    ("ssm", "resourcedatasync"): ["resource-data-sync"],
    # s3: different naming
    ("s3", "storagelensconfiguration"): ["storage-lens"],
    # dms: abbreviated vs full name
    ("dms", "ReplicationSubnetGroup"): ["subgrp"],
}


def pattern_to_regex(arn_pattern: str) -> str:
    """Convert ARN pattern to regex with named capture groups."""
    # Replace ${VarName} with placeholder
    placeholders = []
    def capture_var(m):
        placeholders.append(m.group(1))
        return f"\x00{len(placeholders) - 1}\x00"
    result = re.sub(r"\$\{([^}]+)\}", capture_var, arn_pattern)

    # Replace * with placeholder
    result = result.replace("*", "\x01")

    # Escape remaining regex special chars
    result = re.escape(result)

    # Unescape hyphens (only special inside character classes)
    result = result.replace("\\-", "-")

    # Restore capture groups
    for i, name in enumerate(placeholders):
        result = result.replace(f"\x00{i}\x00", f"(?P<{name}>.+?)")

    # Restore wildcards
    result = result.replace("\x01", ".*")

    return f"^{result}$"


def get_type_names(service: str, resource_type: str) -> list[str]:
    """Get list of type names: primary type + any aliases."""
    types = [resource_type]
    aliases = TYPE_ALIASES.get((service, resource_type), [])
    types.extend(aliases)
    return types


def main():
    df = pd.read_parquet("arn_matcher.parquet")

    # Manual fixes: remove patterns that don't match real-world ARNs
    # EC2 image/snapshot always have account in modern AWS
    df = df[~(
        (df["arn_service"] == "ec2") &
        (df["arn_pattern"].str.contains("::image/|::snapshot/", regex=True))
    )]

    # DMS patterns use wildcards instead of capture groups - replace all
    df = df[df["arn_service"] != "dms"]

    # amplifybackend patterns use wildcards - replace all
    df = df[df["arn_service"] != "amplifybackend"]

    # artifact: only agreement and customer-agreement have wildcards
    df = df[~(
        (df["arn_service"] == "artifact") &
        (df["resource_type"].isin(["agreement", "customer-agreement"]))
    )]

    # backup: recoveryPoint has weird cross-service pattern with wildcard - drop it
    df = df[~(
        (df["arn_service"] == "backup") &
        (df["resource_type"] == "recoveryPoint")
    )]

    # connect: wildcard-* types are IAM permission wildcards, not real resources - drop
    df = df[~(
        (df["arn_service"] == "connect") &
        (df["resource_type"].str.startswith("wildcard-"))
    )]

    # health: only pattern has wildcards - replace all
    df = df[df["arn_service"] != "health"]

    # identitystore: All* types are IAM permission wildcards - drop
    df = df[~(
        (df["arn_service"] == "identitystore") &
        (df["resource_type"].str.startswith("All"))
    )]

    # imagebuilder: all*BuildVersions are "all versions" wildcards - drop
    df = df[~(
        (df["arn_service"] == "imagebuilder") &
        (df["resource_type"].str.startswith("all"))
    )]

    # mobiletargeting: apps and recommenders are plural wildcards - drop
    df = df[~(
        (df["arn_service"] == "mobiletargeting") &
        (df["resource_type"].isin(["apps", "recommenders"]))
    )]

    # neptune-db: only pattern has wildcard - replace
    df = df[df["arn_service"] != "neptune-db"]

    # Manual fixes: add patterns missing from AWS documentation
    manual_patterns = [
        # amplifybackend: wildcards replaced with capture groups (ordered most specific first)
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/api/${ApiId}", "api"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/auth/${AuthId}", "auth"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/challenge/${ChallengeId}", "token"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/config/${ConfigId}", "config"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/environments/${EnvironmentId}", "environment"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/job/${JobId}", "job"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/storage/${StorageId}", "storage"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${AppId}/${SubResourceId}", "backend"),
        ("amplifybackend", "arn:${Partition}:amplifybackend:${Region}:${Account}:/backend/${BackendId}", "created-backend"),
        # artifact: wildcards replaced with capture groups
        ("artifact", "arn:${Partition}:artifact:::agreement/${AgreementId}", "agreement"),
        ("artifact", "arn:${Partition}:artifact::${Account}:customer-agreement/${CustomerAgreementId}", "customer-agreement"),
        # health: wildcards in region and event ID, account may be empty
        ("health", "arn:${Partition}:health:${Region}:${Account}:event/${Service}/${EventTypeCode}/${EventId}", "event"),
        # neptune-db: wildcard replaced with capture group
        ("neptune-db", "arn:${Partition}:neptune-db:${Region}:${Account}:${ClusterResourceId}/${DatabaseId}", "database"),
        # EC2 with account (modern format)
        ("ec2", "arn:${Partition}:ec2:${Region}:${Account}:image/${ImageId}", "image"),
        ("ec2", "arn:${Partition}:ec2:${Region}:${Account}:snapshot/${SnapshotId}", "snapshot"),
        # EKS Kubernetes resources (from Resource Explorer)
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:deployment/${ClusterName}/${Namespace}/${DeploymentName}/${UUID}", "deployment"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:replicaset/${ClusterName}/${Namespace}/${ReplicaSetName}/${UUID}", "replicaset"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:service/${ClusterName}/${Namespace}/${ServiceName}/${UUID}", "service"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:endpointslice/${ClusterName}/${Namespace}/${EndpointSliceName}/${UUID}", "endpointslice"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:namespace/${ClusterName}/${NamespaceName}/${UUID}", "namespace"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:ingress/${ClusterName}/${Namespace}/${IngressName}/${UUID}", "ingress"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:statefulset/${ClusterName}/${Namespace}/${StatefulSetName}/${UUID}", "statefulset"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:persistentvolume/${ClusterName}/${PersistentVolumeName}/${UUID}", "persistentvolume"),
        ("eks", "arn:${Partition}:eks:${Region}:${Account}:daemonset/${ClusterName}/${Namespace}/${DaemonSetName}/${UUID}", "daemonset"),
        # Inspector (legacy)
        ("inspector", "arn:${Partition}:inspector:${Region}:${Account}:target/${TargetId}/template/${TemplateId}", "target-template"),
        # DMS: AWS docs use wildcards, add proper capture groups
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:assessment-run:${AssessmentRunId}", "ReplicationTaskAssessmentRun"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:cert:${CertificateId}", "Certificate"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:data-migration:${DataMigrationId}", "DataMigration"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:data-provider:${DataProviderId}", "DataProvider"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:endpoint:${EndpointId}", "Endpoint"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:es:${EventSubscriptionId}", "EventSubscription"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:individual-assessment:${IndividualAssessmentId}", "ReplicationTaskIndividualAssessment"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:instance-profile:${InstanceProfileId}", "InstanceProfile"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:migration-project:${MigrationProjectId}", "MigrationProject"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:rep:${ReplicationInstanceId}", "ReplicationInstance"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:replication-config:${ReplicationConfigId}", "ReplicationConfig"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:task:${TaskId}", "ReplicationTask"),
        ("dms", "arn:${Partition}:dms:${Region}:${Account}:subgrp:${SubnetGroupName}", "ReplicationSubnetGroup"),
    ]
    manual_df = pd.DataFrame(manual_patterns, columns=["arn_service", "arn_pattern", "resource_type"])
    df = pd.concat([df, manual_df], ignore_index=True)

    # Group by service, maintain order
    services = {}
    for _, row in df.iterrows():
        service = row["arn_service"]
        if service not in services:
            services[service] = []
        regex = pattern_to_regex(row["arn_pattern"])
        type_names = get_type_names(service, row["resource_type"])
        services[service].append((regex, type_names))

    # Write Python file
    with open("arn_patterns.py", "w") as f:
        f.write("# Auto-generated ARN patterns for matching\n")
        f.write("# Patterns are ordered: most specific first\n")
        f.write("# Type names: first is canonical (AWS docs), rest are aliases (Resource Explorer, etc.)\n")
        f.write("import re\n\n")
        f.write("ARN_PATTERNS: dict[str, list[tuple[re.Pattern, list[str]]]] = {\n")

        for service, patterns in services.items():
            f.write(f"    {service!r}: [\n")
            for regex, type_names in patterns:
                f.write(f"        (re.compile(r{regex!r}), {type_names!r}),\n")
            f.write("    ],\n")

        f.write("}\n")

    print(f"Wrote {len(df)} patterns for {len(services)} services to arn_patterns.py")


if __name__ == "__main__":
    main()
