# Copyright (c) ONNX Project Contributors
# SPDX-License-Identifier: Apache-2.0
"""Wrappers for IR classes to enable journaling.

This module provides wrapper functions that enable journaling for ONNX IR classes.
The wrappers are applied when a Journal context is active, and they record operations
to the journal for debugging and analysis purposes.
"""
# mypy: disable-error-code="attr-defined,method-assign,assignment"

from __future__ import annotations

import functools
from typing import Any, Callable, TypeVar

from typing_extensions import Concatenate, ParamSpec

from onnx_ir import _core, _graph_containers
from onnx_ir.journaling import _journaling

_P = ParamSpec("_P")
_T = TypeVar("_T")
_SelfT = TypeVar("_SelfT")
_ValueT = TypeVar("_ValueT")


# =============================================================================
# Generic wrapper factories
# =============================================================================


def _init_wrapper(
    journal: _journaling.Journal,
    original_init: Callable[Concatenate[_SelfT, _P], None],
    *,
    details_func: Callable[[_SelfT], str | None] = repr,
) -> Callable[Concatenate[_SelfT, _P], None]:
    """Generic wrapper factory for __init__ methods.

    Args:
        journal: The journal to record to.
        original_init: The original __init__ method.
        details_func: A function that takes self and returns details string.
    """

    @functools.wraps(original_init)
    def wrapper(self: _SelfT, *args: _P.args, **kwargs: _P.kwargs) -> None:
        original_init(self, *args, **kwargs)
        journal.record(self, "init", details=details_func(self))

    return wrapper


def _setter_wrapper(
    journal: _journaling.Journal,
    original_setter: Callable[[_SelfT, _ValueT], None],
    property_name: str,
    operation: str,
) -> Callable[[_SelfT, _ValueT], None]:
    """Generic wrapper factory for property setters.

    Args:
        journal: The journal to record to.
        original_setter: The original setter method.
        property_name: The private attribute name (e.g., "_name").
        operation: The operation name for the journal (e.g., "set_name").
    """

    @functools.wraps(original_setter)
    def wrapper(self: _SelfT, value: _ValueT) -> None:
        old_value = getattr(self, property_name)
        journal.record(self, operation, details=f"{old_value!r} -> {value!r}")
        original_setter(self, value)

    return wrapper


def _method_wrapper(
    journal: _journaling.Journal,
    original_method: Callable[Concatenate[_SelfT, _P], _T],
    operation: str,
    *,
    details_func: Callable[Concatenate[_SelfT, _P], str | None],
) -> Callable[Concatenate[_SelfT, _P], _T]:
    """Generic wrapper factory for methods.

    Args:
        journal: The journal to record to.
        original_method: The original method.
        operation: The operation name for the journal.
        details_func: A function that takes (self, *args, **kwargs) and returns details.
    """

    @functools.wraps(original_method)
    def wrapper(self: _SelfT, *args: _P.args, **kwargs: _P.kwargs) -> _T:
        journal.record(self, operation, details=details_func(self, *args, **kwargs))
        return original_method(self, *args, **kwargs)

    return wrapper


def _container_method_wrapper(
    journal: _journaling.Journal,
    original_method: Callable[Concatenate[_SelfT, _P], _T],
    operation: str,
    *,
    target_attr: str,
    details_func: Callable[Concatenate[_SelfT, _P], str | None],
) -> Callable[Concatenate[_SelfT, _P], _T]:
    """Generic wrapper factory for container methods that record on a parent object.

    Args:
        journal: The journal to record to.
        original_method: The original method.
        operation: The operation name for the journal.
        target_attr: The attribute name to get the target object from self (e.g., "_graph").
        details_func: A function that takes (self, *args, **kwargs) and returns details.
    """

    @functools.wraps(original_method)
    def wrapper(self: _SelfT, *args: _P.args, **kwargs: _P.kwargs) -> _T:
        target = getattr(self, target_attr)
        journal.record(target, operation, details=details_func(self, *args, **kwargs))
        return original_method(self, *args, **kwargs)

    return wrapper


# =============================================================================
# Store and restore original methods
# =============================================================================


