import inspect
from collections import defaultdict
from pathlib import Path
from types import FunctionType

import libcst

from localstack.testing.aws.asf_utils import collect_implemented_provider_operations

# collect all provider operations
provider_ops = collect_implemented_provider_operations()

# iterate over all operations and group all operations which need to be modified per provider file
operations_per_provider_file: dict[str, list[str]] = defaultdict(list)
for provider_class, base_class, provider_function in provider_ops:
    try:
        sub_function = getattr(provider_class, provider_function)
    except AttributeError:
        raise AttributeError(
            f"Given method name ('{provider_function}') is not a method of the sub class ('{provider_class.__name__}')."
        )

    if not isinstance(sub_function, FunctionType):
        raise AttributeError(
            f"Given method name ('{provider_function}') is not a method of the sub class ('{provider_class.__name__}')."
        )

    if not getattr(sub_function, "expand_parameters", True):
        # if the operation on the subclass has the "expand_parameters" attribute (it has a handler decorator) set to False, we don't care
        continue

    if wrapped := getattr(sub_function, "__wrapped__", False):
        # if the operation on the subclass has a decorator, unwrap it
        sub_function = wrapped

    operations_per_provider_file[inspect.getfile(provider_class)].append(sub_function.__name__)


class ProviderMethodTransformer(libcst.CSTTransformer):
    def __init__(self, operations: list[str]):
        self.operations = operations
        super().__init__()

    def leave_FunctionDef(
        self, original_node: libcst.FunctionDef, updated_node: libcst.FunctionDef
    ) -> libcst.CSTNode:
        if original_node.name.value in self.operations and not original_node.params.star_kwarg:
            print(f"  -> {original_node.name.value}")
            new_params = updated_node.params.with_changes(
                star_kwarg=libcst.Param(name=libcst.Name("kwargs"))
            )
            return updated_node.with_changes(params=new_params)

        return updated_node


# iterate over all providers and tranform the found operations
for provider_file, ops in operations_per_provider_file.items():
    print(f"- {provider_file}")
    module = libcst.parse_module(Path(provider_file).read_bytes())
    updated_module = module.visit(ProviderMethodTransformer(ops))
    Path(provider_file).write_text(updated_module.code)