"""
2D Vector Edge Extraction Module for PDF Exporter.
Extracts top-view wireframe, lineart, boundary, contour, and crease edges from 3D SlicingObjects.
"""
import math
from typing import Any
import numpy as np
from numpy.typing import NDArray


def deduplicate_2d_edges(edges_2d: NDArray[np.float64]) -> NDArray[np.float64]:
    """
    Deduplicates 2D edge segments regardless of direction (A->B or B->A).
    """
    if len(edges_2d) == 0:
        return edges_2d

    p0 = edges_2d[:, 0, :]
    p1 = edges_2d[:, 1, :]

    swap_mask = (p0[:, 0] > p1[:, 0]) | ((p0[:, 0] == p1[:, 0]) & (p0[:, 1] > p1[:, 1]))

    canonical_p0 = np.where(swap_mask[:, None], p1, p0)
    canonical_p1 = np.where(swap_mask[:, None], p0, p1)

    quantized_p0 = np.round(canonical_p0, decimals=4)
    quantized_p1 = np.round(canonical_p1, decimals=4)

    stacked = np.column_stack([quantized_p0, quantized_p1])

    _, unique_indices = np.unique(stacked, axis=0, return_index=True)
    return edges_2d[unique_indices]


def extract_unique_top_view_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts 2D top-view unique edges (X, Y projection) from mesh_data.
    """
    from ...infra.blender.blender_mesh_capture import SlicingGroup, SlicingCollection, SlicingObject

    all_edge_segments: list[NDArray[np.float64]] = []

    def process_item(item: Any):
        if item is None:
            return
        elif isinstance(item, (list, tuple)):
            for sub in item:
                process_item(sub)
        elif isinstance(item, SlicingGroup):
            for coll in item.collections.values():
                process_item(coll)
        elif isinstance(item, SlicingCollection):
            all_verts, all_idxs = item.unique_verts
            if len(all_idxs) > 0:
                edges = np.vstack([
                    all_idxs[:, [0, 1]],
                    all_idxs[:, [1, 2]],
                    all_idxs[:, [2, 0]]
                ])
                sorted_edges = np.sort(edges, axis=1)
                unique_edge_indices = np.unique(sorted_edges, axis=0)
                edge_3d = all_verts[unique_edge_indices]
                all_edge_segments.append(edge_3d[:, :, :2])
        elif isinstance(item, SlicingObject):
            if item.mesh.size > 0:
                verts_3d = item.mesh[:, :3, :3]
                verts_2d = verts_3d[:, :, :2]
                e0 = verts_2d[:, [0, 1], :]
                e1 = verts_2d[:, [1, 2], :]
                e2 = verts_2d[:, [2, 0], :]
                edges_2d = np.concatenate([e0, e1, e2], axis=0)
                all_edge_segments.append(edges_2d)

    process_item(mesh_data)

    if not all_edge_segments:
        return np.empty((0, 2, 2), dtype=np.float64)

    combined_edges = np.concatenate(all_edge_segments, axis=0)
    if combined_edges.shape[0] == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    dists_sq = np.sum((combined_edges[:, 0, :] - combined_edges[:, 1, :]) ** 2, axis=1)
    valid_mask = dists_sq > 1e-8
    valid_edges = combined_edges[valid_mask]
    if valid_edges.shape[0] == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    return deduplicate_2d_edges(valid_edges)


def extract_lineart_edges(mesh_data: Any, max_angle_deg: float = 30.0) -> NDArray[np.float64]:
    """
    Fully vectorized NumPy Lineart edge extraction:
    1. Silhouette edge: shared between at least one CCW face and at least one CW face in 2D top view.
    2. Boundary edge: owned by only 1 face.
    3. Crease edge: owned by 2 or more faces where 3D dihedral angle between normals > max_angle_deg.
    Returns 2D edge segments array of shape (N, 2, 2).
    """
    from ...infra.blender.blender_mesh_capture import SlicingGroup, SlicingCollection, SlicingObject

    all_tris_3d = []
    all_normals = []

    def process_item(item: Any):
        if item is None:
            return
        elif isinstance(item, (list, tuple)):
            for sub in item:
                process_item(sub)
        elif isinstance(item, SlicingGroup):
            for coll in item.collections.values():
                process_item(coll)
        elif isinstance(item, SlicingCollection):
            for so in item.objects:
                process_item(so)
        elif isinstance(item, SlicingObject):
            if item.mesh.size > 0:
                verts_3d = item.mesh[:, :3, :3]  # (N, 3, 3)
                if item.mesh.shape[1] >= 4:
                    normals = item.mesh[:, 3, :]  # (N, 3)
                else:
                    v0_3d = verts_3d[:, 0, :]
                    v1_3d = verts_3d[:, 1, :]
                    v2_3d = verts_3d[:, 2, :]
                    normals = np.cross(v1_3d - v0_3d, v2_3d - v0_3d)

                norms_len = np.linalg.norm(normals, axis=1, keepdims=True)
                norms_len = np.where(norms_len == 0, 1.0, norms_len)
                normals_unit = normals / norms_len

                zero_norm_mask = (normals[:, 0] == 0) & (normals[:, 1] == 0) & (normals[:, 2] == 0)
                if np.any(zero_norm_mask):
                    v0_3d = verts_3d[:, 0, :]
                    v1_3d = verts_3d[:, 1, :]
                    v2_3d = verts_3d[:, 2, :]
                    computed_normals = np.cross(v1_3d - v0_3d, v2_3d - v0_3d)
                    cn_len = np.linalg.norm(computed_normals, axis=1, keepdims=True)
                    cn_len = np.where(cn_len == 0, 1.0, cn_len)
                    normals_unit[zero_norm_mask] = (computed_normals / cn_len)[zero_norm_mask]

                all_tris_3d.append(verts_3d)
                all_normals.append(normals_unit)

    process_item(mesh_data)

    if not all_tris_3d:
        return np.empty((0, 2, 2), dtype=np.float64)

    tris_3d = np.concatenate(all_tris_3d, axis=0)  # (N, 3, 3)
    normals_unit = np.concatenate(all_normals, axis=0)  # (N, 3)
    num_tris = len(tris_3d)

    if num_tris == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    # Top-view 2D vertices
    verts_2d = tris_3d[:, :, :2]  # (N, 3, 2)

    # 2D orientation per triangle
    v0_2d = verts_2d[:, 0, :]
    v1_2d = verts_2d[:, 1, :]
    v2_2d = verts_2d[:, 2, :]
    cross_2d = (v1_2d[:, 0] - v0_2d[:, 0]) * (v2_2d[:, 1] - v0_2d[:, 1]) - (v1_2d[:, 1] - v0_2d[:, 1]) * (v2_2d[:, 0] - v0_2d[:, 0])
    is_ccw = cross_2d > 0
    is_cw = cross_2d < 0

    # Extract all 3 edges per triangle
    e0_3d = tris_3d[:, [0, 1], :]
    e1_3d = tris_3d[:, [1, 2], :]
    e2_3d = tris_3d[:, [2, 0], :]
    edges_3d = np.concatenate([e0_3d, e1_3d, e2_3d], axis=0)  # (3N, 2, 3)

    e0_2d = verts_2d[:, [0, 1], :]
    e1_2d = verts_2d[:, [1, 2], :]
    e2_2d = verts_2d[:, [2, 0], :]
    edges_2d = np.concatenate([e0_2d, e1_2d, e2_2d], axis=0)  # (3N, 2, 2)

    edge_ccw = np.tile(is_ccw, 3)  # (3N,)
    edge_cw = np.tile(is_cw, 3)    # (3N,)
    edge_normals = np.tile(normals_unit, (3, 1))  # (3N, 3)

    # Filter out near zero-length 2D edges
    dists_sq_2d = np.sum((edges_2d[:, 0, :] - edges_2d[:, 1, :]) ** 2, axis=1)
    valid_len_mask = dists_sq_2d >= 1e-8

    edges_3d = edges_3d[valid_len_mask]
    edges_2d = edges_2d[valid_len_mask]
    edge_ccw = edge_ccw[valid_len_mask]
    edge_cw = edge_cw[valid_len_mask]
    edge_normals = edge_normals[valid_len_mask]

    if len(edges_3d) == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    # Canonical undirected edge sorting in 3D (rounded to 4 decimal places)
    rounded_3d = np.round(edges_3d, decimals=4)
    pA = rounded_3d[:, 0, :]
    pB = rounded_3d[:, 1, :]
    swap_mask = (pA[:, 0] > pB[:, 0]) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] > pB[:, 1])
    ) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] == pB[:, 1]) & (pA[:, 2] > pB[:, 2])
    )

    canonical_3d = rounded_3d.copy()
    canonical_3d[swap_mask] = canonical_3d[swap_mask][:, ::-1, :]

    canonical_2d = edges_2d.copy()
    canonical_2d[swap_mask] = canonical_2d[swap_mask][:, ::-1, :]

    # Unique 3D edge identification using NumPy structured view
    flat_3d = canonical_3d.reshape(len(canonical_3d), 6)
    view_bytes = np.ascontiguousarray(flat_3d).view(np.dtype((np.void, 48))).ravel()
    unique_keys, first_indices, inverse_indices, counts = np.unique(
        view_bytes, return_index=True, return_inverse=True, return_counts=True
    )

    num_unique = len(unique_keys)

    # Rule 1: Silhouette edge (has CCW and CW face)
    has_ccw = np.bincount(inverse_indices, weights=edge_ccw.astype(int), minlength=num_unique) > 0
    has_cw = np.bincount(inverse_indices, weights=edge_cw.astype(int), minlength=num_unique) > 0
    rule1_unique = has_ccw & has_cw

    # Rule 2: Boundary edge (owned by 1 face)
    rule2_unique = counts == 1

    # Rule 3: Crease edge (owned by 2+ faces with normal angle > max_angle_deg)
    cos_threshold = math.cos(math.radians(max_angle_deg))
    rule3_unique = np.zeros(num_unique, dtype=bool)

    multi_edge_mask = counts >= 2
    if np.any(multi_edge_mask):
        sort_by_edge = np.argsort(inverse_indices)
        sorted_edge_ids = inverse_indices[sort_by_edge]
        sorted_normals = edge_normals[sort_by_edge]

        splits = np.where(sorted_edge_ids[:-1] != sorted_edge_ids[1:])[0] + 1
        edge_groups = np.split(sorted_normals, splits)
        unique_group_ids = np.split(sorted_edge_ids, splits)

        for grp_normals, grp_ids in zip(edge_groups, unique_group_ids):
            edge_id = grp_ids[0]
            if len(grp_normals) >= 2:
                dots = grp_normals @ grp_normals.T
                if np.any(dots < cos_threshold - 1e-6):
                    rule3_unique[edge_id] = True

    marked_unique_mask = rule1_unique | rule2_unique | rule3_unique
    marked_first_indices = first_indices[marked_unique_mask]

    return canonical_2d[marked_first_indices]


def extract_boundary_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts 2D edge segments owned by only 1 face (single-face boundary edges).
    """
    from ...infra.blender.blender_mesh_capture import SlicingGroup, SlicingCollection, SlicingObject

    all_tris_3d = []

    def process_item(item: Any):
        if item is None:
            return
        elif isinstance(item, (list, tuple)):
            for sub in item:
                process_item(sub)
        elif isinstance(item, SlicingGroup):
            for coll in item.collections.values():
                process_item(coll)
        elif isinstance(item, SlicingCollection):
            for so in item.objects:
                process_item(so)
        elif isinstance(item, SlicingObject):
            if item.mesh.size > 0:
                all_tris_3d.append(item.mesh[:, :3, :3])

    process_item(mesh_data)

    if not all_tris_3d:
        return np.empty((0, 2, 2), dtype=np.float64)

    tris_3d = np.concatenate(all_tris_3d, axis=0)
    if len(tris_3d) == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    verts_2d = tris_3d[:, :, :2]
    e0_3d = tris_3d[:, [0, 1], :]
    e1_3d = tris_3d[:, [1, 2], :]
    e2_3d = tris_3d[:, [2, 0], :]
    edges_3d = np.concatenate([e0_3d, e1_3d, e2_3d], axis=0)

    e0_2d = verts_2d[:, [0, 1], :]
    e1_2d = verts_2d[:, [1, 2], :]
    e2_2d = verts_2d[:, [2, 0], :]
    edges_2d = np.concatenate([e0_2d, e1_2d, e2_2d], axis=0)

    dists_sq_2d = np.sum((edges_2d[:, 0, :] - edges_2d[:, 1, :]) ** 2, axis=1)
    valid_mask = dists_sq_2d >= 1e-8

    edges_3d = edges_3d[valid_mask]
    edges_2d = edges_2d[valid_mask]

    if len(edges_3d) == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    rounded_3d = np.round(edges_3d, decimals=4)
    pA = rounded_3d[:, 0, :]
    pB = rounded_3d[:, 1, :]
    swap_mask = (pA[:, 0] > pB[:, 0]) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] > pB[:, 1])
    ) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] == pB[:, 1]) & (pA[:, 2] > pB[:, 2])
    )

    canonical_3d = rounded_3d.copy()
    canonical_3d[swap_mask] = canonical_3d[swap_mask][:, ::-1, :]

    canonical_2d = edges_2d.copy()
    canonical_2d[swap_mask] = canonical_2d[swap_mask][:, ::-1, :]

    flat_3d = canonical_3d.reshape(len(canonical_3d), 6)
    view_bytes = np.ascontiguousarray(flat_3d).view(np.dtype((np.void, 48))).ravel()
    unique_keys, first_indices, inverse_indices, counts = np.unique(
        view_bytes, return_index=True, return_inverse=True, return_counts=True
    )

    boundary_mask = counts == 1
    return canonical_2d[first_indices[boundary_mask]]