def get_original_methods() -> dict[str, Any]:
    """Obtain original methods for later restoration.

    Returns:
        A dictionary mapping method names to their original implementations.
    """
    original_methods = {
        # TensorBase
        "TensorBase.__init__": _core.TensorBase.__init__,
        # Node
        "Node.__init__": _core.Node.__init__,
        "Node.name.fset": _core.Node.name.fset,
        "Node.domain.fset": _core.Node.domain.fset,
        "Node.version.fset": _core.Node.version.fset,
        "Node.op_type.fset": _core.Node.op_type.fset,
        "Node.overload.fset": _core.Node.overload.fset,
        "Node.resize_inputs": _core.Node.resize_inputs,
        "Node.prepend": _core.Node.prepend,
        "Node.append": _core.Node.append,
        "Node.resize_outputs": _core.Node.resize_outputs,
        "Node.graph.fset": _core.Node.graph.fset,
        # Value
        "Value.__init__": _core.Value.__init__,
        "Value.name.fset": _core.Value.name.fset,
        "Value.type.fset": _core.Value.type.fset,
        "Value.shape.fset": _core.Value.shape.fset,
        "Value.const_value.fset": _core.Value.const_value.fset,
        "Value.replace_all_uses_with": _core.Value.replace_all_uses_with,
        "Value.merge_shapes": _core.Value.merge_shapes,
        # Graph
        "Graph.__init__": _core.Graph.__init__,
        "Graph.register_initializer": _core.Graph.register_initializer,
        "Graph.append": _core.Graph.append,
        "Graph.extend": _core.Graph.extend,
        "Graph.remove": _core.Graph.remove,
        "Graph.insert_after": _core.Graph.insert_after,
        "Graph.insert_before": _core.Graph.insert_before,
        "Graph.sort": _core.Graph.sort,
        # Model
        "Model.__init__": _core.Model.__init__,
        # Function
        "Function.__init__": _core.Function.__init__,
        "Function.name.fset": _core.Function.name.fset,
        "Function.domain.fset": _core.Function.domain.fset,
        "Function.overload.fset": _core.Function.overload.fset,
        # Attr
        "Attr.__init__": _core.Attr.__init__,
        # _GraphIO (GraphInputs/GraphOutputs)
        "_GraphIO.append": _graph_containers._GraphIO.append,
        "_GraphIO.extend": _graph_containers._GraphIO.extend,
        "_GraphIO.insert": _graph_containers._GraphIO.insert,
        "_GraphIO.pop": _graph_containers._GraphIO.pop,
        "_GraphIO.remove": _graph_containers._GraphIO.remove,
        "_GraphIO.clear": _graph_containers._GraphIO.clear,
        "_GraphIO.__setitem__": _graph_containers._GraphIO.__setitem__,
        # GraphInitializers
        "GraphInitializers.__setitem__": _graph_containers.GraphInitializers.__setitem__,
        "GraphInitializers.__delitem__": _graph_containers.GraphInitializers.__delitem__,
        # Attributes
        "Attributes.__setitem__": _graph_containers.Attributes.__setitem__,
    }

    return original_methods


