import logging
from collections import defaultdict
from typing import Any

import libcst as cst
from libcst.metadata import MetadataWrapper, PositionProvider

logger = logging.getLogger(__name__)


class MethodInfo:
    def __init__(
        self,
        name: str,
        return_type: str | None,
        method_code: str,
        call_tree: set[str],
        call_function_position: dict[str, tuple[int, int]],
        call_position: dict[str, tuple[int, int]],
    ):
        self.name = name
        self.return_type = return_type
        self.method_code = method_code
        self.method_call_tree = call_tree
        self.call_position = call_position
        self.call_function_position = call_function_position

    def __repr__(self) -> str:
        return f"MethodInfo(name={self.name}, return_type={self.return_type}, method_call_tree={self.method_call_tree})"


class MethodsList:
    def __init__(self, methods_dict: dict[str, Any]) -> None:
        self.methods_dict = methods_dict

    def __repr__(self) -> str:
        return f"MethodsList(methods_dict={self.methods_dict})"


class CodeInfo:
    def __init__(
        self,
        method_info: MethodInfo | None,
        methods_list: MethodsList | None,
        external_method_map: dict[str, Any],
    ):
        self.method_info = method_info
        self.all_methods_list = methods_list
        self.external_method_map = external_method_map

    def __repr__(self) -> str:
        return (
            f"CodeInfo(method_info={self.method_info})\n"
            f"CodeInfo(methods_list={self.all_methods_list})\n"
            f"CodeInfo(external_method_map={self.external_method_map})"
        )


class ClearCodeTransformer(cst.CSTTransformer):
    def __init__(self, methods_to_keep: set[str], module: cst.Module):
        self.module = module
        self.methods_to_keep = methods_to_keep

    def leave_FunctionDef(
        self, original_node: cst.FunctionDef, updated_node: cst.FunctionDef
    ) -> cst.BaseStatement | cst.FlattenSentinel[cst.BaseStatement] | cst.RemovalSentinel:
        if original_node.name.value not in self.methods_to_keep:
            return cst.RemovalSentinel.REMOVE
        return updated_node


class FunctionOrMethodVisitor(cst.CSTVisitor):
    METADATA_DEPENDENCIES = (PositionProvider,)

    def __init__(self, target_name: str, wrapper: MetadataWrapper, class_name: str | None = None):
        self.target_name = target_name
        self.class_name = class_name
        self.wrapper = wrapper
        self.module = wrapper.module
        self.is_in_target_class = False
        self.method_info: MethodInfo | None = None
        self.all_methods: dict[str, Any] = defaultdict(set)
        self.curr_visited_method = ""
        self.curr_visited_method_call_tree: set[str] = set()
        self.call_function_position: dict[str, tuple[int, int]] = {}
        self.call_position: dict[str, tuple[int, int]] = {}

    def visit_Call(self, node: cst.Call) -> None:
        call_name = ""
        if isinstance(node.func, cst.Name):
            call_name = node.func.value
        elif isinstance(node.func, cst.Attribute):
            call_name = self.module.code_for_node(node.func)
        self.curr_visited_method_call_tree.add(call_name)

        function_position = self.get_metadata(PositionProvider, node.func, None)
        call_position = self.get_metadata(PositionProvider, node, None)

        if function_position:
            self.call_function_position[call_name] = (
                function_position.start.line,
                function_position.end.column,
            )
        if call_position:
            self.call_position[call_name] = (call_position.start.line, call_position.end.column)

    def visit_ClassDef(self, node: cst.ClassDef) -> None:
        if self.class_name and node.name.value == self.class_name:
            self.is_in_target_class = True
        else:
            self.is_in_target_class = False

    def visit_FunctionDef(self, node: cst.FunctionDef) -> None:
        if self.class_name:
            if self.is_in_target_class and node.name.value == self.target_name:
                self.method_info = self._extract_function_info(node)
        elif node.name.value == self.target_name:
            self.method_info = self._extract_function_info(node)
        self.curr_visited_method = node.name.value

    def leave_FunctionDef(self, node: cst.FunctionDef) -> None:
        self.all_methods[node.name.value] = self._extract_function_info(node)
        self.curr_visited_method = ""
        self.curr_visited_method_call_tree = set()
        self.call_function_position = {}
        self.call_position = {}

    def _extract_function_info(self, node: cst.FunctionDef) -> MethodInfo:
        return_type = (
            node.returns.annotation.value
            if node.returns and isinstance(node.returns.annotation, cst.Name)
            else None
        )
        method_code = self.module.code_for_node(node)
        return MethodInfo(
            node.name.value,
            return_type,
            method_code,
            self.curr_visited_method_call_tree,
            call_function_position=self.call_function_position,
            call_position=self.call_position,
        )


class RemoveUnusedImportTransformer(cst.CSTTransformer):
    def __init__(self, unused_imports: dict[cst.Import | cst.ImportFrom, set[str]]) -> None:
        self.unused_imports = unused_imports

    def leave_import_alike(
        self,
        original_node: cst.Import | cst.ImportFrom,
        updated_node: cst.Import | cst.ImportFrom,
    ) -> cst.Import | cst.ImportFrom | cst.RemovalSentinel:
        if original_node not in self.unused_imports:
            return updated_node
        if isinstance(updated_node.names, cst.ImportStar):
            return updated_node
        names_to_keep = []
        for name in updated_node.names:
            asname = name.asname
            name_value = asname.name.value if asname is not None else name.name.value  # type: ignore[union-attr]
            if name_value not in self.unused_imports[original_node]:
                names_to_keep.append(name.with_changes(comma=cst.MaybeSentinel.DEFAULT))
        if len(names_to_keep) == 0:
            return cst.RemoveFromParent()
        return updated_node.with_changes(names=names_to_keep)

    def leave_Import(self, original_node: cst.Import, updated_node: cst.Import) -> cst.Import:
        return self.leave_import_alike(original_node, updated_node)  # type: ignore[return-value]

    def leave_ImportFrom(
        self, original_node: cst.ImportFrom, updated_node: cst.ImportFrom
    ) -> cst.ImportFrom:
        return self.leave_import_alike(original_node, updated_node)  # type: ignore[return-value]
