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

import pandas as pd


def extract_arn_service(arn_pattern: str) -> str:
    # arn:${Partition}:SERVICE:${Region}:${Account}:resource
    parts = arn_pattern.split(":")
    if len(parts) >= 3:
        return parts[2]
    return ""


ARN_SERVICE_OVERRIDES = {
    "arn:${Partition}:${Vendor}:${Region}:*:${ResourceType}:${RecoveryPointId}": "backup",
}

EXCLUDED_ARNS = {
    "arn:${Partition}:rtbfabric:${Region}:${Account}:gateway/${GatewayId}/link/${LinkId}",
    "arn:${Partition}:rtbfabric:${Region}:${Account}:gateway/${GatewayId}",
    "arn:${Partition}:aws-marketplace::${Account}:${Catalog}/ReportingData/${FactTable}/Dashboard/${DashboardName}",
    "arn:${Partition}:iot:${Region}:${Account}:thinggroup/${ThingGroupName}",
    "arn:${Partition}:mediapackagev2:${Region}:${Account}:channelGroup/${ChannelGroupName}/channel/${ChannelName}",
    "arn:${Partition}:mediapackagev2:${Region}:${Account}:channelGroup/${ChannelGroupName}/channel/${ChannelName}/originEndpoint/${OriginEndpointName}",
}


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

    # Extract service from ARN pattern
    df["arn_service"] = df["arn_pattern"].apply(extract_arn_service)

    # Apply manual overrides
    for arn, service in ARN_SERVICE_OVERRIDES.items():
        df.loc[df["arn_pattern"] == arn, "arn_service"] = service

    # For ARNs that appear 2+ times, keep only where arn_service == service_prefix
    # If no match, keep first (arn_service is the authoritative service)
    results = []
    for arn, group in df.groupby("arn_pattern"):
        if len(group) == 1:
            results.append(group.iloc[0])
        else:
            matches = group[group["arn_service"] == group["service_prefix"]]
            if matches.empty:
                # No match - keep first, arn_service is authoritative
                results.append(group.iloc[0])
            else:
                for _, row in matches.iterrows():
                    results.append(row)

    output = pd.DataFrame(results).drop_duplicates()

    # Remove excluded ARNs (have multiple resource types)
    output = output[~output["arn_pattern"].isin(EXCLUDED_ARNS)]

    # Verify no ARNs were lost (except excluded)
    missing = set(df["arn_pattern"]) - set(output["arn_pattern"]) - EXCLUDED_ARNS
    if missing:
        raise ValueError(f"ARNs with no matches:\n{missing}")

    # Reorder and sort
    output = output[["arn_pattern", "arn_service", "service_prefix", "resource_type"]]
    output = output.sort_values("arn_pattern")

    output.to_parquet("arn_index.parquet", index=False)
    output.to_excel("arn_index.xlsx", index=False)
    print(f"Wrote {len(output)} resources to arn_index.parquet and arn_index.xlsx")


if __name__ == "__main__":
    main()
