from __future__ import annotations
from bpy.types import Object, Mesh
import bpy
import zlib, struct
from functools import cached_property

import numpy as np
from numpy import dtype, float64
from numpy.typing import NDArray

from typing import Any

from ...core.geometry import crc32_array
from .blender_bridge import get_all_children

TYPES_NAME = 'DEFAULT'

class MeshCaptureException(Exception): pass

class SlicingObject():
    name: str
    parent: str
    mesh: NDArray[float64]
    object_type: str = "ModelPart"
    modifiers: list[dict] = []
    extruder: int = 1

    def __init__(self, obj: Object | None = None, parent: str = "") -> None:
        if not bpy.context.scene: raise Exception('No scene currently open!')

        if obj: self.name = str(obj.name)
        self.parent = parent

        depsgraph = bpy.context.evaluated_depsgraph_get()
        assert bpy.context.scene.unit_settings
        scene_scale: float = bpy.context.scene.unit_settings.scale_length

        if ns := getattr(obj, 'nodeslicer', None): 
            self.object_type = ns.object_type
            self.extruder = ns.extruder
            self.modifiers = [{'param_id': m.field, 'param_value': m.value_str} for m in ns.modifiers]

        if obj:
            if obj.mode == 'EDIT':
                obj.update_from_editmode()
            eval_objects = obj.evaluated_get(depsgraph)
            self.mesh = objects_to_tris([eval_objects], 1000 * scene_scale)
        else:
            self.mesh = np.empty((0, 4, 3))

    def offset(self, offset: NDArray):
        self.mesh[:, :3, :] += offset

    def scale(self, factor: float):
        self.mesh[:, :3, :] *= factor

    def assign_modifiers(self, modifiers: dict[str, str]):
        self.modifiers = [
            {"param_id": k, "param_value": v}
            for k, v in modifiers.items()
        ]

    @property
    def checksum(self) -> int:
        import json

        buf = bytearray()
        buf.extend(struct.pack(">I", crc32_array(self.mesh)))
        buf.extend(struct.pack(">I", zlib.crc32(self.name.encode("utf-8"))))
        buf.extend(struct.pack(">I", zlib.crc32(self.parent.encode("utf-8"))))

        return zlib.crc32(buf) & 0xFFFFFFFF

    @property
    def height(self) -> float: return self.mesh[:, :3, 2].max()

    @property
    def min_x(self) -> float: return self.mesh[:, :3, 0].min()

    @property
    def max_x(self) -> float: return self.mesh[:, :3, 0].max()

    @property
    def min_y(self) -> float: return self.mesh[:, :3, 1].min()

    @property
    def max_y(self) -> float: return self.mesh[:, :3, 1].max()

    @property
    def min_z(self) -> float: return self.mesh[:, :3, 2].min()

    @property
    def max_z(self) -> float: return self.mesh[:, :3, 2].max()

    @property
    def min_xy(self) -> NDArray: return np.array([self.min_x, self.min_y, 0.0])

    @property
    def max_xy(self) -> NDArray: return np.array([self.max_x, self.max_y, 0.0])

    @property
    def center_xy(self) -> NDArray: return (self.min_xy + self.max_xy) / 2.0

