"""CST pruning for context-type-specific code views."""

from __future__ import annotations

from typing import TYPE_CHECKING

import libcst as cst

from .dependencies import (
    collect_top_level_defs_with_usages,
    get_section_names,
    is_assignment_used,
    recurse_sections,
)
from .models import CodeContextType, PruneConfig

if TYPE_CHECKING:
    from .models import UsageInfo


def maybe_strip_docstring(
    node: cst.FunctionDef | cst.ClassDef,
    cfg: PruneConfig,
) -> cst.FunctionDef | cst.ClassDef:
    """Strip docstring from function or class if configured."""
    if not cfg.remove_docstrings:
        return node
    if not isinstance(node.body, cst.IndentedBlock):
        return node

    body_stmts = node.body.body
    if not body_stmts:
        return node

    first_stmt = body_stmts[0]
    if (
        isinstance(first_stmt, cst.SimpleStatementLine)
        and len(first_stmt.body) == 1
    ):
        expr_stmt = first_stmt.body[0]
        if isinstance(expr_stmt, cst.Expr) and isinstance(
            expr_stmt.value,
            (cst.SimpleString, cst.ConcatenatedString),
        ):
            new_body = body_stmts[1:] or [
                cst.SimpleStatementLine(body=[cst.Pass()])
            ]
            return node.with_changes(
                body=node.body.with_changes(body=new_body)
            )

    return node


def prune_cst(  # noqa: C901, PLR0911, PLR0912
    node: cst.CSTNode,
    target_functions: set[str],
    cfg: PruneConfig,
    prefix: str = "",
) -> tuple[cst.CSTNode | None, bool]:
    """Recursively filter the CST, keeping relevant nodes."""
    if isinstance(node, (cst.Import, cst.ImportFrom)):
        return None, False

    if isinstance(node, cst.FunctionDef):
        name = node.name.value
        qualified_name = f"{prefix}.{name}" if prefix else name

        if cfg.helpers and qualified_name in cfg.helpers:
            return maybe_strip_docstring(node, cfg), True

        if qualified_name in target_functions:
            if cfg.exclude_init_from_targets and name == "__init__":
                return None, False
            if cfg.include_target_in_output:
                return (
                    maybe_strip_docstring(node, cfg),
                    True,
                )
            return None, True

        if cfg.keep_class_init and name == "__init__":
            return node, False

        if (
            cfg.include_dunder_methods
            and len(name) > 4  # noqa: PLR2004
            and name.startswith("__")
            and name.endswith("__")
        ):
            if not cfg.include_init_dunder and name == "__init__":
                return None, False
            return maybe_strip_docstring(node, cfg), False

        return None, False

    if isinstance(node, cst.ClassDef):
        if prefix:
            return None, False
        if not isinstance(node.body, cst.IndentedBlock):
            msg = "ClassDef body is not an IndentedBlock"
            raise TypeError(msg)
        class_name = node.name.value

        if cfg.defs_with_usages:
            has_targets = any(
                isinstance(stmt, cst.FunctionDef)
                and f"{class_name}.{stmt.name.value}" in target_functions
                for stmt in node.body.body
            )
            if (
                not has_targets
                and class_name in cfg.defs_with_usages
                and cfg.defs_with_usages[class_name].used_by_qualified_function
            ):
                return node, True

        new_class_body: list[cst.CSTNode] = []
        found_in_class = False

        for stmt in node.body.body:
            filtered, found_target = prune_cst(
                stmt, target_functions, cfg, class_name
            )
            found_in_class |= found_target
            if filtered:
                new_class_body.append(filtered)

        if not found_in_class:
            return None, False
        if not new_class_body:
            return None, True
        updated = node.with_changes(
            body=node.body.with_changes(body=new_class_body)
        )
        return maybe_strip_docstring(updated, cfg), True

    # Handle assignments for READ_WRITABLE mode
    if cfg.defs_with_usages is not None and isinstance(
        node, (cst.Assign, cst.AnnAssign, cst.AugAssign)
    ):
        if is_assignment_used(node, cfg.defs_with_usages):
            return node, True
        return None, False

    # For other nodes, recursively process children
    section_names = get_section_names(node)
    if not section_names:
        return node, False

    return recurse_sections(
        node,
        section_names,
        lambda child: prune_cst(child, target_functions, cfg, prefix),
        keep_non_target_children=cfg.helpers is not None,
    )


def parse_code_and_prune_cst(  # noqa: PLR0913
    code: str | cst.Module,
    code_context_type: CodeContextType,
    target_functions: set[str],
    helpers_of_helper_functions: set[str] | None = None,
    *,
    remove_docstrings: bool = False,
    defs_with_usages: dict[str, UsageInfo] | None = None,
) -> cst.Module:
    """Parse and filter the code CST, returning the pruned Module."""
    if helpers_of_helper_functions is None:
        helpers_of_helper_functions = set()

    module = code if isinstance(code, cst.Module) else cst.parse_module(code)

    if code_context_type == CodeContextType.READ_WRITABLE:
        if defs_with_usages is None:
            defs_with_usages = collect_top_level_defs_with_usages(
                module,
                target_functions | helpers_of_helper_functions,
            )
        cfg = PruneConfig(
            defs_with_usages=defs_with_usages,
            keep_class_init=True,
        )
    elif code_context_type == CodeContextType.READ_ONLY:
        cfg = PruneConfig(
            helpers=helpers_of_helper_functions,
            remove_docstrings=remove_docstrings,
            include_target_in_output=False,
            include_dunder_methods=True,
        )
    elif code_context_type == CodeContextType.TESTGEN:
        cfg = PruneConfig(
            helpers=helpers_of_helper_functions,
            remove_docstrings=remove_docstrings,
            include_dunder_methods=True,
            include_init_dunder=True,
        )
    elif code_context_type == CodeContextType.HASHING:
        cfg = PruneConfig(
            remove_docstrings=True,
            exclude_init_from_targets=True,
        )
    else:
        msg = f"Unknown code_context_type: {code_context_type}"
        raise ValueError(msg)

    filtered_node, found_target = prune_cst(module, target_functions, cfg)

    if not found_target:
        msg = "No target functions found in the provided code"
        raise ValueError(msg)
    if filtered_node and isinstance(filtered_node, cst.Module):
        return filtered_node
    msg = "Pruning produced no module"
    raise ValueError(msg)
