from __future__ import annotations
from contextlib import ContextDecorator, contextmanager, nullcontext
from typing import Any

from .formatter import _format_node_description, get_node_input_names
from .proxy import NodeErrorReportingProxy


def attach_node_failure_info(node: Any, exc: Exception, fallback_inputs: list[str] | None = None) -> Exception:
    """Enriches an exception in-place with the failing node and input socket names, preventing duplicate wrapping."""
    if getattr(exc, "_nodeslicer_node", None) is not None:
        return exc
    try:
        setattr(exc, "_nodeslicer_node", node)
        input_names = get_node_input_names(node, fallback_inputs=fallback_inputs)
        setattr(exc, "_nodeslicer_inputs", input_names)

        node_str = _format_node_description(node) or "Unknown Node"
        inputs_str = ", ".join(input_names) if input_names else "(None)"
        prefix = f"Execution failed at node {node_str} [Inputs: {inputs_str}]: "

        if exc.args:
            orig_msg = str(exc.args[0])
            if not orig_msg.startswith("Execution failed at node "):
                exc.args = (prefix + orig_msg, *exc.args[1:])
        else:
            exc.args = (prefix + type(exc).__name__,)

        print(f"[NodeSlicer ERROR] {str(exc)}")
    except Exception:
        pass
    return exc


class NodeErrorHandler(ContextDecorator):
    """Context manager and decorator that intercepts exceptions in node execution
    and enriches them with node failure info (node name, inputs) without intruding
    into the node's internal logic."""

    def __init__(self, node: Any, fallback_inputs: list[str] | None = None):
        self.node = node
        self.fallback_inputs = fallback_inputs

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if exc_val is not None and isinstance(exc_val, Exception):
            attach_node_failure_info(self.node, exc_val, fallback_inputs=self.fallback_inputs)
        return False  # Do not suppress exception


def handle_node_errors(node: Any, fallback_inputs: list[str] | None = None) -> NodeErrorHandler:
    """Creates a modular context decorator to handle and enrich node execution errors."""
    return NodeErrorHandler(node, fallback_inputs=fallback_inputs)


@contextmanager
def node_execution_scope(
    node: Any,
    operator: Any = None,
    fallback_inputs: list[str] | None = None,
):
    """Modular, non-intrusive execution scope that combines profiling (if active),
    operator error reporting proxy (if operator provided), and exception enrichment
    without tightly coupling any of them into node classes."""
    try:
        from ..profiling import profile_node_execution
        prof_ctx = profile_node_execution(node)
    except Exception:
        prof_ctx = nullcontext()

    wrapped_op = NodeErrorReportingProxy(operator, node) if operator is not None else None

    class ExecutionScope:
        def __init__(self, op: Any):
            self.operator = op

    with prof_ctx:
        with handle_node_errors(node, fallback_inputs=fallback_inputs):
            yield ExecutionScope(wrapped_op)