"""Replace function definitions in Python source code."""

from __future__ import annotations

import ast
import logging
import os
from collections import defaultdict
from itertools import chain
from typing import TYPE_CHECKING, TypeVar

import libcst as cst

from .._model import FunctionParent
from ..context.imports import (
    gather_source_imports as gather_source_imports,  # noqa: PLC0414
)
from ..verification._ranking import (
    normalize_code as normalize_code,  # noqa: PLC0414
)
from ..verification._ranking import (
    normalize_node as normalize_node,  # noqa: PLC0414
)

# -- Re-exports from sub-modules ----------------------------------
# Every name that was previously importable from this module is
# re-exported here so external callers need no changes.
from ._global_defs import (
    GlobalAssignmentCollector as GlobalAssignmentCollector,  # noqa: PLC0414
)
from ._global_defs import (
    GlobalAssignmentTransformer as GlobalAssignmentTransformer,  # noqa: PLC0414
)
from ._global_defs import (
    GlobalFunctionCollector as GlobalFunctionCollector,  # noqa: PLC0414
)
from ._global_defs import (
    GlobalFunctionTransformer as GlobalFunctionTransformer,  # noqa: PLC0414
)
from ._global_defs import (
    GlobalStatementCollector as GlobalStatementCollector,  # noqa: PLC0414
)
from ._global_defs import (
    GlobalStatementTransformer as GlobalStatementTransformer,  # noqa: PLC0414
)
from ._global_defs import (
    add_global_assignments as add_global_assignments,  # noqa: PLC0414
)
from ._global_defs import (
    collect_referenced_names as collect_referenced_names,  # noqa: PLC0414
)
from ._global_defs import (
    extract_global_statements as extract_global_statements,  # noqa: PLC0414
)
from ._global_defs import (
    find_insertion_index_after_imports as find_insertion_index_after_imports,  # noqa: PLC0414
)
from ._import_management import (
    DottedImportCollector as DottedImportCollector,  # noqa: PLC0414
)
from ._import_management import (
    FutureAliasedImportTransformer as FutureAliasedImportTransformer,  # noqa: PLC0414
)
from ._import_management import (
    add_needed_imports_from_module as add_needed_imports_from_module,  # noqa: PLC0414
)
from ._import_management import (
    delete_future_aliased_imports as delete_future_aliased_imports,  # noqa: PLC0414
)
from ._import_management import (
    resolve_star_import as resolve_star_import,  # noqa: PLC0414
)
from ._pytest_transforms import (
    AddRequestArgument as AddRequestArgument,  # noqa: PLC0414
)
from ._pytest_transforms import (
    AutouseFixtureModifier as AutouseFixtureModifier,  # noqa: PLC0414
)
from ._pytest_transforms import (
    PytestMarkAdder as PytestMarkAdder,  # noqa: PLC0414
)
from ._pytest_transforms import (
    add_custom_marker_to_all_tests as add_custom_marker_to_all_tests,  # noqa: PLC0414
)
from ._pytest_transforms import (
    disable_autouse as disable_autouse,  # noqa: PLC0414
)
from ._pytest_transforms import (
    has_autouse_fixture as has_autouse_fixture,  # noqa: PLC0414
)
from ._pytest_transforms import (
    modify_autouse_fixture as modify_autouse_fixture,  # noqa: PLC0414
)

if TYPE_CHECKING:
    from pathlib import Path

    from .._model import FunctionToOptimize
    from ..context.models import CodeStringsMarkdown

log = logging.getLogger(__name__)

ASTNodeT = TypeVar("ASTNodeT", bound=ast.AST)


# -- Core replacement functions ------------------------------------


