from __future__ import annotations

import ast
import atexit
import inspect
import os.path
import re
import shutil
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, Mapping, Sequence

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).

    """

    trgt: str | re.Pattern[str] | tuple[str | re.Pattern[str], ...]
    repl: str | Sequence[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 load_stmts(code: str) -> list[ast.stmt]:
    return ast.parse(code).body


def _remove_source_indent(lines: Sequence[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]


_AWEPATCH_DEBUG = int(os.getenv("AWEPATCH_DEBUG") or 0)


_awepatch_temp_dir: str | None = None
_awepatch_temp_dir_lock = threading.Lock()


def _get_awepatch_temp_dir() -> str:
    """Get or create the temporary directory for awepatch."""
    global _awepatch_temp_dir

    if _awepatch_temp_dir is not None:
        return _awepatch_temp_dir

    with _awepatch_temp_dir_lock:
        if _awepatch_temp_dir is not None:
            return _awepatch_temp_dir

        _awepatch_temp_dir = tempfile.mkdtemp(prefix="awepatch-")
        if not _AWEPATCH_DEBUG:
            # Register cleanup at exit
            atexit.register(
                lambda: (
                    _awepatch_temp_dir
                    and shutil.rmtree(_awepatch_temp_dir, ignore_errors=True)
                )
            )
        return _awepatch_temp_dir


_patch_count = 0
_patch_count_lock = threading.Lock()


def load_function_code(
    func: ast.FunctionDef | ast.AsyncFunctionDef,
    origin: str = "",
) -> 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 _patch_count_lock:
        global _patch_count
        _patch_count += 1

    temp_dir = _get_awepatch_temp_dir()
    file_path = os.path.join(temp_dir, f"patch_{_patch_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 _is_match_node(
    node: ast.AST,
    source: Sequence[str],
    pattern: str | re.Pattern[str],
) -> bool:
    """Check if the AST node matches the target pattern.

    Args:
        node: The AST node to check
        source: The source code lines
        pattern: The target pattern

    Returns:
        True if the AST node matches the target pattern, False otherwise

    """
    if not isinstance(node, ast.stmt):
        return False

    node_lines = source[node.lineno - 1 : node.end_lineno]
    node_source = "".join(_remove_source_indent(node_lines))

    return bool(
        isinstance(pattern, str)
        and node_source.startswith(pattern)
        or isinstance(pattern, re.Pattern)
        and pattern.match(node_source)
    )


def _find_matched_node(
    node: ast.AST,
    source: Sequence[str],
    target: tuple[str | re.Pattern[str], ...],
) -> tuple[ast.AST, str, int] | None:
    """Recursively find the target AST node matching the target patterns.

    Args:
        node: The AST node to search
        source: The source code lines
        target: The target patterns

    Returns:
        The AST node matching the target patterns, or None if not found

    """
    matched: tuple[ast.AST, str, int] | None = None

    for key, field in ast.iter_fields(node):
        if isinstance(field, list):
            for index, item in enumerate(field):  # pyright: ignore[reportUnknownVariableType, reportUnknownArgumentType]
                if not isinstance(item, ast.AST):
                    continue

                if _is_match_node(item, source, target[0]):
                    if len(target) == 1:
                        tmp_matched = (node, key, index)
                    else:
                        tmp_matched = _find_matched_node(item, source, target[1:])
                else:
                    tmp_matched = _find_matched_node(item, source, target)

                if tmp_matched is not None:
                    if matched is not None:
                        raise ValueError(
                            f"Multiple matches found for target pattern {target}"
                        )
                    matched = tmp_matched

        elif isinstance(field, ast.AST):
            if len(target) == 1:
                tmp_matched = _find_matched_node(field, source, target)
            elif _is_match_node(field, source, target[0]):
                tmp_matched = _find_matched_node(field, source, target[1:])
            else:
                tmp_matched = _find_matched_node(field, source, target)

            if tmp_matched is not None:
                if matched is not None:
                    raise ValueError(
                        f"Multiple matches found for target pattern {target}"
                    )
                matched = tmp_matched

    return matched


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

    Applies patches in order: before, replace, after.

    Returns:
        The total number of statements added (can be negative if replaced with fewer).

    """
    i = index

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

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

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

    return i - index  # Total number of statements added


def _check_targets(
    target: str | re.Pattern[str] | tuple[str | re.Pattern[str], ...],
) -> tuple[str | re.Pattern[str], ...]:
    """Check the targets of a Patch object.

    Args:
        target: The target pattern(s) to check.

    Raises:
        ValueError: If the targets are invalid.

    """
    if not target:
        raise ValueError("Target pattern cannot be empty")

    if isinstance(target, (str, re.Pattern)):
        return (target,)
    elif (
        isinstance(target, tuple)  # pyright: ignore[reportUnnecessaryIsInstance]
        and all(isinstance(t, (str, re.Pattern)) for t in target)  # pyright: ignore[reportUnnecessaryIsInstance]
    ):
        return target
    else:
        raise TypeError("Unknown target pattern type")


def _compile_patches(
    func: ast.FunctionDef | ast.AsyncFunctionDef,
    source: list[str],
    patches: list[Patch],
) -> defaultdict[tuple[ast.AST, str], defaultdict[int, dict[Mode, Sequence[ast.stmt]]]]:
    """Compile patches into a dictionary mapping line numbers to mode-specific
    statements.

    Args:
        func: The AST function definition
        source: 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[
        tuple[ast.AST, str], defaultdict[int, dict[Mode, Sequence[ast.stmt]]]
    ] = defaultdict(lambda: defaultdict(dict))

    for patch in patches:
        # Get target line number
        target = _find_matched_node(
            func,
            source,
            _check_targets(patch.trgt),
        )
        if target is None:
            raise ValueError(f"No match found for target pattern {patch.trgt!r}")

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

        if target_patches := compiled_patches[target[:2]][target[2]]:
            # Check for duplicate mode on the same line
            if patch.mode in target_patches:
                raise ValueError(
                    f"Multiple '{patch.mode}' patches on the same target {patch.trgt!r}"
                )
            # Check for conflicting patches: replace cannot be combined with other modes
            if patch.mode == "replace" or "replace" in target_patches:
                raise ValueError(
                    "Cannot combine 'replace' with other modes "
                    f"on target {patch.trgt!r}"
                )
        target_patches[patch.mode] = repl_stmts

    return compiled_patches


def _apply_compiled_patches(
    compiled_patches: Mapping[
        tuple[ast.AST, str], Mapping[int, dict[Mode, Sequence[ast.stmt]]]
    ],
) -> None:
    """Apply compiled patches to the AST function definition.

    Args:
        func: The AST function definition
        compiled_patches: Compiled patches mapping line numbers to mode-specific
          statements

    Returns:
        True if any patches were applied, False otherwise

    """
    for (node, key), index_patches in compiled_patches.items():
        trgt_items = getattr(node, key)
        offset = 0
        for index, patches in sorted(index_patches.items()):
            offset += _apply_stmts_patches(trgt_items, index + offset, 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 = get_source_lines(func)

    # 2. Parse AST and validate
    func_ast = ast.parse("".join(source))
    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 to get plain function code
    func_ast.decorator_list.clear()

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

    # 5. Apply all compiled patches to the AST
    _apply_compiled_patches(compiled_patches)

    return func_ast
