import logging
from collections import defaultdict
from typing import Any

import libcst as cst

from py_agent.models.enum_metadata import EnumMetadata
from py_agent.models.focused_code import ExternalMethodInfo, FocusedCodeDTO
from py_agent.services.import_jedi_extractor import ImportJediExtractor
from py_agent.visitors.enums_collector_visitor import EnumCollector
from py_agent.visitors.focused_code_visitor import (
    ClearCodeTransformer,
    CodeInfo,
    FunctionOrMethodVisitor,
    MethodInfo,
    MethodsList,
    RemoveUnusedImportTransformer,
)
from py_agent.visitors.import_path_visitor import ImportPathVisitor

logger = logging.getLogger(__name__)


def is_method_in_call_tree(method_name: str, call_tree: set[str]) -> bool:
    return method_name in call_tree


def is_constructor(method_name: str) -> bool:
    return method_name == "__init__"


def methods_scope_check(methods_list: MethodsList, target_method: MethodInfo) -> set[str]:
    methods_in_scope: set[str] = set()
    methods_name_checked: list[str] = []
    method_to_check: list[str] = []

    method_to_check.append(target_method.name)
    methods_name_checked.append(target_method.name)
    methods_in_scope.add(target_method.name)

    methods_names = methods_list.methods_dict.keys()

    while method_to_check:
        curr_method = method_to_check.pop()
        curr_method_call_tree = methods_list.methods_dict[curr_method].method_call_tree

        for method in methods_names:
            if method == curr_method:
                continue
            if is_method_in_call_tree(method, curr_method_call_tree) or is_constructor(method):
                methods_in_scope.add(method)
                if method not in methods_name_checked:
                    method_to_check.append(method)
        methods_name_checked.append(curr_method)

    return methods_in_scope


def get_functions_and_methods_info(
    file_path: str, target_name: str, class_name: str | None = None
) -> CodeInfo:
    with open(file_path, encoding="utf-8") as f:
        source_code = f.read()

    if not source_code:
        logger.error(f"failed to read file {file_path}")
        return CodeInfo(method_info=None, methods_list=None, external_method_map={})

    module = cst.parse_module(source_code)
    wrapper = cst.MetadataWrapper(module)
    visitor = FunctionOrMethodVisitor(
        target_name=target_name, class_name=class_name, wrapper=wrapper
    )
    wrapper.visit(visitor)

    jedi_extractor = ImportJediExtractor(absolute_file_path=file_path)
    external_methods_map: dict[str, ExternalMethodInfo] = {}

    if visitor.method_info:
        for func in visitor.method_info.method_call_tree:
            if func not in visitor.all_methods:
                pos = visitor.method_info.call_function_position.get(func)
                if not pos:
                    continue
                line, col = pos
                func_metadata = jedi_extractor.get_by_position(line, col)
                if not func_metadata:
                    logger.debug(f"Jedi could not resolve function {func}")
                    continue
                if func_metadata.type != "function":
                    continue
                if func_metadata.is_core:
                    continue
                external_methods_map[func] = ExternalMethodInfo(
                    name=func, jedi_metadata=func_metadata
                )

    return CodeInfo(
        method_info=visitor.method_info,
        methods_list=MethodsList(methods_dict=visitor.all_methods),
        external_method_map=external_methods_map,
    )


def prepare_code(file_path: str, target_methods_list: set[str]) -> str:
    with open(file_path, encoding="utf-8") as f:
        source_code = f.read()
    if not source_code:
        return ""
    module = cst.parse_module(source_code)
    transformer = ClearCodeTransformer(methods_to_keep=target_methods_list, module=module)
    modified_module = module.visit(transformer)
    return modified_module.code


