from __future__ import annotations
from typing import Any
import numpy as np

from ..bpy_ns_types import NfBoolProperty, NfEnumProperty, notify_tree_changed
from ....registry import register_class
from .base import SliceNode, SlicingCollection, get_all_children, SliceNodeJoiner, SlicingGroup, SlicingObject

@register_class
class SLICENODE_GetMesh(SliceNode):
    bl_label = "Get Mesh"

    include_children: NfBoolProperty(name="Include Children", default=True) # pyright: ignore[reportInvalidTypeForm]

    def init(self, context):
        self.add_input("SLICESOCK_Object")
        self.add_output("SLICESOCK_Mesh")

    def draw_buttons(self, context, layout):
        assert layout
        layout.prop(self, "include_children")

    def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection | None:
        if socket_name != "Mesh": return None

        obj: Any = self.resolve("Object")
        if not obj: return None

        family = [obj] + (get_all_children(obj) if self.include_children else [])

        sc: SlicingCollection = SlicingCollection(family, obj.name)
        
        return sc

@register_class
class SLICENODE_JoinMesh(SliceNodeJoiner):
    bl_label = "Join Mesh Parts"
    socket_type = "SLICESOCK_Mesh"

    def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection:
        from ....infra.blender.blender_slice_nodes_generic import resolve_input
        merged_coll = resolve_input(self.inputs[0]) or SlicingCollection([],'')
        
        for input in self.inputs[1:]:
            coll = resolve_input(input)
            if coll: 
                merged_coll.objects += coll.objects

        return merged_coll

@register_class
class SLICENODE_JoinMeshGroups(SliceNodeJoiner):
    bl_label = "Join Mesh Groups"
    socket_type = "SLICESOCK_Mesh"

    def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingGroup:
        from ....infra.blender.blender_slice_nodes_generic import resolve_input
        group: SlicingGroup = SlicingGroup([])
        
        for input in self.inputs:
            coll = resolve_input(input)
            if coll: 
                group.collections[coll.objects[0].name] = coll

        return group

@register_class
class SLICENODE_SeparateMesh(SliceNode):
    bl_label = "Separate Mesh"

    comparison: NfEnumProperty(
        name="Comparison",
        items=[
            ('EQUAL', "Equal", "=="),
            ('NOT_EQUAL', "Not Equal", "!="),
            ('GREATER_THAN', "Greater Than", ">"),
            ('LESS_THAN', "Less Than", "<"),
            ('GREATER_EQUAL', "Greater Equal", ">="),
            ('LESS_EQUAL', "Less Equal", "<="),
        ],
        default='EQUAL',
        update=notify_tree_changed
    ) # pyright: ignore[reportInvalidTypeForm]

    def init(self, context):
        self.add_input("SLICESOCK_Mesh")
        self.add_input("SLICESOCK_String", "Attribute")
        self.add_input("SLICESOCK_Float", "Value")
        self.add_output("SLICESOCK_Mesh", "Selection")
        self.add_output("SLICESOCK_Mesh", "Inverted")

    def draw_buttons(self, context, layout):
        assert layout
        layout.prop(self, "comparison", text="")

    def evaluate_output(self, socket_name: str, **_kwargs) -> SlicingCollection | None:
        coll: SlicingCollection | None = self.resolve("Mesh")
        attr_name: str = self.resolve("Attribute")
        value: float = self.resolve("Value")

        if not coll: return None

        res_selection = SlicingCollection([], coll.name + "_selection")
        res_inverted = SlicingCollection([], coll.name + "_inverted")

        for so in coll.objects:
            if attr_name not in so.attributes:
                # If attribute missing, entire object goes to inverted
                res_inverted.objects.append(so)
                continue
            
            attr = so.attributes[attr_name]
            vals = attr["values"]
            
            mask = np.zeros(len(so.mesh), dtype=bool)
            if self.comparison == 'EQUAL': mask = vals == value
            elif self.comparison == 'NOT_EQUAL': mask = vals != value
            elif self.comparison == 'GREATER_THAN': mask = vals > value
            elif self.comparison == 'LESS_THAN': mask = vals < value
            elif self.comparison == 'GREATER_EQUAL': mask = vals >= value
            elif self.comparison == 'LESS_EQUAL': mask = vals <= value
            
            if np.all(mask):
                res_selection.objects.append(so)
            elif not np.any(mask):
                res_inverted.objects.append(so)
            else:
                # Split the object
                so_sel = SlicingObject()
                so_sel.name = so.name + "_sel"
                so_sel.parent = so.parent
                so_sel.object_type = so.object_type
                so_sel.extruder = so.extruder
                so_sel.mesh = so.mesh[mask]
                for k, v in so.attributes.items():
                    so_sel.attributes[k] = {
                        "domain": v["domain"],
                        "type": v["type"],
                        "values": v["values"][mask]
                    }
                res_selection.objects.append(so_sel)

                so_inv = SlicingObject()
                so_inv.name = so.name + "_inv"
                so_inv.parent = so.parent
                so_inv.object_type = so.object_type
                so_inv.extruder = so.extruder
                so_inv.mesh = so.mesh[~mask]
                for k, v in so.attributes.items():
                    so_inv.attributes[k] = {
                        "domain": v["domain"],
                        "type": v["type"],
                        "values": v["values"][~mask]
                    }
                res_inverted.objects.append(so_inv)

        if socket_name == "Selection":
            return res_selection
        return res_inverted