def wrap_ir_classes(journal: _journaling.Journal) -> dict[str, Any]:
    """Wrap IR classes with journaling-enabled versions.

    This function replaces methods on IR classes with wrapped versions that
    record operations to the active journal.
    """
    original_methods = get_original_methods()

    # TensorBase
    _core.TensorBase.__init__ = _init_wrapper(
        journal, original_methods["TensorBase.__init__"], details_func=lambda self: None
    )

    # Node
    _core.Node.__init__ = _init_wrapper(journal, original_methods["Node.__init__"])
    _core.Node.name = property(
        _core.Node.name.fget,
        _setter_wrapper(journal, original_methods["Node.name.fset"], "_name", "set_name"),
    )
    _core.Node.domain = property(
        _core.Node.domain.fget,
        _setter_wrapper(
            journal, original_methods["Node.domain.fset"], "_domain", "set_domain"
        ),
    )
    _core.Node.version = property(
        _core.Node.version.fget,
        _setter_wrapper(
            journal, original_methods["Node.version.fset"], "_version", "set_version"
        ),
    )
    _core.Node.op_type = property(
        _core.Node.op_type.fget,
        _setter_wrapper(
            journal, original_methods["Node.op_type.fset"], "_op_type", "set_op_type"
        ),
    )
    _core.Node.overload = property(
        _core.Node.overload.fget,
        _setter_wrapper(
            journal, original_methods["Node.overload.fset"], "_overload", "set_overload"
        ),
    )
    _core.Node.resize_inputs = _method_wrapper(
        journal,
        original_methods["Node.resize_inputs"],
        "resize_inputs",
        details_func=lambda self, new_size: f"{len(self._inputs)} -> {new_size}",
    )
    _core.Node.prepend = _method_wrapper(
        journal,
        original_methods["Node.prepend"],
        "prepend",
        details_func=lambda self, nodes: repr(nodes),
    )
    _core.Node.append = _method_wrapper(
        journal,
        original_methods["Node.append"],
        "append",
        details_func=lambda self, nodes: repr(nodes),
    )
    _core.Node.resize_outputs = _method_wrapper(
        journal,
        original_methods["Node.resize_outputs"],
        "resize_outputs",
        details_func=lambda self, new_size: f"{len(self._outputs)} -> {new_size}",
    )
    _core.Node.graph = property(
        _core.Node.graph.fget,
        _method_wrapper(
            journal,
            original_methods["Node.graph.fset"],
            "set_graph",
            details_func=lambda self, value: (
                f"{(value.name if isinstance(value, _core.Graph) else value)!r}"
            ),
        ),
    )

    # Value
    _core.Value.__init__ = _init_wrapper(journal, original_methods["Value.__init__"])
    _core.Value.name = property(
        _core.Value.name.fget,
        _setter_wrapper(journal, original_methods["Value.name.fset"], "_name", "set_name"),
    )
    _core.Value.type = property(
        _core.Value.type.fget,
        _setter_wrapper(journal, original_methods["Value.type.fset"], "_type", "set_type"),
    )
    _core.Value.shape = property(
        _core.Value.shape.fget,
        _setter_wrapper(journal, original_methods["Value.shape.fset"], "_shape", "set_shape"),
    )
    _core.Value.const_value = property(
        _core.Value.const_value.fget,
        _setter_wrapper(
            journal,
            original_methods["Value.const_value.fset"],
            "_const_value",
            "set_const_value",
        ),
    )
    _core.Value.replace_all_uses_with = _method_wrapper(
        journal,
        original_methods["Value.replace_all_uses_with"],
        "replace_all_uses_with",
        details_func=lambda self, replacement, replace_graph_outputs=False: (
            f"replacement={replacement!r}, replace_graph_outputs={replace_graph_outputs}"
        ),
    )
    _core.Value.merge_shapes = _method_wrapper(
        journal,
        original_methods["Value.merge_shapes"],
        "merge_shapes",
        details_func=lambda self, other: f"original={self._shape!r}, other={other!r}",
    )

    # Graph
    _core.Graph.__init__ = _init_wrapper(
        journal, original_methods["Graph.__init__"], details_func=lambda self: str(self.name)
    )
    _core.Graph.register_initializer = _method_wrapper(
        journal,
        original_methods["Graph.register_initializer"],
        "register_initializer",
        details_func=lambda self, value: repr(value),
    )
    _core.Graph.append = _method_wrapper(
        journal,
        original_methods["Graph.append"],
        "append",
        details_func=lambda self, node: repr(node),
    )
    _core.Graph.extend = _method_wrapper(
        journal,
        original_methods["Graph.extend"],
        "extend",
        details_func=lambda self, nodes: repr(nodes),
    )
    _core.Graph.remove = _method_wrapper(
        journal,
        original_methods["Graph.remove"],
        "remove",
        details_func=lambda self, nodes, safe=False: f"nodes={nodes!r}, safe={safe}",
    )
    _core.Graph.insert_after = _method_wrapper(
        journal,
        original_methods["Graph.insert_after"],
        "insert_after",
        details_func=lambda self, node, new_nodes: f"node={node!r}, new_nodes={new_nodes!r}",
    )
    _core.Graph.insert_before = _method_wrapper(
        journal,
        original_methods["Graph.insert_before"],
        "insert_before",
        details_func=lambda self, node, new_nodes: f"node={node!r}, new_nodes={new_nodes!r}",
    )
    _core.Graph.sort = _method_wrapper(
        journal,
        original_methods["Graph.sort"],
        "sort",
        details_func=lambda self: None,
    )

    # Model
    _core.Model.__init__ = _init_wrapper(journal, original_methods["Model.__init__"])

    # Function
    _core.Function.__init__ = _init_wrapper(journal, original_methods["Function.__init__"])
    _core.Function.name = property(
        _core.Function.name.fget,
        _setter_wrapper(journal, original_methods["Function.name.fset"], "_name", "set_name"),
    )
    _core.Function.domain = property(
        _core.Function.domain.fget,
        _setter_wrapper(
            journal, original_methods["Function.domain.fset"], "_domain", "set_domain"
        ),
    )
    _core.Function.overload = property(
        _core.Function.overload.fget,
        _setter_wrapper(
            journal, original_methods["Function.overload.fset"], "_overload", "set_overload"
        ),
    )

    # Attr
    _core.Attr.__init__ = _init_wrapper(journal, original_methods["Attr.__init__"])

    # _GraphIO (GraphInputs/GraphOutputs)
    _graph_containers._GraphIO.append = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.append"],
        "append_io",
        target_attr="_graph",
        details_func=lambda self, item: f"[{self.__class__.__name__}] {item!r}",
    )
    _graph_containers._GraphIO.extend = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.extend"],
        "extend_io",
        target_attr="_graph",
        details_func=lambda self, other: f"[{self.__class__.__name__}] {other!r}",
    )
    _graph_containers._GraphIO.insert = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.insert"],
        "insert_io",
        target_attr="_graph",
        details_func=lambda self, i, item: f"[{self.__class__.__name__}] {item!r}",
    )
    _graph_containers._GraphIO.pop = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.pop"],
        "pop_io",
        target_attr="_graph",
        details_func=lambda self, i=-1: f"[{self.__class__.__name__}] index={i}",
    )
    _graph_containers._GraphIO.remove = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.remove"],
        "remove_io",
        target_attr="_graph",
        details_func=lambda self, item: f"[{self.__class__.__name__}] {item!r}",
    )
    _graph_containers._GraphIO.clear = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.clear"],
        "clear_io",
        target_attr="_graph",
        details_func=lambda self: f"[{self.__class__.__name__}]",
    )
    _graph_containers._GraphIO.__setitem__ = _container_method_wrapper(
        journal,
        original_methods["_GraphIO.__setitem__"],
        "set_io",
        target_attr="_graph",
        details_func=lambda self, i, item: (
            f"[{self.__class__.__name__}] index={i}, item={item!r}"
        ),
    )

    # GraphInitializers
    _graph_containers.GraphInitializers.__setitem__ = _container_method_wrapper(
        journal,
        original_methods["GraphInitializers.__setitem__"],
        "set_initializer",
        target_attr="_graph",
        details_func=lambda self, key, value: f"key={key!r}, value={value!r}",
    )
    _graph_containers.GraphInitializers.__delitem__ = _container_method_wrapper(
        journal,
        original_methods["GraphInitializers.__delitem__"],
        "delete_initializer",
        target_attr="_graph",
        details_func=lambda self, key: f"key={key!r}",
    )

    # Attributes
    _graph_containers.Attributes.__setitem__ = _container_method_wrapper(
        journal,
        original_methods["Attributes.__setitem__"],
        "set_attribute",
        target_attr="_owner",
        details_func=lambda self, key, value: f"key={key!r}, value={value!r}",
    )

    return original_methods