def remove_unused_imports(code: str) -> str:
    wrapper = cst.metadata.MetadataWrapper(cst.parse_module(code))
    scopes = {s for s in wrapper.resolve(cst.metadata.ScopeProvider).values() if s is not None}
    unused_imports: dict[cst.Import | cst.ImportFrom, set[str]] = defaultdict(set)

    for scope in scopes:
        for assignment in scope.assignments:
            if not isinstance(assignment, cst.metadata.Assignment):
                continue
            node = assignment.node
            if isinstance(node, (cst.Import, cst.ImportFrom)):
                if len(assignment.references) == 0:
                    unused_imports[node].add(assignment.name)

    fixed_module = wrapper.module.visit(RemoveUnusedImportTransformer(unused_imports))
    return fixed_module.code


def get_imports_path(
    code: str, project_path: str, absolute_file_path: str, relative_file_path: str
) -> dict[str, Any]:
    module = cst.parse_module(code)
    visitor = ImportPathVisitor(project_path, absolute_file_path, module)
    module.visit(visitor)
    return visitor.local_imports


def get_enums_from_imports(imports: dict[str, Any]) -> list[EnumMetadata]:
    enums_in_scope: list[EnumMetadata] = []
    for _key, item in imports.items():
        if item.module_path and len(item.module_path) > 0:
            with open(item.module_path, encoding="utf-8") as f:
                source_code = f.read()
            module = cst.parse_module(source_code)
            visitor = EnumCollector(module)
            module.visit(visitor)
            enums = visitor.enum_collection
            if enums:
                if item.type == "full":
                    for enum_name, enum_code in enums.items():
                        enums_in_scope.append(
                            EnumMetadata(
                                name=enum_name,
                                code=enum_code,
                                module_path=item.module_path,
                                module_name=item.module_name,
                            )
                        )
                else:
                    for imported_item in item.imported_items:
                        if imported_item in enums:
                            enums_in_scope.append(
                                EnumMetadata(
                                    name=imported_item,
                                    code=enums[imported_item],
                                    module_path=item.module_path,
                                    module_name=item.module_name,
                                )
                            )
    return enums_in_scope


class FocusedCodeExtractor:
    @staticmethod
    def get_focused_code_by_function_name(
        absolute_file_path: str,
        relative_file_path: str,
        project_path: str,
        function_name: str,
        class_name: str | None,
    ) -> FocusedCodeDTO:
        logger.debug("get_focused_code_by_function_name: get code info")
        code_info = get_functions_and_methods_info(
            absolute_file_path, target_name=function_name, class_name=class_name
        )

        all_methods_list = (
            list(code_info.all_methods_list.methods_dict.keys())
            if code_info.all_methods_list
            else []
        )
        all_external_methods_list = code_info.external_method_map

        if not code_info.method_info:
            return FocusedCodeDTO(
                is_extracted=False,
                file_path=absolute_file_path,
                function_name=function_name,
                class_name=class_name,
                focused_code=None,
                called_methods=None,
                all_methods_list=all_methods_list,
                all_external_methods_list=all_external_methods_list,
                enums=[],
            )

        logger.debug("get_focused_code_by_function_name: checking scope")
        assert code_info.all_methods_list is not None
        methods_to_keep = methods_scope_check(code_info.all_methods_list, code_info.method_info)

        logger.debug(
            f"get_focused_code_by_function_name: remove unused methods, keep {methods_to_keep}"
        )
        focused_code = prepare_code(absolute_file_path, methods_to_keep)

        logger.debug("get_focused_code_by_function_name: remove unused imports")
        fixed_code = remove_unused_imports(focused_code)

        logger.debug("get_focused_code_by_function_name: get_imports_path")
        imports = get_imports_path(fixed_code, project_path, absolute_file_path, relative_file_path)

        enums = get_enums_from_imports(imports)
        logger.debug(f"get_focused_code_by_function_name: enums={enums}")

        return FocusedCodeDTO(
            is_extracted=True,
            file_path=absolute_file_path,
            function_name=function_name,
            class_name=class_name,
            focused_code=fixed_code,
            called_methods=list(methods_to_keep),
            all_methods_list=all_methods_list,
            all_external_methods_list=all_external_methods_list,
            enums=enums,
        )
