from __future__ import annotations

import ast
import inspect
import os.path
import re
import sys
import tempfile
import threading
from collections import defaultdict
from functools import partial
from types import CodeType, FunctionType
from typing import TYPE_CHECKING, Any, Literal, NamedTuple, TypeAlias

if TYPE_CHECKING:
    from collections.abc import Callable

Mode: TypeAlias = Literal["before", "after", "replace"]


class Patch(NamedTuple):
    """A single patch operation.

    Attributes:
        pattern: The pattern to search for in the source code.
        repl: The replacement code or AST statements.
        mode: The mode of patching (before/after/replace).

    """

    pattern: str | re.Pattern[str]
    repl: str | list[ast.stmt]
    mode: Mode = "before"


def get_origin_function(func: Callable[..., Any]) -> Callable[..., Any]:
    """Get the object wrapped by *func*.

    Follows the chain of :attr:`__wrapped__` or :attr:`__func__` attributes or
    :class:`functools.partial` objects, returning the last object in the chain.

    :exc:`ValueError` is raised if a cycle is encountered.

    """
    f = func  # remember the original func for error reporting
    # Memoise by id to tolerate non-hashable objects, but store objects to
    # ensure they aren't destroyed, which would allow their IDs to be reused.
    memo = {id(f): f}
    recursion_limit = sys.getrecursionlimit()
    while not isinstance(func, type):
        if hasattr(func, "__wrapped__"):
            func = func.__wrapped__  # pyright: ignore[reportFunctionMemberAccess]
        elif hasattr(func, "__func__"):
            func = func.__func__  # pyright: ignore[reportFunctionMemberAccess]
        elif isinstance(func, partial):
            func = func.func
        else:
            break
        id_func = id(func)
        if (id_func in memo) or (len(memo) >= recursion_limit):
            raise ValueError(f"wrapper loop when unwrapping {f!r}")
        memo[id_func] = func
    return func


def find_line_number(lines: list[str], pattern: str | re.Pattern[str]) -> int:
    if not isinstance(pattern, (str, re.Pattern)):  # pyright: ignore[reportUnnecessaryIsInstance]
        raise TypeError(f"Unknown pattern type: {type(pattern)}")

    lineno = None
    for index, line in enumerate(lines):
        line = line.strip()

        if (isinstance(pattern, str) and line == pattern) or (
            isinstance(pattern, re.Pattern) and pattern.search(line)
        ):
            if lineno is not None:
                raise ValueError(f"Multiple matches found for pattern: {pattern}")
            lineno = index + 1

    if lineno is None:
        raise ValueError(f"No match found for pattern: {pattern}")

    return lineno


def load_stmts(code: str) -> list[ast.stmt]:
    return ast.parse(code).body


def _remove_source_indent(lines: list[str]) -> list[str]:
    """Remove common leading indentation from code lines."""

    indent = len(lines[0]) - len(lines[0].lstrip())
    # for blank lines indent may be larger than line length, so we replace with '\n'
    lines = [line[indent:] if len(line) > indent else "\n" for line in lines]
    return lines


def get_source_lines(
    obj: CodeType | FunctionType,
) -> list[str]:
    """Get the source lines of a function or code object.

    Args:
        obj (CodeType | FunctionType): The function or code object to get the source
          lines.

    """
    source_lines = inspect.getsourcelines(obj)[0]
    return _remove_source_indent(source_lines)


def _find_function_code(module: CodeType) -> CodeType:
    funcs = [
        func
        for func in module.co_consts
        if isinstance(func, CodeType) and func.co_name != "<lambda>"
    ]
    if len(funcs) != 1:
        raise ValueError("Only single function definitions are supported")
    return funcs[0]


patched_count = 0
patched_count_lock = threading.Lock()

_awepatch_temp_dir = os.path.abspath(tempfile.mkdtemp(prefix="awepatch_"))


def load_function_code(
    func: ast.FunctionDef | ast.AsyncFunctionDef,
    origin: str = "",
    temp_dir: str = _awepatch_temp_dir,
) -> CodeType:
    """Load a function's code object from its AST module.

    Args:
        func (ast.Module): The AST module containing the function definition.
        origin (str, optional): The origin location for the function. Defaults to "".
        temp_dir (str): Temporary directory to store the generated file.

    Returns:
        CodeType: The code object of the function.

    """

    source = ast.unparse(func)
    origin = f" (patched from {origin})" if origin else ""
    source = f"# generated by awepatch{origin}\n{source}"

    with patched_count_lock:
        global patched_count
        patched_count += 1

    file_path = os.path.join(temp_dir, f"{patched_count:03d}__{func.name}.py")
    with open(file_path, "w") as f:
        f.write(source)
    module_code = compile(source, filename=file_path, mode="exec")
    func_code = _find_function_code(module_code)
    if func_code.co_name != func.name:
        # This should never happen though better to be safe than sorry
        raise ValueError("Function name mismatch!")

    return func_code


def _pop_matched_patches(
    target: ast.AST,
    patches: dict[int, dict[Mode, list[ast.stmt]]],
) -> dict[Mode, list[ast.stmt]]:
    """Pop matched statements from the AST based on the target line number."""
    if not isinstance(target, ast.stmt):
        return {}

    if target.end_lineno is None:
        return {}

    matched_patches: dict[Mode, list[ast.stmt]] = {}
    for lineno in list(patches):
        if lineno < target.lineno or lineno > target.end_lineno:
            continue

        for mode, stmts in patches.pop(lineno).items():
            if matched_patches:
                # Check for duplicate mode on the same line
                if mode in matched_patches:
                    raise ValueError(
                        f"Multiple '{mode}' patches on "
                        f"the same target `{ast.unparse(target)}`"
                    )

                # replace cannot be combined with other modes
                if mode == "replace" or "replace" in matched_patches:
                    raise ValueError(
                        f"Cannot combine 'replace' with other modes on"
                        f"the same target `{ast.unparse(target)}`"
                    )
            matched_patches[mode] = stmts

    return matched_patches