def replace_function_source(
    source: str,
    function: FunctionToOptimize,
    new_source: str,
) -> str:
    """Replace *function* in *source* with *new_source*.

    Uses libcst to preserve formatting of surrounding code.
    The optimized code's body and decorators replace the
    original; everything else in the file stays untouched.
    """
    class_name = function.class_name
    func_name = function.function_name

    try:
        optimized_func = _find_function(
            cst.parse_module(new_source),
            class_name,
            func_name,
        )
    except cst.ParserSyntaxError:
        msg = f"Failed to parse new_source for {function.qualified_name!r}"
        raise ValueError(msg) from None
    if optimized_func is None:
        msg = f"Function {function.qualified_name!r} not found in new_source"
        raise ValueError(msg)

    try:
        original = cst.parse_module(source)
    except cst.ParserSyntaxError:
        msg = f"Failed to parse source for {function.qualified_name!r}"
        raise ValueError(msg) from None
    new_body: list[cst.BaseStatement] = []

    for node in original.body:
        if (
            class_name is None
            and isinstance(node, cst.FunctionDef)
            and node.name.value == func_name
        ):
            new_body.append(
                node.with_changes(
                    body=optimized_func.body,
                    decorators=optimized_func.decorators,
                ),
            )
        elif (
            class_name is not None
            and isinstance(node, cst.ClassDef)
            and node.name.value == class_name
        ):
            new_body.append(
                _replace_method_in_class(
                    node,
                    func_name,
                    optimized_func,
                ),
            )
        else:
            new_body.append(node)

    return original.with_changes(body=new_body).code


def _replace_method_in_class(
    cls: cst.ClassDef,
    method_name: str,
    optimized_func: cst.FunctionDef,
) -> cst.ClassDef:
    """Return *cls* with *method_name* replaced."""
    new_members: list[cst.BaseStatement | cst.BaseSmallStatement] = []
    for child in cls.body.body:
        if (
            isinstance(child, cst.FunctionDef)
            and child.name.value == method_name
        ):
            new_members.append(
                child.with_changes(
                    body=optimized_func.body,
                    decorators=optimized_func.decorators,
                ),
            )
        else:
            new_members.append(child)
    return cls.with_changes(
        body=cls.body.with_changes(body=new_members),
    )


def _find_function(
    module: cst.Module,
    class_name: str | None,
    func_name: str,
) -> cst.FunctionDef | None:
    """Find a function in *module* by class and function name."""
    for node in module.body:
        if (
            class_name is None
            and isinstance(node, cst.FunctionDef)
            and node.name.value == func_name
        ):
            return node
        if (
            class_name is not None
            and isinstance(node, cst.ClassDef)
            and node.name.value == class_name
        ):
            for child in node.body.body:
                if (
                    isinstance(child, cst.FunctionDef)
                    and child.name.value == func_name
                ):
                    return child
    return None


def is_zero_diff(original_code: str, new_code: str) -> bool:
    """Return True when the optimization changed nothing."""
    return normalize_code(original_code) == normalize_code(
        new_code,
    )


def extract_function_names(code: str) -> list[str]:
    """Extract top-level and class-level function names.

    Returns dotted names for class methods
    (e.g. ``"Cls.method"``) and bare names for
    module-level functions.
    """
    names: list[str] = []
    tree = ast.parse(code)
    for node in tree.body:
        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)):
            names.append(node.name)
        elif isinstance(node, ast.ClassDef):
            cls_name = node.name
            names.extend(
                f"{cls_name}.{child.name}"
                for child in node.body
                if isinstance(
                    child,
                    (
                        ast.FunctionDef,
                        ast.AsyncFunctionDef,
                    ),
                )
            )
    return names


# -- Multi-function replacement ------------------------------------


def _parse_function_names(
    names: list[str],
) -> list[tuple[str | None, str]] | None:
    """Parse dotted function names into (class, func) tuples.

    Returns None if any name has unsupported nesting.
    """
    result: list[tuple[str | None, str]] = []
    for name in names:
        if "." not in name:
            result.append((None, name))
        elif name.count(".") == 1:
            cls, fn = name.split(".")
            result.append((cls, fn))
        else:
            log.error(
                "Unable to find %s. Returning unchanged source code.",
                name,
            )
            return None
    return result


