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]

NATIVE_SOCKET_MAP: dict[str, str] = {}

NATIVE_SOCKET_LABELS: dict[str, str] = {
    "NodeSocketFloat": "Float",
    "NodeSocketInt": "Integer",
    "NodeSocketVector": "Vector",
    "NodeSocketRotation": "Rotation",
    "NodeSocketString": "String",
    "NodeSocketObject": "Object",
}

if hasattr(bpy.types, "NodeSocket") and not hasattr(bpy.types.NodeSocket, "value"):
    def _get_socket_value(self):
        try:
            return self.default_value
        except (AttributeError, TypeError):
            pass
        return getattr(self, "_value", getattr(self, "default", None))

    def _set_socket_value(self, val):
        try:
            self.default_value = val
            return
        except (AttributeError, TypeError):
            pass
        self._value = val

    try:
        bpy.types.NodeSocket.value = property(_get_socket_value, _set_socket_value)
    except Exception:
        pass

class AutoBlIdName:
    def __init_subclass__(cls, **kwargs):
        if kwargs:
            try:
                super().__init_subclass__(**kwargs)
            except TypeError:
                super().__init_subclass__()
        else:
            super().__init_subclass__()
        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, register_class
class SliceNode(Node, AutoBlIdName):
    def get(self, key: str, default: Any = None) -> Any:
        try:
            return super().get(key, default)  # type: ignore
        except (AttributeError, TypeError):
            pass
        if hasattr(self, "_data") and isinstance(self._data, dict):
            return self._data.get(key, default)
        return getattr(self, key, default)

    def __setitem__(self, key: str, value: Any) -> None:
        try:
            super().__setitem__(key, value)  # type: ignore
        except (AttributeError, TypeError):
            if not hasattr(self, "_data") or self._data is None:
                self._data = {}
            self._data[key] = value

    def __getitem__(self, key: str) -> Any:
        try:
            return super().__getitem__(key)  # type: ignore
        except (AttributeError, TypeError):
            if hasattr(self, "_data") and isinstance(self._data, dict):
                return self._data[key]
            raise KeyError(key)

    def __contains__(self, key: str) -> bool:
        try:
            return super().__contains__(key)  # type: ignore
        except (AttributeError, TypeError):
            if hasattr(self, "_data") and isinstance(self._data, dict):
                return key in self._data
            return False

    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)
        inputs = getattr(self, "inputs", None)
        if inputs is None:
            return None
        return resolve_input(inputs.get(fieldname), cache=cache, visited=visited)

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

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

    def add_multi_socket_group(self, in_type: str, out_type: str = "", name: str = "", in_name: str = "", out_name: str = "", in_display_shape: str = "", out_display_shape: 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 ""),
            'in_display_shape': in_display_shape,
            'out_display_shape': out_display_shape,
        }
        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']
            in_display_shape = cfg.get('in_display_shape', '')
            out_display_shape = cfg.get('out_display_shape', '')
            
            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):
                if getattr(s, 'is_linked', False):
                    return True
                val = getattr(s, 'value', None)
                if val is None or val == "":
                    val = getattr(s, 'default_value', None)
                return bool(val)

            ins = get_ins()
            
            if not ins:
                self.add_input(in_type, in_name, identifier=f"{group_id}_ms_in_0", display_shape=in_display_shape)
                if out_type:
                    self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_0", display_shape=out_display_shape)
                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}", display_shape=in_display_shape)
                if out_type:
                    self.add_output(out_type, out_name, identifier=f"{group_id}_ms_out_{idx}", display_shape=out_display_shape)
                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 = getattr(node, "use_custom_color", False)
    old_color = getattr(node, "color", (0.0, 0.0, 0.0))[:] if hasattr(node, "color") else None
    
    if hasattr(node, "use_custom_color"):
        node.use_custom_color = True
    if hasattr(node, "color"):
        node.color = (75/255, 73/255, 19/255) # Green
    
    # Refresh UI
    try:
        bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
    except Exception:
        pass
    try:
        with execution_cache():
            yield
    finally:
        if hasattr(node, "use_custom_color"):
            node.use_custom_color = old_use_color
        if hasattr(node, "color") and old_color is not None:
            node.color = old_color
        try:
            bpy.ops.wm.redraw_timer(type='DRAW_WIN_SWAP', iterations=1)
        except Exception:
            pass

@contextmanager
def _group_context(node: bpy.types.Node, cache: _EvalCache | None = None):
    stack = _get_group_stack(cache)
    stack.append(node)
    try:
        yield
    finally:
        stack.pop()

def _get_group_stack(cache: _EvalCache | None = None) -> list[bpy.types.Node]:
    if cache is None:
        cache = _current_execution_cache
    if cache is None: return [] 
    if '_group_stack' not in cache:
        cache['_group_stack'] = []
    return cache['_group_stack']

def _get_current_group_node(cache: _EvalCache | None = None) -> bpy.types.Node | None:
    stack = _get_group_stack(cache)
    return stack[-1] if stack else None

