from __future__ import annotations
import typing

from typing import Any
from bpy.types import NodeSocket, Node

class NodeException(Exception): ...

_EvalCache = dict[tuple[int, str], Any]

class AutoBlIdName:
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)
        if "bl_idname" not in cls.__dict__:
            cls.bl_idname = cls.__name__

class SliceNodeSocket(NodeSocket, AutoBlIdName):
    default: Any = None
    color = (1., 1., 1., 1.)

    def draw_color(self, context, node) -> tuple[float, float, float, float]: return self.color # pyright: ignore[reportIncompatibleMethodOverride]

class SliceNode(Node, AutoBlIdName):
    def evaluate_output(self, socket_name: str, **_kwargs) -> Any: ...

    def add_output(self, bl_idname: str, name: str | None = None):
        from ..registry import get_label_by_name
        return self.outputs.new(bl_idname, name or get_label_by_name(bl_idname))

    def add_input(self, bl_idname: str, name: str | None = None, multi: bool = False):
        from ..registry import get_label_by_name
        sock = self.inputs.new(bl_idname, name or get_label_by_name(bl_idname))
        if multi:
            sock.link_limit = 0
            sock.display_shape = 'LIST'
        return sock

## Evaluation helpers

def _follow_reroute(socket: NodeSocket) -> NodeSocket:
    current: NodeSocket = socket
    seen: set[int] = set()

    while getattr(current.node, "bl_idname", "") == "NodeReroute":
        assert current.node
        ptr = current.node.as_pointer()
        if ptr in seen:
            break
        seen.add(ptr)

        reroute_input = current.node.inputs[0] if current.node.inputs else None
        if not reroute_input or not reroute_input.is_linked:
            raise NodeException("Loose reroute node found: ensure all reroute nodes are correctly linked")

        current = reroute_input.links[0].from_socket

    return current

def resolve_input(
    socket: SliceNodeSocket | None,
    *,
    cache: _EvalCache | None = None,
    visited: set[tuple[int, str]] | None = None,
):
    if socket is None:
        raise Exception(f"The queried socket is non-existant")

    assert socket.node

    if not socket.is_linked:
        val = getattr(socket, "value", None)
        if val is not None:
            return val
        return socket.default

    assert socket.is_linked
    assert socket.links
    from_socket: NodeSocket = _follow_reroute(socket.links[0].from_socket)

    return _resolve_output(
        from_socket,
        default=None,
        cache=cache,
        visited=visited,
    )

def _resolve_output(
    socket: NodeSocket,
    *,
    default: Any = None,
    cache: _EvalCache | None = None,
    visited: set[tuple[int, str]] | None = None,
):

    cache = {} if cache is None else cache
    visited = set() if visited is None else visited

    assert socket.node
    key = (socket.node.as_pointer(), socket.node.name)
    if key in cache: return cache[key]
    if key in visited: return default  # cycle guard

    visited.add(key)
    try:
        if not isinstance(socket.node, SliceNode):
            return default

        if getattr(socket.node, "mute", False):
            # Attempt bypass: find a matching input socket by name or type
            input_to_bypass = socket.node.inputs.get(socket.name)
            if not input_to_bypass:
                for inp in socket.node.inputs:
                    if inp.bl_idname == socket.bl_idname:
                        input_to_bypass = inp
                        break
            
            if input_to_bypass:
                value = resolve_input(typing.cast(Any, input_to_bypass), cache=cache, visited=visited)
                cache[key] = value
                return value
            return default

        assert socket.node.evaluate_output
        value = socket.node.evaluate_output(socket.name, cache=cache, visited=visited)
        cache[key] = value
        return value
    finally:
        visited.remove(key)