class SlicingCollection():
    name: str
    objects: list[SlicingObject]

    def __init__(self, objs: list[Object], parent: str):
        self.objects = [SlicingObject(obj, parent) for obj in objs]
        self.objects = [o for o in self.objects if o.mesh.size]
        self.name = parent

    def offset(self, offset: NDArray):
        for so in self.objects: so.offset(offset)

    def scale(self, factor: float):
        for so in self.objects: so.scale(factor)

    @property
    def checksum(self) -> int:
        return crc32_array(np.array([o.checksum for o in self.objects]))

    @property
    def height(self) -> float:
        height_all = [so.height for so in self.objects if np.any(so.mesh)]
        if not height_all: return 0
        return max(height_all)

    @property
    def meshes(self) -> list[NDArray]:
        return [o.mesh for o in self.objects]

    @cached_property
    def mesh_lengths_ids(self) -> NDArray:
        return np.array([len(m) for m in self.meshes])

    @cached_property
    def mesh_start_ids(self) -> NDArray:
        lengths = self.mesh_lengths_ids
        return np.insert(np.cumsum(lengths)[:-1], 0, 0)

    @cached_property
    def mesh_end_ids(self) -> NDArray:
        starts = self.mesh_start_ids
        lengths = self.mesh_lengths_ids
        return starts + lengths - 1

    @cached_property
    def all_verts(self) -> NDArray:
        all_verts = np.vstack(self.meshes)[:, :3, :]
        return all_verts.reshape(-1, 3)

    @cached_property
    def unique_verts(self) -> tuple[np.ndarray, np.ndarray]:
        verts_list = []
        idxs = []
        offset = 0

        for mesh in self.meshes:
            verts = mesh[:, :3, :].reshape(-1, 3)

            void_dtype = np.dtype((np.void, verts.dtype.itemsize * verts.shape[1]))
            verts_void = verts.view(void_dtype).ravel()

            unique_voids, inv = np.unique(verts_void, return_inverse=True)

            unique_verts = unique_voids.view(verts.dtype).reshape(-1, 3)

            verts_list.append(unique_verts)
            idxs.append(inv.reshape(-1, 3) + offset)

            offset += unique_verts.shape[0]

        all_verts = np.concatenate(verts_list, axis=0)
        all_idxs  = np.vstack(idxs)
        return all_verts, all_idxs

    @property
    def min_x(self) -> float | None:
        return min([so.min_x for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def max_x(self) -> float | None:
        return max([so.max_x for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def min_y(self) -> float | None:
        return min([so.min_y for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def max_y(self) -> float | None:
        return max([so.max_y for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def min_z(self) -> float | None:
        return min([so.min_z for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def max_z(self) -> float | None:
        return max([so.max_z for so in self.objects if np.any(so.mesh)], default=None)

    @property
    def min_xy(self) -> NDArray | None:
        if self.min_x is None or self.min_y is None: return None
        return np.array([self.min_x, self.min_y, 0.0])

    @property
    def max_xy(self) -> NDArray | None:
        if self.max_x is None or self.max_y is None: return None
        return np.array([self.max_x, self.max_y, 0.0])

    @property
    def center_xy(self) -> NDArray | None:
        if self.min_xy is None or self.max_xy is None: return None
        return (self.min_xy + self.max_xy) / 2.0

class SlicingGroup():

    collections: dict[str, SlicingCollection]

    # Metadata
    wipe_tower_xy: NDArray = np.array([0, 0])
    wipe_tower_rotation_deg: float = 0

    def __init__(self, selected_objs: list[Object]) -> None:

        family = []
        for obj in selected_objs:
            family.append(obj)
            family.extend(get_all_children(obj))

        filtered_objs = [o for o in family] # if getattr(o, TYPES_NAME).object_type != 'Ignore'

        def find_top_parent(o):
            current = o
            while current.parent is not None:
                current = current.parent
            return current
        
        parents: dict[str, list[Object]] = {}
        [parents.setdefault(find_top_parent(o).name, []).append(o) for o in filtered_objs]

        self.collections = {}
        for k, objects in parents.items():
            self.collections[k] = SlicingCollection(objects, k)
        
        self._extract_metadata(family)

    def _extract_metadata(self, objs) -> None:
        import math
        depsgraph = bpy.context.evaluated_depsgraph_get()
        scale_tx = 1000. * bpy.context.scene.unit_settings.scale_length #type: ignore

        self.wipe_tower_xy = self.center_xy[0:2]
        for obj in objs:
            if obj.mode == 'EDIT':
                obj.update_from_editmode()
            eval_obj = obj.evaluated_get(depsgraph)
            if hasattr(eval_obj, 'nodeslicer'):
                if eval_obj.nodeslicer.object_type == 'WipeTower':
                    self.wipe_tower_xy = np.array(eval_obj.location[0:2]) * scale_tx
                    self.wipe_tower_rotation_deg = round(eval_obj.rotation_euler[2]*180/math.pi,5)
                    break

        return

    def offset(self, offset: NDArray):
        for k, so in self.collections.items(): so.offset(offset)
        self.wipe_tower_xy = self.wipe_tower_xy + offset[0:2]

    def scale(self, factor: float):
        for k, so in self.collections.items(): so.scale(factor)
        self.wipe_tower_xy *= factor

    @property
    def checksum(self):
        buf = bytearray()
        for key, col in sorted(self.collections.items(), key=lambda kv: kv[0]):
            key_bytes = key.encode('utf-8')
            buf.extend(struct.pack(">I", len(key_bytes)))
            buf.extend(key_bytes)
            buf.extend(struct.pack(">I", col.checksum))

        buf.extend(struct.pack(">I", crc32_array(self.wipe_tower_xy)))
        buf.extend(struct.pack(">I", crc32_array(np.array(self.wipe_tower_rotation_deg))))

        return zlib.crc32(buf) & 0xFFFFFFFF

    @property
    def height(self) -> float: return max([so.height for k, so in self.collections.items() if so.height], default=0)

    @property
    def min_x(self) -> float: return min([so.min_x for k, so in self.collections.items() if so.min_x], default=0)

    @property
    def max_x(self) -> float: return max([so.max_x for k, so in self.collections.items() if so.max_x], default=0)

    @property
    def min_y(self) -> float: return min([so.min_y for k, so in self.collections.items() if so.min_y], default=0)

    @property
    def max_y(self) -> float: return max([so.max_y for k, so in self.collections.items() if so.max_y is not None], default=0)

    @property
    def min_z(self) -> float: return min([so.min_z for k, so in self.collections.items() if so.min_z is not None], default=0)

    @property
    def max_z(self) -> float: return max([so.max_z for k, so in self.collections.items() if so.max_z is not None], default=0)

    @property
    def min_xy(self) -> NDArray:
        if self.min_x is None or self.min_y is None: return np.array([.0, .0, .0])
        return np.array([self.min_x, self.min_y, 0.0])

    @property
    def max_xy(self) -> NDArray:
        if self.max_x is None or self.max_y is None: return np.array([.0, .0, .0])
        return np.array([self.max_x, self.max_y, 0.0])

    @property
    def center_xy(self) -> NDArray:
        if self.min_xy is None or self.max_xy is None: return np.array([.0, .0, .0])
        return (self.min_xy + self.max_xy) / 2.0

class TriMesh():
    mesh: Mesh
    length_verts: int
    length_tris: int
    matrix_world: Any

    def __init__(self, mesh: Mesh, v: int, t: int, m) -> None:
        self.mesh = mesh
        self.length_verts = v
        self.length_tris = t
        self.matrix_world = m

def _cached_objects_to_tris(obj_names: tuple[str, ...], matrices: tuple[tuple[float, ...], ...], scale: float) -> NDArray[float64]:
    # Implementation remains largely the same but uses names to fetch objects
    tris_count: int = 0
    meshes: list[TriMesh] = []

    for name in obj_names:
        obj = bpy.data.objects.get(name)
        if not obj: continue
        
        try:
            depsgraph = bpy.context.evaluated_depsgraph_get()
            # We must use the evaluated object to get modifiers
            obj_eval = obj.evaluated_get(depsgraph)
            mesh: Mesh = obj_eval.to_mesh(preserve_all_data_layers=True, depsgraph=depsgraph)
        except: continue
        mesh.calc_loop_triangles()

        cur_length = len(mesh.loop_triangles)
        if cur_length == 0: continue

        tris_count += cur_length
        
        meshes += [TriMesh(
            mesh,
            len(mesh.vertices),
            len(mesh.loop_triangles),
            np.array(obj.matrix_world.transposed())
        )]

    tris_flat: np.ndarray = np.empty(tris_count * 4 * 3, dtype=np.float64)
    tris: np.ndarray = tris_flat.reshape(-1,  4,  3)

    col_idx = 0
    for trimesh in meshes:
        tris_v_i_flat: np.ndarray = np.empty(trimesh.length_tris * 3, dtype=np.int32)
        tris_v_n_flat: np.ndarray = np.empty(trimesh.length_tris * 3, dtype=np.float64)
        tris_verts_flat: np.ndarray = np.empty(trimesh.length_verts * 3, dtype=np.float64)

        trimesh.mesh.loop_triangles.foreach_get("vertices", tris_v_i_flat)
        trimesh.mesh.loop_triangles.foreach_get("normal", tris_v_n_flat)
        trimesh.mesh.vertices.foreach_get("co", tris_verts_flat)

        tris_v_i: np.ndarray = tris_v_i_flat.reshape((-1, 3))
        tris_v_n: np.ndarray = tris_v_n_flat.reshape((-1, 3))
        tris_verts: np.ndarray = tris_verts_flat.reshape((-1, 3))

        homogeneous_verts: NDArray[float64] = np.hstack((tris_verts, np.ones((tris_verts.shape[0], 1)))) # type: ignore
        tx_verts = homogeneous_verts @ trimesh.matrix_world
        tx_verts = (tx_verts[:, :3]) * scale

        homogeneous_norm: NDArray[float64] = np.hstack((tris_v_n, np.ones((tris_v_n.shape[0], 1)))) # type: ignore
        tx_norm = homogeneous_norm @ trimesh.matrix_world.T
        tx_norm = tx_norm[:, :3]
        tx_norm = tx_norm / np.linalg.norm(tx_norm, axis=1, keepdims=True)
        tx_norm = tx_norm[:, np.newaxis, :]

        tx_tris = tx_verts[tris_v_i]
        
        tris_coords_and_norm: NDArray[float64] = np.concatenate((tx_tris, tx_norm), axis=1)
        
        tris[col_idx:col_idx + trimesh.length_tris,:] = tris_coords_and_norm
        
        col_idx += trimesh.length_tris

    return tris

def objects_to_tris(objects: list[Object], scale) -> NDArray[float64]:
    obj_names = tuple(o.name for o in objects)
    matrices = tuple(tuple(val for row in o.matrix_world for val in row) for o in objects)
    return _cached_objects_to_tris(obj_names, matrices, scale)

def _cached_capture_mesh_edges(
    obj_name: str, 
    mesh_name: str, 
    matrix_tuple: tuple[float, ...], 
    scene_scale: float,
    attr_name: str = "e_mult",
    f_attr_name: str = "f_rate"
) -> tuple[list[NDArray[float64]], list[NDArray[float64]], list[NDArray[float64]]]:
    # This helper expects that the object and mesh exist and are accessible via bpy.data
    obj = bpy.data.objects.get(obj_name)
    if not obj: return [], [], []

    # We don't use obj_eval here directly as it's not hashable, 
    # but we assume this is called when we know we need to re-capture.
    # The actual capture logic stays the same but is now inside this cached function.
    try:
        depsgraph = bpy.context.evaluated_depsgraph_get()
        obj_eval = obj.evaluated_get(depsgraph)
        mesh = obj_eval.to_mesh(preserve_all_data_layers=True, depsgraph=depsgraph)
    except Exception:
        return [], [], []

    if not mesh or not mesh.vertices: 
        return [], [], []

    try:
        world_mat = np.array(matrix_tuple).reshape(4, 4)
        num_verts = len(mesh.vertices)
        verts = np.empty((num_verts, 3), dtype=np.float64)
        mesh.vertices.foreach_get("co", verts.ravel())

        e_mults = np.ones(num_verts, dtype=np.float64)
        if attr_name in mesh.attributes:
            mesh.attributes[attr_name].data.foreach_get("value", e_mults)

        f_rates = np.zeros(num_verts, dtype=np.float64)
        if f_attr_name in mesh.attributes:
            mesh.attributes[f_attr_name].data.foreach_get("value", f_rates)

        tx_verts = (verts @ world_mat[:3, :3].T + world_mat[:3, 3]) * scene_scale

        edges_set = {tuple(sorted(e.vertices)) for e in mesh.edges}
        connected = np.array([(i, i + 1) in edges_set for i in range(num_verts - 1)])

        split_at = np.where(~connected)[0] + 1
        raw_paths = np.split(tx_verts, split_at)
        raw_e_mults = np.split(e_mults, split_at)
        raw_f_rates = np.split(f_rates, split_at)
        indices_split = np.split(np.arange(num_verts), split_at)

        final_paths = []
        final_e_mults = []
        final_f_rates = []
        for i, path_coords in enumerate(raw_paths):
            if path_coords.shape[0] == 0:
                continue
            idx = indices_split[i]
            e_vals = raw_e_mults[i]
            f_vals = raw_f_rates[i]
            if len(idx) > 2 and tuple(sorted((idx[0], idx[-1]))) in edges_set:
                path_coords = np.vstack((path_coords, path_coords[0]))
                e_vals = np.concatenate((e_vals, [e_vals[0]]))
                f_vals = np.concatenate((f_vals, [f_vals[0]]))
            final_paths.append(path_coords)
            final_e_mults.append(e_vals)
            final_f_rates.append(f_vals)

        return final_paths, final_e_mults, final_f_rates

    finally:
        if hasattr(obj_eval, "to_mesh_clear"):
            obj_eval.to_mesh_clear()

def capture_mesh_edges(obj: Object, attr_name: str = "e_mult", f_attr_name: str = "f_rate") -> tuple[list[NDArray[float64]], list[NDArray[float64]], list[NDArray[float64]]]:
    if not obj or not obj.data: return [], [], []

    if obj.mode == 'EDIT':
        obj.update_from_editmode()

    assert bpy.context.scene and bpy.context.scene.unit_settings
    scene_scale = 1000.0 * bpy.context.scene.unit_settings.scale_length

    # Create a hashable key from object state
    matrix_tuple = tuple(val for row in obj.matrix_world for val in row) # Flattened tuple

    return _cached_capture_mesh_edges(
        obj.name, 
        obj.data.name, 
        matrix_tuple, 
        scene_scale,
        attr_name,
        f_attr_name
    )