def extract_crease_edges(mesh_data: Any, max_angle_deg: float = 30.0) -> NDArray[np.float64]:
    """
    Extracts 2D edge segments owned by 2 or more faces where the 3D dihedral angle between face normals > max_angle_deg.
    """
    return extract_lineart_edges(mesh_data, max_angle_deg=max_angle_deg)


def extract_silhouette_contour_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts top-view silhouette contour edges using CW vs CCW triangle orientation in +Z top-view projection.
    """
    from ...infra.blender.blender_mesh_capture import SlicingGroup, SlicingCollection, SlicingObject

    all_tris_3d = []

    def process_item(item: Any):
        if item is None:
            return
        elif isinstance(item, (list, tuple)):
            for sub in item:
                process_item(sub)
        elif isinstance(item, SlicingGroup):
            for coll in item.collections.values():
                process_item(coll)
        elif isinstance(item, SlicingCollection):
            for so in item.objects:
                process_item(so)
        elif isinstance(item, SlicingObject):
            if item.mesh.size > 0:
                all_tris_3d.append(item.mesh[:, :3, :3])

    process_item(mesh_data)

    if not all_tris_3d:
        return np.empty((0, 2, 2), dtype=np.float64)

    tris_3d = np.concatenate(all_tris_3d, axis=0)  # (N, 3, 3)
    num_tris = len(tris_3d)

    if num_tris == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    verts_2d = tris_3d[:, :, :2]  # (N, 3, 2)

    v0_2d = verts_2d[:, 0, :]
    v1_2d = verts_2d[:, 1, :]
    v2_2d = verts_2d[:, 2, :]
    cross_2d = (v1_2d[:, 0] - v0_2d[:, 0]) * (v2_2d[:, 1] - v0_2d[:, 1]) - (v1_2d[:, 1] - v0_2d[:, 1]) * (v2_2d[:, 0] - v0_2d[:, 0])
    is_ccw = cross_2d > 0
    is_cw = cross_2d < 0

    e0_3d = tris_3d[:, [0, 1], :]
    e1_3d = tris_3d[:, [1, 2], :]
    e2_3d = tris_3d[:, [2, 0], :]
    edges_3d = np.concatenate([e0_3d, e1_3d, e2_3d], axis=0)  # (3N, 2, 3)

    e0_2d = verts_2d[:, [0, 1], :]
    e1_2d = verts_2d[:, [1, 2], :]
    e2_2d = verts_2d[:, [2, 0], :]
    edges_2d = np.concatenate([e0_2d, e1_2d, e2_2d], axis=0)  # (3N, 2, 2)

    edge_ccw = np.tile(is_ccw, 3)
    edge_cw = np.tile(is_cw, 3)

    dists_sq_2d = np.sum((edges_2d[:, 0, :] - edges_2d[:, 1, :]) ** 2, axis=1)
    valid_len_mask = dists_sq_2d >= 1e-8

    edges_3d = edges_3d[valid_len_mask]
    edges_2d = edges_2d[valid_len_mask]
    edge_ccw = edge_ccw[valid_len_mask]
    edge_cw = edge_cw[valid_len_mask]

    if len(edges_3d) == 0:
        return np.empty((0, 2, 2), dtype=np.float64)

    rounded_3d = np.round(edges_3d, decimals=4)
    pA = rounded_3d[:, 0, :]
    pB = rounded_3d[:, 1, :]
    swap_mask = (pA[:, 0] > pB[:, 0]) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] > pB[:, 1])
    ) | (
        (pA[:, 0] == pB[:, 0]) & (pA[:, 1] == pB[:, 1]) & (pA[:, 2] > pB[:, 2])
    )

    canonical_3d = rounded_3d.copy()
    canonical_3d[swap_mask] = canonical_3d[swap_mask][:, ::-1, :]

    canonical_2d = edges_2d.copy()
    canonical_2d[swap_mask] = canonical_2d[swap_mask][:, ::-1, :]

    flat_3d = canonical_3d.reshape(len(canonical_3d), 6)
    view_bytes = np.ascontiguousarray(flat_3d).view(np.dtype((np.void, 48))).ravel()
    unique_keys, first_indices, inverse_indices, counts = np.unique(
        view_bytes, return_index=True, return_inverse=True, return_counts=True
    )

    num_unique = len(unique_keys)

    has_ccw = np.bincount(inverse_indices, weights=edge_ccw.astype(int), minlength=num_unique) > 0
    has_cw = np.bincount(inverse_indices, weights=edge_cw.astype(int), minlength=num_unique) > 0
    silhouette_mask = has_ccw & has_cw

    return canonical_2d[first_indices[silhouette_mask]]


def extract_outer_contour_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts top-view contour edges using backfacing CW vs CCW triangle orientation.
    """
    return extract_silhouette_contour_edges(mesh_data)