from __future__ import annotations
from typing import Any
import os
import math
import numpy as np
from numpy.typing import NDArray
import shapely
from shapely.geometry import Polygon, MultiPolygon

MM_TO_PT = 72.0 / 25.4  # 1 mm = 2.8346456692913386 PostScript points


def extract_unique_top_view_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts unique 2D line segments projected orthographically from top (+Z).
    Handles SlicingGroup, SlicingCollection, SlicingObject, or lists/dicts of them.
    Returns array of shape (N, 2, 2) where each element is [(x0, y0), (x1, y1)] in mm.
    """
    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]  # shape (E, 2, 3)
                all_edge_segments.append(edge_3d[:, :, :2])  # top view (X, Y)
        elif isinstance(item, SlicingObject):
            if item.mesh.size > 0:
                verts_3d = item.mesh[:, :3, :3]  # (N, 3, 3)
                verts_2d = verts_3d[:, :, :2]  # (N, 3, 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)  # (3N, 2, 2)
                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)

    # Filter out near-zero-length 2D edges (e.g. vertical edges seen from top)
    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)

    # Deduplicate 2D edges regardless of endpoint orientation
    rounded = np.round(valid_edges, decimals=5)
    swap_mask = (rounded[:, 0, 0] > rounded[:, 1, 0]) | (
        (rounded[:, 0, 0] == rounded[:, 1, 0]) & (rounded[:, 0, 1] > rounded[:, 1, 1])
    )
    ordered = valid_edges.copy()
    ordered[swap_mask] = ordered[swap_mask][:, ::-1, :]
    ordered_rounded = rounded.copy()
    ordered_rounded[swap_mask] = ordered_rounded[swap_mask][:, ::-1, :]

    view = ordered_rounded.reshape(len(ordered_rounded), -1)
    _, unique_indices = np.unique(view, axis=0, return_index=True)
    return ordered[unique_indices]


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)
                normals = item.mesh[:, 3, :]  # (N, 3)

                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_outer_contour_edges(mesh_data: Any) -> NDArray[np.float64]:
    """
    Extracts top-view outer boundary / silhouette contour line segments using Shapely 2D polygon union.
    """
    from ..infra.blender.blender_mesh_capture import SlicingGroup, SlicingCollection, SlicingObject

    all_tris_2d = []

    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)
                verts_2d = verts_3d[:, :, :2]  # (N, 3, 2)
                all_tris_2d.append(verts_2d)

    process_item(mesh_data)

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

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

    polys = []
    for tri in tris_2d:
        p = Polygon(tri)
        if p.is_valid and p.area > 1e-6:
            polys.append(p)

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

    try:
        union_poly = shapely.union_all(polys)
    except Exception:
        return np.empty((0, 2, 2), dtype=np.float64)

    boundary_edges = []

    def process_geom(geom):
        if geom.is_empty:
            return
        if isinstance(geom, Polygon):
            coords = np.array(geom.exterior.coords)
            for i in range(len(coords) - 1):
                boundary_edges.append([coords[i], coords[i + 1]])
            for interior in geom.interiors:
                icoords = np.array(interior.coords)
                for i in range(len(icoords) - 1):
                    boundary_edges.append([icoords[i], icoords[i + 1]])
        elif isinstance(geom, MultiPolygon):
            for sub in geom.geoms:
                process_geom(sub)
        elif hasattr(geom, 'geoms'):
            for sub in geom.geoms:
                process_geom(sub)

    process_geom(union_poly)
    if not boundary_edges:
        return np.empty((0, 2, 2), dtype=np.float64)

    return np.array(boundary_edges, dtype=np.float64)


def _format_pdf_lines(lines_arr: NDArray[np.float64]) -> list[str]:
    """
    Vectorized formatting of 2D line segments to PDF operator strings.
    lines_arr shape: (N, 2, 2)
    """
    if len(lines_arr) == 0:
        return []
    flat = lines_arr.reshape(-1, 4)
    return [f"{row[0]:.2f} {row[1]:.2f} m {row[2]:.2f} {row[3]:.2f} l" for row in flat]


def export_wireframe_to_pdf(
    edges_2d: NDArray[np.float64],
    filepath: str,
    style: str = 'WIREFRAME',
    crease_angle: float = 30.0,
    margin_mm: float = 10.0,
    line_width_mm: float = 0.5,
    stroke_color: tuple[float, float, float] = (0.0, 0.0, 0.0),
    bg_color: tuple[float, float, float] = (1.0, 1.0, 1.0),
    raw_geometry: Any = None
) -> None:
    """
    Exports 2D vector geometry to a 1:1 true physical scale vector PDF file.
    Styles:
    - WIREFRAME: All projected 3D wireframe edges
    - LINEART: Silhouette, boundary, and crease edges above threshold angle
    - OUTER_CONTOUR: Silhouette boundary contour outline
    """
    if style == 'LINEART' and raw_geometry is not None:
        lineart_edges = extract_lineart_edges(raw_geometry, max_angle_deg=crease_angle)
        if len(lineart_edges) > 0:
            edges_to_draw = lineart_edges
        else:
            edges_to_draw = edges_2d
    elif style == 'OUTER_CONTOUR' and raw_geometry is not None:
        contour_edges = extract_outer_contour_edges(raw_geometry)
        if len(contour_edges) > 0:
            edges_to_draw = contour_edges
        else:
            edges_to_draw = edges_2d
    else:
        edges_to_draw = edges_2d

    if len(edges_to_draw) > 0:
        min_x = float(edges_to_draw[:, :, 0].min())
        max_x = float(edges_to_draw[:, :, 0].max())
        min_y = float(edges_to_draw[:, :, 1].min())
        max_y = float(edges_to_draw[:, :, 1].max())
    elif len(edges_2d) > 0:
        min_x = float(edges_2d[:, :, 0].min())
        max_x = float(edges_2d[:, :, 0].max())
        min_y = float(edges_2d[:, :, 1].min())
        max_y = float(edges_2d[:, :, 1].max())
    else:
        min_x = max_x = min_y = max_y = 0.0

    w_geom = max_x - min_x
    h_geom = max_y - min_y

    L = max(w_geom, h_geom)
    if L > 1e-6:
        step_raw = L / 10.0
        exponent = int(np.round(np.log10(step_raw)))
        main_step = 10.0 ** exponent
    else:
        main_step = 10.0

    minor_step = main_step / 10.0

    if len(edges_to_draw) > 0 or len(edges_2d) > 0:
        grid_min_x = math.floor(min_x / main_step) * main_step
        grid_max_x = math.ceil(max_x / main_step) * main_step
        grid_min_y = math.floor(min_y / main_step) * main_step
        grid_max_y = math.ceil(max_y / main_step) * main_step

        if grid_min_x == grid_max_x:
            grid_max_x += main_step
        if grid_min_y == grid_max_y:
            grid_max_y += main_step
    else:
        grid_min_x, grid_max_x = -50.0, 50.0
        grid_min_y, grid_max_y = -50.0, 50.0

    w_grid_mm = grid_max_x - grid_min_x
    h_grid_mm = grid_max_y - grid_min_y

    w_total_mm = w_grid_mm + 2.0 * margin_mm
    h_total_mm = h_grid_mm + 2.0 * margin_mm

    page_w_pt = w_total_mm * MM_TO_PT
    page_h_pt = h_total_mm * MM_TO_PT

    line_width_pt = line_width_mm * MM_TO_PT

    def to_pdf_pt(pt_mm):
        x_pt = (pt_mm[:, 0] - grid_min_x + margin_mm) * MM_TO_PT
        y_pt = (pt_mm[:, 1] - grid_min_y + margin_mm) * MM_TO_PT
        return np.column_stack([x_pt, y_pt])

    max_grid_steps = 2000

    # Vectorized vertical grid lines
    k_x_start = int(np.round(grid_min_x / minor_step))
    k_x_end = int(np.round(grid_max_x / minor_step))
    k_x = np.arange(k_x_start, k_x_end + 1)

    if len(k_x) <= max_grid_steps:
        gx = k_x * minor_step
        pts0_v = np.column_stack([gx, np.full_like(gx, grid_min_y)])
        pts1_v = np.column_stack([gx, np.full_like(gx, grid_max_y)])

        p0_v = to_pdf_pt(pts0_v)
        p1_v = to_pdf_pt(pts1_v)
        is_main_x = (k_x % 10 == 0)

        main_v_lines = np.column_stack([p0_v[is_main_x], p1_v[is_main_x]]).reshape(-1, 2, 2)
        minor_v_lines = np.column_stack([p0_v[~is_main_x], p1_v[~is_main_x]]).reshape(-1, 2, 2)
    else:
        main_v_lines = np.empty((0, 2, 2))
        minor_v_lines = np.empty((0, 2, 2))

    # Vectorized horizontal grid lines
    k_y_start = int(np.round(grid_min_y / minor_step))
    k_y_end = int(np.round(grid_max_y / minor_step))
    k_y = np.arange(k_y_start, k_y_end + 1)

    if len(k_y) <= max_grid_steps:
        gy = k_y * minor_step
        pts0_h = np.column_stack([np.full_like(gy, grid_min_x), gy])
        pts1_h = np.column_stack([np.full_like(gy, grid_max_x), gy])

        p0_h = to_pdf_pt(pts0_h)
        p1_h = to_pdf_pt(pts1_h)
        is_main_y = (k_y % 10 == 0)

        main_h_lines = np.column_stack([p0_h[is_main_y], p1_h[is_main_y]]).reshape(-1, 2, 2)
        minor_h_lines = np.column_stack([p0_h[~is_main_y], p1_h[~is_main_y]]).reshape(-1, 2, 2)
    else:
        main_h_lines = np.empty((0, 2, 2))
        minor_h_lines = np.empty((0, 2, 2))

    main_lines_arr = np.concatenate([main_v_lines, main_h_lines], axis=0) if len(main_v_lines) or len(main_h_lines) else np.empty((0, 2, 2))
    minor_lines_arr = np.concatenate([minor_v_lines, minor_h_lines], axis=0) if len(minor_v_lines) or len(minor_h_lines) else np.empty((0, 2, 2))

    stream_cmds: list[str] = []

    # 1. Solid background fill
    br, bg, bb = bg_color
    sr, sg, sb = stroke_color
    stream_cmds.append(f"{br:.3f} {bg:.3f} {bb:.3f} rg")
    stream_cmds.append(f"0 0 {page_w_pt:.2f} {page_h_pt:.2f} re f")

    # 2. Wireframe stroke color
    stream_cmds.append(f"{sr:.3f} {sg:.3f} {sb:.3f} RG")

    # 3. Minor Grid Lines (25% Opacity)
    if len(minor_lines_arr) > 0:
        stream_cmds.append("/GS25 gs")
        minor_grid_pt = min(0.3 * MM_TO_PT, line_width_pt)
        stream_cmds.append(f"{minor_grid_pt:.2f} w")
        lines_str = _format_pdf_lines(minor_lines_arr)
        for start_idx in range(0, len(lines_str), 5000):
            stream_cmds.extend(lines_str[start_idx:start_idx + 5000])
            stream_cmds.append("S")

    # 4. Main Grid Lines (50% Opacity)
    if len(main_lines_arr) > 0:
        stream_cmds.append("/GS50 gs")
        main_grid_pt = min(0.4 * MM_TO_PT, line_width_pt)
        stream_cmds.append(f"{main_grid_pt:.2f} w")
        lines_str = _format_pdf_lines(main_lines_arr)
        for start_idx in range(0, len(lines_str), 5000):
            stream_cmds.extend(lines_str[start_idx:start_idx + 5000])
            stream_cmds.append("S")

    # 5. Object Lines (100% Opacity)
    if len(edges_to_draw) > 0:
        stream_cmds.append("/GS100 gs")
        stream_cmds.append(f"{line_width_pt:.2f} w")
        stream_cmds.append("1 j 1 J")  # round join and cap

        p0_pt = to_pdf_pt(edges_to_draw[:, 0, :])
        p1_pt = to_pdf_pt(edges_to_draw[:, 1, :])
        obj_lines_arr = np.column_stack([p0_pt, p1_pt]).reshape(-1, 2, 2)
        lines_str = _format_pdf_lines(obj_lines_arr)

        for start_idx in range(0, len(lines_str), 5000):
            stream_cmds.extend(lines_str[start_idx:start_idx + 5000])
            stream_cmds.append("S")

    stream_content = "\n".join(stream_cmds) + "\n"
    stream_bytes = stream_content.encode("utf-8")

    objects: list[bytes] = []
    objects.append(b"1 0 obj\n<< /Type /Catalog /Pages 2 0 R >>\nendobj\n")
    objects.append(b"2 0 obj\n<< /Type /Pages /Kids [3 0 R] /Count 1 >>\nendobj\n")
    page_header = (
        f"3 0 obj\n"
        f"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 {page_w_pt:.2f} {page_h_pt:.2f}] "
        f"/Resources << /ExtGState << /GS25 5 0 R /GS50 6 0 R /GS100 7 0 R >> >> "
        f"/Contents 4 0 R >>\n"
        f"endobj\n"
    )
    objects.append(page_header.encode("utf-8"))
    content_header = f"4 0 obj\n<< /Length {len(stream_bytes)} >>\nstream\n".encode("utf-8")
    content_footer = b"endstream\nendobj\n"
    objects.append(content_header + stream_bytes + content_footer)
    objects.append(b"5 0 obj\n<< /Type /ExtGState /CA 0.25 /ca 0.25 >>\nendobj\n")
    objects.append(b"6 0 obj\n<< /Type /ExtGState /CA 0.5 /ca 0.5 >>\nendobj\n")
    objects.append(b"7 0 obj\n<< /Type /ExtGState /CA 1.0 /ca 1.0 >>\nendobj\n")

    header = b"%PDF-1.4\n%\xe2\xe3\xcf\xd3\n"

    parent_dir = os.path.dirname(filepath)
    if parent_dir:
        os.makedirs(parent_dir, exist_ok=True)

    with open(filepath, "wb") as f:
        f.write(header)
        offsets = [0]
        current_offset = len(header)

        for obj_bytes in objects:
            offsets.append(current_offset)
            f.write(obj_bytes)
            current_offset += len(obj_bytes)

        xref_offset = current_offset
        f.write(b"xref\n")
        f.write(f"0 {len(offsets)}\n".encode("utf-8"))
        f.write(b"0000000000 65535 f \n")
        for off in offsets[1:]:
            f.write(f"{off:010d} 00000 n \n".encode("utf-8"))

        f.write(b"trailer\n")
        f.write(f"<< /Size {len(offsets)} /Root 1 0 R >>\n".encode("utf-8"))
        f.write(b"startxref\n")
        f.write(f"{xref_offset}\n".encode("utf-8"))
        f.write(b"%%EOF\n")