"""Detect and revert unused helper functions in optimized code."""

from __future__ import annotations

import ast
import logging
from collections import defaultdict
from itertools import chain
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionSource, FunctionToOptimize
    from ..context.models import CodeOptimizationContext, CodeStringsMarkdown

log = logging.getLogger(__name__)


def find_target_node(
    root: ast.AST,
    function_to_optimize: FunctionToOptimize,
) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
    """Find the AST node for the target function inside its parent scopes."""
    parents = function_to_optimize.parents
    node = root
    for parent in parents:
        body = getattr(node, "body", None)
        if not body:
            return None
        for child in body:
            if isinstance(child, ast.ClassDef) and child.name == parent.name:
                node = child
                break
        else:
            return None
    body = getattr(node, "body", None)
    if not body:
        return None
    target_name = function_to_optimize.function_name
    for child in body:
        if (
            isinstance(child, (ast.FunctionDef, ast.AsyncFunctionDef))
            and child.name == target_name
        ):
            return child
    return None


def _collect_attr_names(
    value_id: str,
    attr_name: str,
    class_name: str | None,
    names: set[str],
    imported_names_map: dict[str, set[str]],
) -> None:
    """Collect attribute reference names including self and class lookups."""
    if value_id == "self":
        names.add(attr_name)
        if class_name:
            names.add(f"{class_name}.{attr_name}")
    else:
        names.add(attr_name)
        full_ref = f"{value_id}.{attr_name}"
        names.add(full_ref)
        mapped_names = imported_names_map.get(full_ref)
        if mapped_names:
            names.update(mapped_names)


def _collect_called_names(
    entrypoint_ast: ast.FunctionDef | ast.AsyncFunctionDef,
    function_to_optimize: FunctionToOptimize,
    imported_names_map: dict[str, set[str]],
) -> set[str]:
    """Collect all function names called within the entrypoint AST node."""
    called = {function_to_optimize.function_name}
    class_name = (
        function_to_optimize.parents[0].name
        if function_to_optimize.parents
        else None
    )

    for node in ast.walk(entrypoint_ast):
        if isinstance(node, ast.Call):
            if isinstance(node.func, ast.Name):
                called.add(node.func.id)
                mapped_names = imported_names_map.get(node.func.id)
                if mapped_names:
                    called.update(mapped_names)
            elif isinstance(node.func, ast.Attribute):
                if isinstance(node.func.value, ast.Name):
                    _collect_attr_names(
                        node.func.value.id,
                        node.func.attr,
                        class_name,
                        called,
                        imported_names_map,
                    )
                else:
                    called.add(node.func.attr)
        elif isinstance(node, ast.Attribute) and isinstance(
            node.value,
            ast.Name,
        ):
            _collect_attr_names(
                node.value.id,
                node.attr,
                class_name,
                called,
                imported_names_map,
            )

    return called


def _collect_from_import_helpers(
    node: ast.ImportFrom,
    helpers_by_file_and_func: dict[str, dict[str, list[FunctionSource]]],
    imported_names_map: dict[str, set[str]],
) -> None:
    """Process a from-import node to map imported names to helper FQNs."""
    module_name = node.module
    if not module_name:
        return
    file_entry = helpers_by_file_and_func.get(module_name)
    if not file_entry:
        return
    for alias in node.names:
        imported_name = alias.asname or alias.name
        original_name = alias.name
        helpers = file_entry.get(original_name)
        if helpers:
            imported_set = imported_names_map[imported_name]
            for helper in helpers:
                imported_set.add(helper.qualified_name)
                imported_set.add(helper.fully_qualified_name)


def _collect_plain_import_helpers(
    node: ast.Import,
    helpers_by_file: dict[str, list[FunctionSource]],
    imported_names_map: dict[str, set[str]],
) -> None:
    """Process a plain import node to map imported names to helper FQNs."""
    for alias in node.names:
        imported_name = alias.asname or alias.name
        module_name_imp = alias.name
        helpers = helpers_by_file.get(module_name_imp)
        if helpers:
            for helper in helpers:
                full_call = f"{imported_name}.{helper.only_function_name}"
                full_call_set = imported_names_map[full_call]
                full_call_set.add(helper.qualified_name)
                full_call_set.add(helper.fully_qualified_name)


def _analyze_imports_in_optimized_code(
    optimized_ast: ast.AST,
    code_context: CodeOptimizationContext,
) -> dict[str, set[str]]:
    """Map imported names to qualified helper names based on import statements."""
    imported_names_map: dict[str, set[str]] = defaultdict(set)

    helpers_by_file_and_func: dict[
        str,
        dict[str, list[FunctionSource]],
    ] = defaultdict(dict)
    helpers_by_file: dict[str, list[FunctionSource]] = defaultdict(list)
    for helper in code_context.helper_functions:
        jedi_type = helper.definition_type
        if jedi_type != "class":
            func_name = helper.only_function_name
            module_name = helper.file_path.stem
            if func_name is not None:
                helpers_by_file_and_func[module_name].setdefault(
                    func_name,
                    [],
                ).append(helper)
            helpers_by_file[module_name].append(helper)

    for node in ast.walk(optimized_ast):
        if isinstance(node, ast.ImportFrom):
            _collect_from_import_helpers(
                node,
                helpers_by_file_and_func,
                imported_names_map,
            )
        elif isinstance(node, ast.Import):
            _collect_plain_import_helpers(
                node,
                helpers_by_file,
                imported_names_map,
            )

    return dict(imported_names_map)