def _modify_ast_node(
    ast_node: ast.AST,
    patches: dict[int, dict[Mode, list[ast.stmt]]],
) -> bool:
    """Find and modify the statement at the specified line number in the AST.

    Args:
        ast_node: The AST node to traverse
        target: The target line number
        patches: Dictionary mapping modes to their replacement statements

    Returns:
        True if the target statement was found and modified, False otherwise

    """
    for _, field in ast.iter_fields(ast_node):
        if isinstance(field, list):
            # Search for target statement in list fields
            if _modify_ast_list(field, patches):  # pyright: ignore[reportUnknownArgumentType]
                return True
        elif isinstance(field, ast.AST):  # noqa: SIM102
            # Recursively process child nodes
            if _modify_ast_node(field, patches):
                return True

    return False


def _modify_ast_list(
    ast_list: list[Any],
    patches: dict[int, dict[Mode, list[ast.stmt]]],
) -> bool:
    """Find and modify the target statement in a list of statements.

    Args:
        ast_list: The list of AST statements
        target: The target line number
        patches: Dictionary mapping modes to their replacement statements

    Returns:
        True if the target statement was found and modified, False otherwise

    """
    index = 0
    while index < len(ast_list):
        item = ast_list[index]
        if not isinstance(item, ast.AST):
            index += 1
            continue

        # First recursively check child nodes
        if _modify_ast_node(item, patches):
            return True

        # Check if current statement matches target line number
        if matched := _pop_matched_patches(item, patches):
            index = _apply_stmts_patches(ast_list, index, matched)
            if not patches:
                return True

        index += 1

    return False


def _apply_stmts_patches(
    stmts: list[Any],
    index: int,
    patches: dict[Mode, list[ast.stmt]],
) -> int:
    """Insert or replace statements at the specified position with multiple modes.

    Applies patches in order: before, replace, after.
    """
    # Apply 'before' first
    if "before" in patches:
        stmts[index:index] = patches["before"]
        index += len(patches["before"])  # Adjust index after insertion

    # Apply 'replace' (which removes the original statement)
    if "replace" in patches:
        stmts[index : index + 1] = patches["replace"]
        index += len(patches["replace"]) - 1  # Adjust for replacement

    # Apply 'after' last
    if "after" in patches:
        stmts[index + 1 : index + 1] = patches["after"]
        index += len(patches["after"])  # Adjust index after insertion

    return index


def _compile_patches(
    source_lines: list[str],
    patches: list[Patch],
) -> dict[int, dict[Mode, list[ast.stmt]]]:
    """Compile patches into a dictionary mapping line numbers to mode-specific
    statements.

    Args:
        source_lines: The source code lines of the function
        patches: List of Patch objects to compile

    Returns:
        Dictionary mapping line numbers to dictionaries of mode->statements

    Raises:
        ValueError: If duplicate modes on same line or conflicting patches

    """
    compiled_patches: defaultdict[int, dict[Mode, list[ast.stmt]]] = defaultdict(dict)

    for patch in patches:
        # Get target line number
        target_lineno = find_line_number(source_lines, patch.pattern)

        # Prepare replacement statements
        repl_stmts = (
            load_stmts(patch.repl) if isinstance(patch.repl, str) else patch.repl
        )

        if line_patches := compiled_patches[target_lineno]:
            # Check for duplicate mode on the same line
            if patch.mode in line_patches:
                raise ValueError(
                    f"Multiple '{patch.mode}' patches on the same line {target_lineno}"
                )

            # Check for conflicting patches: replace cannot be combined with other modes
            if patch.mode == "replace" or "replace" in line_patches:
                raise ValueError(
                    f"Cannot combine 'replace' with other modes on line {target_lineno}"
                )

        line_patches[patch.mode] = repl_stmts

    return compiled_patches


def ast_patch(
    func: CodeType | FunctionType,
    patches: list[Patch],
) -> ast.FunctionDef | ast.AsyncFunctionDef:
    """Patch the AST of a function or code object.

    Args:
        func (CodeType | FunctionType): The function or code object to patch.
        patches (list[Patch]): List of Patch objects for applying multiple patches.

    Returns:
        ast.Module: The modified AST module contains only one function definition.

    """
    # Validate arguments
    if not patches:
        raise ValueError("patches list cannot be empty")

    # 1. Get source code
    source_lines = get_source_lines(func)

    # 2. Parse AST and validate
    func_ast = ast.parse("".join(source_lines))
    if len(func_ast.body) != 1:
        raise ValueError("Only single function definitions are supported")
    func_ast = func_ast.body[0]  # type: ignore[assignment]
    if not isinstance(func_ast, (ast.FunctionDef, ast.AsyncFunctionDef)):
        raise ValueError(f"Not a function definition: {type(func_ast)}")

    # 3. Clear decorators
    func_ast.decorator_list.clear()

    # 4. Pre-compile all patches: find line numbers and prepare replacement statements
    compiled_patches = _compile_patches(source_lines, patches)

    # 5. Modify AST nodes in place
    if not _modify_ast_node(func_ast, compiled_patches):
        raise ValueError(f"No ast.stmt found for patches: {compiled_patches}")

    return func_ast
