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

import pandas as pd


def parse_arn_parts(arn_pattern: str) -> dict:
    """Parse ARN into its 6 standard parts."""
    # arn:partition:service:region:account-id:resource-id
    arn, partition, service, region, account, resource = arn_pattern.split(":", 5)
    return {
        "arn": arn,
        "partition": partition,
        "service": service,
        "region": region,
        "account": account,
        "resource": resource,
    }


def parse_segments(resource: str) -> dict:
    """Split resource into segments by /, : and then by variables"""
    import re
    var_pattern = re.compile(r'(\$\{[^}]+\})')

    # First split by / and :
    parts = [s1 for s0 in resource.split('/') for s1 in s0.split(':')]

    # Then split each part by variables
    segments = []
    for part in parts:
        splits = var_pattern.split(part)
        segments.extend([s for s in splits if s])

    # Each segment must have at most one ${ or *
    for seg in segments:
        var_count = seg.count("${") + seg.count("*")
        assert var_count <= 1, f"Segment '{seg}' has {var_count} variables/wildcards"

    return {f"seg_{i}": seg for i, seg in enumerate(segments)}


def normalize_for_sort(value: str) -> str:
    """Replace variables and wildcards so they sort after literals."""
    import re
    # Replace ${...} with ~ (comes after alphanumeric)
    value = re.sub(r"\$\{[^}]+\}", "~", value)
    # Replace * with ~~ (comes after ~, most greedy)
    value = value.replace("*", "~~")
    return value


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

    # Parse ARN parts
    parsed = df["arn_pattern"].apply(parse_arn_parts).apply(pd.Series)
    df = pd.concat([df, parsed], axis=1)

    # Parse resource segments
    segments = df["resource"].apply(parse_segments).apply(pd.Series)
    df = pd.concat([df, segments], axis=1)

    # Count non-empty segments (more segments = more specific = sort first)
    seg_cols = [c for c in df.columns if c.startswith("seg_")]
    df["_seg_count"] = df[seg_cols].notna().sum(axis=1)

    # Create sort keys with normalized values
    for col in ["arn_service", "region", "account"] + seg_cols:
        df[f"_sort_{col}"] = df[col].fillna("").apply(normalize_for_sort)

    # Sort: service, region, account, then segment count DESC, then segments
    sort_cols = [f"_sort_{c}" for c in ["arn_service", "region", "account"]] + seg_cols
    df = df.sort_values(["_sort_arn_service", "_sort_region", "_sort_account", "_seg_count"] + [f"_sort_{c}" for c in seg_cols], ascending=[True, True, True, False] + [True] * len(seg_cols))

    # Drop sort columns
    df = df.drop(columns=[c for c in df.columns if c.startswith("_sort_") or c == "_seg_count"])

    # Output columns
    base_cols = ["arn_pattern", "arn_service", "resource_type", "partition", "service", "region", "account", "resource"]
    output = df[base_cols + seg_cols]

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


if __name__ == "__main__":
    main()
