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

from collections import defaultdict
from dataclasses import dataclass, field
from typing import Callable

import boto3
from rich.progress import track

Setter = Callable[[dict[str, dict[str, str]]], dict[str, bool]]


@dataclass
class TagChange:
    """Tag change to apply."""

    arn: str
    resource_type: str
    tags: dict[str, str]


@dataclass
class SetTagsResult:
    """Result of applying tag changes."""

    success: list[str] = field(default_factory=list)
    failed: dict[str, str] = field(default_factory=dict)


def set_using_tagging_api(arn_to_tags: dict[str, dict[str, str]]) -> dict[str, bool]:
    """Set tags via ResourceGroupsTaggingAPI."""
    if not arn_to_tags:
        return {}

    client = boto3.client("resourcegroupstaggingapi")
    results = {}

    for arn, tags in track(
        arn_to_tags.items(), description="Setting tags (tagging API)..."
    ):
        try:
            client.tag_resources(ResourceARNList=[arn], Tags=tags)
            results[arn] = True
        except Exception:
            results[arn] = False

    return results


def set_using_iam_api(arn_to_tags: dict[str, dict[str, str]]) -> dict[str, bool]:
    """Set tags for IAM resources via dedicated API."""
    if not arn_to_tags:
        return {}

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

    for arn, tags in track(arn_to_tags.items(), description="Setting tags (IAM)..."):
        parts = arn.split(":")
        if len(parts) < 6:
            results[arn] = False
            continue

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

        segment = segment.split("/")
        rt = segment[0]
        name = segment[-1]
        payload = [{"Key": k, "Value": v} for k, v in tags.items()]

        try:
            if rt == "role":
                client.tag_role(RoleName=name, Tags=payload)
                results[arn] = True
            elif rt == "user":
                client.tag_user(UserName=name, Tags=payload)
                results[arn] = True
            else:
                results[arn] = False
        except Exception:
            results[arn] = False

    return results


def set_using_wafv2_api(arn_to_tags: dict[str, dict[str, str]]) -> dict[str, bool]:
    """Set tags for WAFv2 resources."""
    if not arn_to_tags:
        return {}

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

    for arn, tags in track(arn_to_tags.items(), description="Setting tags (WAFv2)..."):
        payload = [{"Key": k, "Value": v} for k, v in tags.items()]
        try:
            client.tag_resource(ResourceARN=arn, Tags=payload)
            results[arn] = True
        except Exception:
            results[arn] = False

    return results


def set_using_autoscaling_api(
    arn_to_tags: dict[str, dict[str, str]],
) -> dict[str, bool]:
    """Set tags for Auto Scaling Groups."""
    if not arn_to_tags:
        return {}

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

    for arn, tags in track(arn_to_tags.items(), description="Setting tags (ASG)..."):
        parts = arn.split(":")
        if len(parts) < 8:
            results[arn] = False
            continue

        segment = parts[7]
        if not segment.startswith("autoScalingGroupName/"):
            results[arn] = False
            continue

        name = segment.split("/", 1)[1]
        payload = [
            {
                "ResourceId": name,
                "ResourceType": "auto-scaling-group",
                "Key": k,
                "Value": v,
                "PropagateAtLaunch": True,
            }
            for k, v in tags.items()
        ]

        try:
            client.create_or_update_tags(Tags=payload)
            results[arn] = True
        except Exception:
            results[arn] = False

    return results


SETTERS: dict[str, Setter] = {
    "iam:role": set_using_iam_api,
    "iam:user": set_using_iam_api,
    "wafv2:ipset": set_using_wafv2_api,
    "wafv2:rulegroup": set_using_wafv2_api,
    "wafv2:webacl": set_using_wafv2_api,
    "autoscaling:autoScalingGroup": set_using_autoscaling_api,
}


def set_tags(changes: list[TagChange]) -> SetTagsResult:
    """Apply tag changes, routing by resource type."""
    if not changes:
        return SetTagsResult()

    tagging: dict[str, dict[str, str]] = {}
    custom: dict[Setter, dict[str, dict[str, str]]] = defaultdict(dict)

    for c in changes:
        setter = SETTERS.get(c.resource_type)
        if setter:
            custom[setter][c.arn] = c.tags
        else:
            tagging[c.arn] = c.tags

    result = SetTagsResult()

    if tagging:
        for arn, ok in set_using_tagging_api(tagging).items():
            if ok:
                result.success.append(arn)
            else:
                result.failed[arn] = "Failed to set tags"

    for setter, arn_to_tags in custom.items():
        for arn, ok in setter(arn_to_tags).items():
            if ok:
                result.success.append(arn)
            else:
                result.failed[arn] = "Failed to set tags"

    return result
