"""Tag fetching with service-specific API support."""

from collections import defaultdict
from typing import Callable

import boto3
from rich.progress import track

Getter = Callable[[list[str]], dict[str, dict]]


def get_using_tagging_api(arns: list[str], batch_size: int = 100) -> dict[str, dict]:
    """Fetch tags via ResourceGroupsTaggingAPI (batch of 100)."""
    if not arns:
        return {}

    client = boto3.client("resourcegroupstaggingapi")
    tags = {}
    batches = list(range(0, len(arns), batch_size))

    for start in track(batches, description="Fetching tags (tagging API)..."):
        batch = arns[start : start + batch_size]
        response = client.get_resources(ResourceARNList=batch)
        for r in response["ResourceTagMappingList"]:
            tags[r["ResourceARN"]] = {t["Key"]: t["Value"] for t in r.get("Tags", [])}

    return tags


def get_using_iam_api(arns: list[str]) -> dict[str, dict]:
    """Fetch tags for IAM resources via dedicated API."""
    if not arns:
        return {}

    client = boto3.client("iam")
    tags = {}

    for arn in track(arns, description="Fetching tags (IAM)..."):
        parts = arn.split(":")
        if len(parts) < 6:
            continue

        segment = parts[5]
        if "/" not in segment:
            continue

        segment = segment.split("/")
        rt = segment[0]
        name = segment[-1]

        try:
            if rt == "role":
                response = client.list_role_tags(RoleName=name)
                tags[arn] = {t["Key"]: t["Value"] for t in response.get("Tags", [])}
            elif rt == "user":
                response = client.list_user_tags(UserName=name)
                tags[arn] = {t["Key"]: t["Value"] for t in response.get("Tags", [])}
        except client.exceptions.NoSuchEntityException:
            pass

    return tags


def get_using_wafv2_api(arns: list[str]) -> dict[str, dict]:
    """Fetch tags for WAFv2 resources via dedicated API."""
    if not arns:
        return {}

    client = boto3.client("wafv2")
    tags = {}

    for arn in track(arns, description="Fetching tags (WAFv2)..."):
        try:
            response = client.list_tags_for_resource(ResourceARN=arn)
            tags[arn] = {
                t["Key"]: t["Value"]
                for t in response.get("TagInfoForResource", {}).get("TagList", [])
            }
        except client.exceptions.WAFNonexistentItemException:
            pass

    return tags


def get_using_autoscaling_api(arns: list[str]) -> dict[str, dict]:
    """Fetch tags for Auto Scaling Groups via dedicated API."""
    if not arns:
        return {}

    client = boto3.client("autoscaling")
    tags = {}

    names = []
    by_arn = {}
    for arn in arns:
        parts = arn.split(":")
        if len(parts) >= 8:
            segment = parts[7]
            if segment.startswith("autoScalingGroupName/"):
                name = segment.split("/", 1)[1]
                names.append(name)
                by_arn[arn] = name

    if not names:
        return {}

    paginator = client.get_paginator("describe_tags")
    by_name: dict[str, dict] = defaultdict(dict)

    for page in paginator.paginate(
        Filters=[{"Name": "auto-scaling-group", "Values": names}]
    ):
        for tag in page["Tags"]:
            by_name[tag["ResourceId"]][tag["Key"]] = tag["Value"]

    for arn, name in by_arn.items():
        tags[arn] = by_name.get(name, {})

    return tags


GETTERS: dict[str, Getter] = {
    "iam:role": get_using_iam_api,
    "iam:user": get_using_iam_api,
    "wafv2:ipset": get_using_wafv2_api,
    "wafv2:rulegroup": get_using_wafv2_api,
    "wafv2:webacl": get_using_wafv2_api,
    "autoscaling:autoScalingGroup": get_using_autoscaling_api,
}


def get_tags(arn_to_type: dict[str, str], batch_size: int) -> dict[str, dict]:
    """Fetch tags for ARNs, routing by resource type."""
    if not arn_to_type:
        raise ValueError("arn_to_type cannot be empty")

    tagging_arns: list[str] = []
    custom: dict[Getter, list[str]] = defaultdict(list)

    for arn, rt in arn_to_type.items():
        getter = GETTERS.get(rt)
        if getter:
            custom[getter].append(arn)
        else:
            tagging_arns.append(arn)

    tags = get_using_tagging_api(tagging_arns, batch_size)

    for getter, arns in custom.items():
        tags.update(getter(arns))

    return tags