def _classify_optimized_nodes(  # noqa: C901
    optimized_module: cst.Module,
    names_set: set[tuple[str | None, str]],
    preexisting: set[tuple[str, tuple[FunctionParent, ...]]],
) -> tuple[
    dict[tuple[str | None, str], cst.FunctionDef],
    list[cst.FunctionDef],
    dict[str, list[cst.FunctionDef]],
    list[cst.ClassDef],
    dict[str, cst.FunctionDef],
]:
    """Classify optimized code nodes into categories.

    Returns modified functions, new functions, new class
    methods, new classes, and modified ``__init__`` methods.
    """
    modified: dict[tuple[str | None, str], cst.FunctionDef] = {}
    new_funcs: list[cst.FunctionDef] = []
    new_cls_funcs: dict[str, list[cst.FunctionDef]] = defaultdict(list)
    new_classes: list[cst.ClassDef] = []
    modified_inits: dict[str, cst.FunctionDef] = {}

    for node in optimized_module.body:
        if isinstance(node, cst.FunctionDef):
            key = (None, node.name.value)
            if key in names_set:
                modified[key] = node
            elif preexisting and ((node.name.value, ()) not in preexisting):
                new_funcs.append(node)
        elif isinstance(node, cst.ClassDef):
            cls_name = node.name.value
            parents = (
                FunctionParent(
                    name=cls_name,
                    type="ClassDef",
                ),
            )
            if (cls_name, ()) not in preexisting:
                new_classes.append(node)
            for child in node.body.body:
                if not isinstance(
                    child,
                    cst.FunctionDef,
                ):
                    continue
                mkey = (cls_name, child.name.value)
                if mkey in names_set:
                    modified[mkey] = child
                elif (
                    child.name.value == "__init__"
                    and preexisting
                    and (cls_name, ()) in preexisting
                ):
                    modified_inits[cls_name] = child
                elif preexisting and (
                    (child.name.value, parents) not in preexisting
                ):
                    new_cls_funcs[cls_name].append(
                        child,
                    )

    return (
        modified,
        new_funcs,
        new_cls_funcs,
        new_classes,
        modified_inits,
    )


def _rebuild_body(
    original_module: cst.Module,
    modified: dict[tuple[str | None, str], cst.FunctionDef],
    modified_inits: dict[str, cst.FunctionDef],
    new_class_funcs: dict[str, list[cst.FunctionDef]],
) -> tuple[list[cst.BaseStatement], set[str]]:
    """Rebuild module body with function replacements.

    Returns the new body and set of existing class names.
    """
    new_body: list[cst.BaseStatement] = []
    existing_cls: set[str] = set()

    for node in original_module.body:
        if isinstance(node, cst.FunctionDef):
            key = (None, node.name.value)
            if key in modified:
                mf = modified[key]
                new_body.append(
                    node.with_changes(
                        body=mf.body,
                        decorators=mf.decorators,
                    ),
                )
            else:
                new_body.append(node)
        elif isinstance(node, cst.ClassDef):
            new_body.append(
                _rebuild_class(
                    node,
                    modified,
                    modified_inits,
                    new_class_funcs,
                ),
            )
            existing_cls.add(node.name.value)
        else:
            new_body.append(node)

    return new_body, existing_cls


def _rebuild_class(
    node: cst.ClassDef,
    modified: dict[tuple[str | None, str], cst.FunctionDef],
    modified_inits: dict[str, cst.FunctionDef],
    new_class_funcs: dict[str, list[cst.FunctionDef]],
) -> cst.ClassDef:
    """Rebuild a class with replaced/added methods."""
    cls_name = node.name.value
    members: list[cst.BaseStatement | cst.BaseSmallStatement] = []
    for child in node.body.body:
        if isinstance(child, cst.FunctionDef):
            key = (cls_name, child.name.value)
            if key in modified:
                mf = modified[key]
                members.append(
                    child.with_changes(
                        body=mf.body,
                        decorators=mf.decorators,
                    ),
                )
            elif child.name.value == "__init__" and cls_name in modified_inits:
                members.append(
                    modified_inits[cls_name],
                )
            else:
                members.append(child)
        else:
            members.append(child)

    if cls_name in new_class_funcs:
        members.extend(new_class_funcs[cls_name])

    return node.with_changes(
        body=node.body.with_changes(body=members),
    )