def detect_unused_helper_functions(
    function_to_optimize: FunctionToOptimize,
    code_context: CodeOptimizationContext,
    optimized_code: str | CodeStringsMarkdown,
) -> list[FunctionSource]:
    """Detect helper functions that are not called in the optimized code."""
    from ..context.models import (  # noqa: PLC0415
        CodeStringsMarkdown as CodeStringsMarkdownCls,
    )

    if (
        isinstance(optimized_code, CodeStringsMarkdownCls)
        and len(
            optimized_code.code_strings,
        )
        > 0
    ):
        return list(
            chain.from_iterable(
                detect_unused_helper_functions(
                    function_to_optimize,
                    code_context,
                    code.code,
                )
                for code in optimized_code.code_strings
            ),
        )

    if not isinstance(optimized_code, str):
        return []

    try:
        optimized_ast = ast.parse(optimized_code)
        entrypoint_function_ast = find_target_node(
            optimized_ast,
            function_to_optimize,
        )

        if not entrypoint_function_ast:
            log.debug(
                "Could not find entrypoint function %s in optimized code",
                function_to_optimize.function_name,
            )
            return []

        imported_names_map = _analyze_imports_in_optimized_code(
            optimized_ast,
            code_context,
        )
        called_function_names = _collect_called_names(
            entrypoint_function_ast,
            function_to_optimize,
            imported_names_map,
        )

        log.debug(
            "Functions called in optimized entrypoint: %s",
            called_function_names,
        )

        unused_helpers = _find_unused_helpers(
            code_context,
            called_function_names,
            function_to_optimize.file_path,
        )

    except Exception:  # noqa: BLE001
        log.debug(
            "Error detecting unused helper functions",
            exc_info=True,
        )
        return []
    else:
        return unused_helpers


def _find_unused_helpers(
    code_context: CodeOptimizationContext,
    called_function_names: set[str],
    entrypoint_file_path: Path,
) -> list[FunctionSource]:
    """Return helpers from *code_context* not present in *called_function_names*."""
    unused_helpers: list[FunctionSource] = []
    for helper_function in code_context.helper_functions:
        if helper_function.definition_type == "class":
            continue
        helper_qualified_name = helper_function.qualified_name
        helper_simple_name = helper_function.only_function_name
        helper_fully_qualified_name = helper_function.fully_qualified_name

        is_called = (
            helper_qualified_name in called_function_names
            or helper_simple_name in called_function_names
            or helper_fully_qualified_name in called_function_names
            or (
                helper_function.file_path != entrypoint_file_path
                and f"{helper_function.file_path.stem}.{helper_simple_name}"
                in called_function_names
            )
        )

        if not is_called:
            unused_helpers.append(helper_function)
            log.debug(
                "Helper function %s is not called in optimized code",
                helper_qualified_name,
            )
    return unused_helpers


def revert_unused_helper_functions(
    project_root: Path,
    unused_helpers: list[FunctionSource],
    original_helper_code: dict[Path, str],
) -> None:
    """Revert unused helper functions back to their original definitions."""
    from ..codegen._replacement import (  # noqa: PLC0415
        is_zero_diff,
        replace_functions_and_add_imports,
    )

    if not unused_helpers:
        return

    log.debug(
        "Reverting %d unused helper function(s) to original definitions",
        len(unused_helpers),
    )

    resolved_original_helper_code = {
        p.resolve(): code for p, code in original_helper_code.items()
    }

    unused_helpers_by_file: dict[Path, list[FunctionSource]] = defaultdict(
        list,
    )
    for helper in unused_helpers:
        unused_helpers_by_file[helper.file_path.resolve()].append(helper)

    for file_path, helpers_in_file in unused_helpers_by_file.items():
        if file_path not in resolved_original_helper_code:
            continue
        try:
            original_code = resolved_original_helper_code[file_path]
            helper_names = [
                helper.qualified_name for helper in helpers_in_file
            ]
            source_code = file_path.read_text(encoding="utf8")
            new_code = replace_functions_and_add_imports(
                source_code,
                helper_names,
                original_code,
                file_path,
                set(),
                project_root,
            )
            if not is_zero_diff(source_code, new_code):
                file_path.write_text(new_code, encoding="utf8")
                log.debug(
                    "Reverted unused helpers in %s: %s",
                    file_path,
                    ", ".join(helper_names),
                )
        except Exception:
            log.exception(
                "Error reverting unused helpers in %s",
                file_path,
            )