def restore_ir_classes(original_methods: dict[str, Any]) -> None:
    """Restore IR classes to their original implementations.

    This function undoes the wrapping done by wrap_ir_classes().
    """
    # TensorBase
    _core.TensorBase.__init__ = original_methods["TensorBase.__init__"]

    # Node
    _core.Node.__init__ = original_methods["Node.__init__"]
    _core.Node.name = property(
        _core.Node.name.fget,
        original_methods["Node.name.fset"],
    )
    _core.Node.domain = property(
        _core.Node.domain.fget,
        original_methods["Node.domain.fset"],
    )
    _core.Node.version = property(
        _core.Node.version.fget,
        original_methods["Node.version.fset"],
    )
    _core.Node.op_type = property(
        _core.Node.op_type.fget,
        original_methods["Node.op_type.fset"],
    )
    _core.Node.overload = property(
        _core.Node.overload.fget,
        original_methods["Node.overload.fset"],
    )
    _core.Node.resize_inputs = original_methods["Node.resize_inputs"]
    _core.Node.prepend = original_methods["Node.prepend"]
    _core.Node.append = original_methods["Node.append"]
    _core.Node.resize_outputs = original_methods["Node.resize_outputs"]
    _core.Node.graph = property(
        _core.Node.graph.fget,
        original_methods["Node.graph.fset"],
    )

    # Value
    _core.Value.__init__ = original_methods["Value.__init__"]
    _core.Value.name = property(
        _core.Value.name.fget,
        original_methods["Value.name.fset"],
    )
    _core.Value.type = property(
        _core.Value.type.fget,
        original_methods["Value.type.fset"],
    )
    _core.Value.shape = property(
        _core.Value.shape.fget,
        original_methods["Value.shape.fset"],
    )
    _core.Value.const_value = property(
        _core.Value.const_value.fget,
        original_methods["Value.const_value.fset"],
    )
    _core.Value.replace_all_uses_with = original_methods["Value.replace_all_uses_with"]
    _core.Value.merge_shapes = original_methods["Value.merge_shapes"]

    # Graph
    _core.Graph.__init__ = original_methods["Graph.__init__"]
    _core.Graph.register_initializer = original_methods["Graph.register_initializer"]
    _core.Graph.append = original_methods["Graph.append"]
    _core.Graph.extend = original_methods["Graph.extend"]
    _core.Graph.remove = original_methods["Graph.remove"]
    _core.Graph.insert_after = original_methods["Graph.insert_after"]
    _core.Graph.insert_before = original_methods["Graph.insert_before"]
    _core.Graph.sort = original_methods["Graph.sort"]

    # Model
    _core.Model.__init__ = original_methods["Model.__init__"]

    # Function
    _core.Function.__init__ = original_methods["Function.__init__"]
    _core.Function.name = property(
        _core.Function.name.fget,
        original_methods["Function.name.fset"],
    )
    _core.Function.domain = property(
        _core.Function.domain.fget,
        original_methods["Function.domain.fset"],
    )
    _core.Function.overload = property(
        _core.Function.overload.fget,
        original_methods["Function.overload.fset"],
    )

    # Attr
    _core.Attr.__init__ = original_methods["Attr.__init__"]

    # _GraphIO (GraphInputs/GraphOutputs)
    _graph_containers._GraphIO.append = original_methods["_GraphIO.append"]
    _graph_containers._GraphIO.extend = original_methods["_GraphIO.extend"]
    _graph_containers._GraphIO.insert = original_methods["_GraphIO.insert"]
    _graph_containers._GraphIO.pop = original_methods["_GraphIO.pop"]
    _graph_containers._GraphIO.remove = original_methods["_GraphIO.remove"]
    _graph_containers._GraphIO.clear = original_methods["_GraphIO.clear"]
    _graph_containers._GraphIO.__setitem__ = original_methods["_GraphIO.__setitem__"]

    # GraphInitializers
    _graph_containers.GraphInitializers.__setitem__ = original_methods[
        "GraphInitializers.__setitem__"
    ]
    _graph_containers.GraphInitializers.__delitem__ = original_methods[
        "GraphInitializers.__delitem__"
    ]

    # Attributes
    _graph_containers.Attributes.__setitem__ = original_methods["Attributes.__setitem__"]
