from __future__ import annotations
from dataclasses import field
import typing
import bpy
from contextlib import contextmanager

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

class NodeException(Exception): ...

_EvalCache = dict[tuple[int, str, 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
    bl_color = (1.0, 1.0, 1.0, 1.0)

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

from ...registry import get_label_by_name
class SliceNode(Node, AutoBlIdName):
    def update(self) -> None: 
        self._update_multi_groups()

    def evaluate_output(self, socket_name: str, **_kwargs) -> Any: ...

    def resolve(self, fieldname) -> Any:
        cache = getattr(self, "_current_cache", None)
        visited = getattr(self, "_current_visited", None)
        return resolve_input(self.inputs.get(fieldname), cache=cache, visited=visited)

    def add_output(self, bl_idname: str, name: str = "", identifier: str = "", hide: bool = False):
        eval_name = name or get_label_by_name(bl_idname)
        sock = self.outputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
        sock.hide = hide
        return sock

    def add_input(self, bl_idname: str, name: str = "", identifier: str = "", hide: bool = False):
        eval_name = name or get_label_by_name(bl_idname)
        sock = self.inputs.new(bl_idname, eval_name, identifier=identifier or eval_name)
        sock.hide = hide
        return sock

    def add_multi_socket_group(self, in_type: str, out_type: str = "", name: str = "", in_name: str = "", out_name: str = ""):
        if not self.get('multi_groups'): self['multi_groups'] = {}
        group_id = name or in_name or out_name or in_type
        self['multi_groups'][group_id] = {
            'in_type': in_type,
            'out_type': out_type,
            'in_name': in_name or name or get_label_by_name(in_type),
            'out_name': out_name or name or (get_label_by_name(out_type) if out_type else ""),
        }
        self._update_multi_groups()
    
    def resolve_multi_input(self, fieldname) -> Any:
        cache = getattr(self, "_current_cache", None)
        visited = getattr(self, "_current_visited", None)
        # Search by name or identifier prefix for multi-groups
        inputs = [i for i in list(self.inputs) if i.name == fieldname or i.identifier.startswith(f"{fieldname}_ms_in_")]
        return [resolve_input(r, cache=cache, visited=visited) for r in inputs]

    def _update_multi_groups(self):
        if not self.get('multi_groups'): return
        
        for group_id, cfg in self['multi_groups'].items():
            in_type = cfg['in_type']
            out_type = cfg.get('out_type', '')
            in_name = cfg['in_name']
            out_name = cfg['out_name']
            
            def get_ins(): return [i for i in self.inputs if i.identifier.startswith(f"{group_id}_ms_in_")]
            def get_outs(): return [o for o in self.outputs if o.identifier.startswith(f"{group_id}_ms_out_")]
            def is_active(s): return s.is_linked or bool(getattr(s, 'value', False))

            ins = get_ins()
            
            if not ins:
                self.add_input(in_type, in_name, identifier=f"{group_id}_ms_in_0")
                if out_type:
                    self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_0")
                ins = get_ins()

            if is_active(ins[-1]):
                idx = len(ins)
                self.add_input(in_type, in_name, identifier=f"{group_id}_ms_in_{idx}")
                if out_type:
                    self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_{idx}")
                ins = get_ins()

            while len(ins) > 1:
                if not is_active(ins[-1]) and not is_active(ins[-2]):
                    self.inputs.remove(ins[-1])
                    outs = get_outs()
                    if out_type and len(outs) >= len(ins):
                        self.outputs.remove(outs[-1])
                    ins.pop()
                    continue
                break
    
## Evaluation helpers

_current_execution_cache: _EvalCache | None = None
_feedback_enabled: bool = True

@contextmanager
def execution_cache():
    global _current_execution_cache
    if _current_execution_cache is not None:
        yield _current_execution_cache
        return

    _current_execution_cache = {}
    try:
        yield _current_execution_cache
    finally:
        _current_execution_cache = None

@contextmanager
def feedback_disabled():
    global _feedback_enabled
    old = _feedback_enabled
    _feedback_enabled = False
    try:
        yield
    finally:
        _feedback_enabled = old

@contextmanager
def node_execution_feedback(node: bpy.types.Node):
    if not _feedback_enabled:
        with execution_cache():
            yield
        return

    # Save original color
    old_use_color = node.use_custom_color
    old_color = node.color[:]
    
    node.use_custom_color = True
    node.color = (75/255, 73/255, 19/255) # Green
    
    # Refresh UI
    bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
    try:
        with execution_cache():
            yield
    finally:
        node.use_custom_color = old_use_color
        node.color = old_color
        bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)

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 cache is None:
        cache = getattr(socket.node, "_current_cache", None) or _current_execution_cache
    if visited is None:
        visited = getattr(socket.node, "_current_visited", None)

    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,
):
    if cache is None:
        cache = _current_execution_cache

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

    assert socket.node
    # Use identifier instead of name for the cache key, as multiple sockets 
    # (especially in multi-groups) can share the same name but must have unique identifiers.
    key = (socket.node.as_pointer(), socket.node.name, socket.identifier)
    if key in cache: return cache[key]

    node_key = (socket.node.as_pointer(), socket.node.name)
    if node_key in visited: return default  # cycle guard

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

            # Set these temporarily so self.resolve() can find them
            socket.node._current_cache = cache # type: ignore
            socket.node._current_visited = visited # type: ignore

            try:
                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, 
                    identifier=socket.identifier, 
                    socket=socket, 
                    cache=cache, 
                    visited=visited
                )
                cache[key] = value
                return value
            finally:
                socket.node._current_cache = None # type: ignore
                socket.node._current_visited = None # type: ignore
    finally:
        visited.remove(node_key)