@register_class
class SLICENODE_Group(bpy.types.NodeCustomGroup, AutoBlIdName):
    bl_label = "Group"
    bl_icon = 'NODETREE'
    
    def _tree_updated(self, context):
        try:
            if not self.name: return
            _ = len(self.inputs)
            _ = len(self.outputs)
        except Exception:
            return
            
        self.sync_sockets()

    slicer_tree: bpy.props.PointerProperty(  # pyright: ignore[reportInvalidTypeForm]
        type=bpy.types.NodeTree,
        name="Node Tree",
        update=_tree_updated
    )

    def sync_sockets(self):
        try:
            if not self.name: return
            _ = len(self.inputs)
            _ = len(self.outputs)
        except Exception:
            return
            
        if not self.slicer_tree:
            self.inputs.clear()
            self.outputs.clear()
            return
            
        # 1. Collect expected sockets from interface
        expected_inputs = [] # (identifier, name, type)
        expected_outputs = []
        
        has_interface = hasattr(self.slicer_tree, "interface")
        if has_interface:
            # In Blender 4.0+, the collection of items is items_tree
            items_tree = getattr(self.slicer_tree.interface, "items_tree", [])
            for item in items_tree:
                if getattr(item, "item_type", "") != 'SOCKET': continue
                target_list = expected_inputs if item.in_out == 'INPUT' else expected_outputs
                target_list.append((item.identifier, item.name, item.socket_type))
        else:
            # Fallback for older Blender API without interface
            for sock in getattr(self.slicer_tree, "inputs", []):
                expected_inputs.append((sock.identifier, sock.name, sock.bl_idname))
            for sock in getattr(self.slicer_tree, "outputs", []):
                expected_outputs.append((sock.identifier, sock.name, sock.bl_idname))

        # 2. Sync Inputs
        self._sync_socket_collection(self.inputs, expected_inputs)
        
        # 3. Sync Outputs
        self._sync_socket_collection(self.outputs, expected_outputs)

    def _sync_socket_collection(self, collection, expected):
        # Remove sockets that shouldn't be there
        expected_ids = {e[0] for e in expected}
        to_remove = [s for s in collection if s.identifier not in expected_ids]
                
        for sock in to_remove:
            collection.remove(sock)
            
        # Add or update sockets
        sock_indexes = {s.identifier: i for i, s in enumerate(collection)}
        for ident, name, stype in expected:
            stype = NATIVE_SOCKET_MAP.get(stype, stype)
            if ident not in sock_indexes:
                try:
                    # Explicitly pass identifier to keep it perfectly synced
                    sock = collection.new(stype, name, identifier=ident)
                    sock_indexes[sock.identifier] = len(collection) - 1
                except Exception as e:
                    print(f"Failed to create socket {name} ({stype}): {e}")
            else:
                # Update name if changed
                sock = collection[sock_indexes[ident]]
                if sock.name != name:
                    sock.name = name

    def draw_buttons(self, context, layout):
        assert layout
        row = layout.row(align=True)
        row.prop(self, "slicer_tree", text="")
        if self.slicer_tree:
            op = row.operator("slicenode.group_toggle", text="", icon='NODETREE')
            op.node_name = self.name

    def draw_label(self):
        if self.slicer_tree:
            return self.slicer_tree.name
        return "Group"


def _resolve_group_output(socket: NodeSocket, cache: _EvalCache | None, visited: set[tuple[int, str]] | None):
    node = socket.node
    assert isinstance(node, (bpy.types.NodeGroup, SLICENODE_Group))
    
    if isinstance(node, SLICENODE_Group):
        tree = node.slicer_tree
    else:
        tree = node.node_tree
        
    if not tree: return None
    
    try:
        idx = list(node.outputs).index(socket)
    except (ValueError, IndexError):
        return None

    # Find the NodeGroupOutput node that has this input linked, or the first one
    output_node = None
    for n in tree.nodes:
        if isinstance(n, bpy.types.NodeGroupOutput):
            # NodeGroupOutput inputs usually match the group outputs by index
            if idx < len(n.inputs) and n.inputs[idx].is_linked:
                output_node = n
                break
    
    if not output_node:
        output_node = next((n for n in tree.nodes if isinstance(n, bpy.types.NodeGroupOutput)), None)
    
    if not output_node or idx >= len(output_node.inputs): return None
    
    inner_socket = output_node.inputs[idx]

    with _group_context(node, cache=cache):
        return resolve_input(typing.cast(Any, inner_socket), cache=cache, visited=visited)

def _resolve_group_input(socket: NodeSocket, cache: _EvalCache | None, visited: set[tuple[int, str]] | None):
    stack = _get_group_stack(cache)
    if not stack: return None
    group_node = stack.pop()
    
    # NodeGroupInput outputs usually match the group inputs by index
    try:
        idx = list(socket.node.outputs).index(socket)
        outer_socket = group_node.inputs[idx]
    except (ValueError, IndexError):
        stack.append(group_node)
        return None
        
    try:
        return resolve_input(typing.cast(Any, outer_socket), cache=cache, visited=visited)
    finally:
        stack.append(group_node)

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:
        try:
            val = getattr(socket, "value", None)
            if val is not None and val != "": return val
        except Exception:
            pass

        try:
            val = getattr(socket, "default_value", None)
            if val is not None: return val
        except Exception:
            pass

        return getattr(socket, "default", None)

    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
    # Include group stack in cache key to ensure independent evaluation for multiple instances
    group_stack = tuple(n.as_pointer() for n in _get_group_stack(cache))
    
    # 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 = (group_stack, socket.node.as_pointer(), socket.node.name, socket.identifier)
    if key in cache: return cache[key]

    node_key = (group_stack, 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):
                is_group = isinstance(socket.node, SLICENODE_Group) or (
                    isinstance(getattr(bpy.types, "NodeGroup", None), type) and isinstance(socket.node, bpy.types.NodeGroup)
                ) or getattr(socket.node, "bl_idname", "") in ("NodeGroup", "GeometryNodeGroup")
                if is_group:
                    return _resolve_group_output(socket, cache=cache, visited=visited)
                
                is_group_input = (
                    isinstance(getattr(bpy.types, "NodeGroupInput", None), type) and isinstance(socket.node, bpy.types.NodeGroupInput)
                ) or getattr(socket.node, "bl_idname", "") == "NodeGroupInput"
                if is_group_input:
                    return _resolve_group_input(socket, cache=cache, visited=visited)
                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)