from __future__ import annotations

from typing import TYPE_CHECKING

import libcst as cst

from aiservice.common.markdown_utils import split_markdown_code
from core.languages.python.optimizer.context_utils.context_helpers import is_multi_context
from core.languages.python.testgen.instrumentation.edit_generated_test import replace_definition_with_import
from core.languages.python.testgen.postprocessing.add_missing_imports import add_missing_imports
from core.languages.python.testgen.postprocessing.range_modifier import modify_large_loops
from core.languages.python.testgen.postprocessing.remove_unused_definitions import (
    remove_unused_definitions_from_pytest_file,
)
from core.languages.python.testgen.postprocessing.removeassert_transformer import remove_asserts_from_test
from core.languages.python.testgen.postprocessing.tensor_limit import modify_tensors
from core.languages.python.testgen.postprocessing.topdef_terminator import delete_top_def_nodes

if TYPE_CHECKING:
    from collections.abc import Callable
    from typing import Any

    from libcst import Module

    from aiservice.models.functions_to_optimize import FunctionToOptimize


def postprocessing_testgen_pipeline(
    module: Module,
    helper_function_names: list[str],
    function_to_optimize: FunctionToOptimize,
    module_path: str,
    source_code_being_tested: str,
) -> tuple[Module, Module]:
    """Full postprocessing pipeline for generated test code.

    Applies all CST transformations in sequence:
    1. Clean up definitions (remove helper functions, unused definitions)
    2. Modify constructs (large loops, tensors)
    3. Add missing imports, replace function definition with import
    4. Remove asserts (for instrumentation only)

    Returns (processed_module_with_asserts_removed, processed_module_with_asserts_kept).
    The second module is the "display" version — fully cleaned but with original asserts intact.
    """
    source_cst: Module | dict[str, Module]
    if is_multi_context(source_code_being_tested):
        source_cst = {
            path: cst.parse_module(code) for path, code in split_markdown_code(source_code_being_tested).items()
        }
    else:
        source_cst = cst.parse_module(source_code_being_tested)

    cleanup_pipeline: list[tuple[Callable[..., Module], dict[str, Any]]] = [
        (delete_top_def_nodes, {"deletable_list": helper_function_names}),
        (remove_unused_definitions_from_pytest_file, {}),
        (modify_large_loops, {}),
        (modify_tensors, {}),
        (add_missing_imports, {"source_cst": source_cst, "module_path": module_path}),
        (replace_definition_with_import, {"function": function_to_optimize, "module_path": module_path}),
    ]

    for func, kwargs in cleanup_pipeline:
        module = func(module, **kwargs)

    # Capture the display version (with asserts) before removing them
    display_module = module

    # Remove asserts for the instrumentation version
    module = remove_asserts_from_test(
        module,
        function_to_optimize=function_to_optimize,
        helper_function_names=helper_function_names,
        module_path=module_path,
    )

    return module, display_module