def replace_functions_in_file(
    source_code: str,
    original_function_names: list[str],
    optimized_code: str,
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
) -> str:
    """Replace functions in *source_code* with optimized versions.

    Handles preexisting-object dedup, ``__init__``
    replacement, and insertion of new helper
    functions/classes.
    """
    parsed = _parse_function_names(
        original_function_names,
    )
    if parsed is None:
        return source_code

    optimized_module = cst.parse_module(
        optimized_code,
    )
    (
        modified,
        new_funcs,
        new_cls_funcs,
        new_classes,
        modified_inits,
    ) = _classify_optimized_nodes(
        optimized_module,
        set(parsed),
        preexisting_objects,
    )

    original_module = cst.parse_module(source_code)

    max_fn_idx: int | None = None
    max_cls_idx: int | None = None
    for idx, _node in enumerate(original_module.body):
        if isinstance(_node, cst.FunctionDef):
            max_fn_idx = idx
        if isinstance(_node, cst.ClassDef):
            max_cls_idx = idx

    new_body, existing_cls = _rebuild_body(
        original_module,
        modified,
        modified_inits,
        new_cls_funcs,
    )

    if new_classes:
        unique = [
            nc for nc in new_classes if nc.name.value not in existing_cls
        ]
        if unique:
            ins = (
                max_cls_idx
                if max_cls_idx is not None
                else find_insertion_index_after_imports(
                    original_module,
                )
            )
            new_body = list(
                chain(
                    new_body[:ins],
                    unique,
                    new_body[ins:],
                ),
            )

    if new_funcs:
        if max_fn_idx is not None:
            new_body = [
                *new_body[: max_fn_idx + 1],
                *new_funcs,
                *new_body[max_fn_idx + 1 :],
            ]
        elif max_cls_idx is not None:
            new_body = [
                *new_body[: max_cls_idx + 1],
                *new_funcs,
                *new_body[max_cls_idx + 1 :],
            ]
        else:
            new_body = [*new_funcs, *new_body]

    return original_module.with_changes(
        body=new_body,
    ).code


def replace_functions_and_add_imports(  # noqa: PLR0913
    source_code: str,
    function_names: list[str],
    optimized_code: str,
    module_abspath: Path,
    preexisting_objects: set[tuple[str, tuple[FunctionParent, ...]]],
    project_root_path: Path,
) -> str:
    """Replace functions and add any new imports.

    Combines :func:`replace_functions_in_file` with
    :func:`add_needed_imports_from_module` in one call.
    """
    return add_needed_imports_from_module(
        optimized_code,
        replace_functions_in_file(
            source_code,
            function_names,
            optimized_code,
            preexisting_objects,
        ),
        module_abspath,
        module_abspath,
        project_root_path,
    )


def get_optimized_code_for_module(
    relative_path: Path,
    optimized_code: CodeStringsMarkdown,
) -> str:
    """Return the optimized code block for *relative_path*.

    Tries an exact path match first, then falls back to:
    1. A single ``None``-keyed code block.
    2. Basename matching (the LLM sometimes returns
       wrong directory).
    """
    file_to_code = optimized_code.file_to_code()
    module_optimized_code = file_to_code.get(
        str(relative_path),
    )
    if module_optimized_code is not None:
        return module_optimized_code

    # Fallback 1: single code block with no file path
    if "None" in file_to_code and len(file_to_code) == 1:
        log.debug(
            "Using code block with None file_path for %s",
            relative_path,
        )
        return file_to_code["None"]

    # Fallback 2: match by filename (basename)
    target_name = relative_path.name
    basename_matches = [
        code
        for path, code in file_to_code.items()
        if path != "None" and os.path.basename(path) == target_name  # noqa: PTH119
    ]
    if len(basename_matches) == 1:
        log.debug(
            "Using basename-matched code block for %s",
            relative_path,
        )
        return basename_matches[0]

    log.warning(
        "Optimized code not found for %s, existing files are %s",
        relative_path,
        list(file_to_code.keys()),
    )
    return ""
