"""2.5D Delaunay triangulation — generate meshes from point clouds."""

from __future__ import annotations

from pathlib import Path
from typing import Any

import numpy as np
from scipy.spatial import Delaunay

from pywolken.core.pointcloud import PointCloud


class Mesh:
    """Triangle mesh with vertices and face indices.

    Attributes:
        vertices: (N, 3) float64 array of vertex positions.
        faces: (M, 3) int64 array of triangle vertex indices.
        vertex_colors: Optional (N, 3) uint8 array of RGB colors.
    """

    def __init__(
        self,
        vertices: np.ndarray,
        faces: np.ndarray,
        vertex_colors: np.ndarray | None = None,
    ) -> None:
        self.vertices = np.asarray(vertices, dtype=np.float64)
        self.faces = np.asarray(faces, dtype=np.int64)
        self.vertex_colors = vertex_colors

    @property
    def num_vertices(self) -> int:
        return len(self.vertices)

    @property
    def num_faces(self) -> int:
        return len(self.faces)

    def decimate(self, target_faces: int) -> Mesh:
        """Simple mesh decimation by random face removal.

        For production use, a proper edge-collapse decimation would be
        better. This is a basic approximation.

        Args:
            target_faces: Target number of faces to keep.

        Returns:
            New Mesh with reduced face count.
        """
        if target_faces >= self.num_faces:
            return self

        rng = np.random.default_rng(42)
        indices = rng.choice(self.num_faces, size=target_faces, replace=False)
        indices.sort()
        new_faces = self.faces[indices]

        # Remap vertices (keep only those referenced by remaining faces)
        used = np.unique(new_faces)
        remap = np.full(self.num_vertices, -1, dtype=np.int64)
        remap[used] = np.arange(len(used))

        new_vertices = self.vertices[used]
        new_faces = remap[new_faces]

        new_colors = None
        if self.vertex_colors is not None:
            new_colors = self.vertex_colors[used]

        return Mesh(new_vertices, new_faces, new_colors)

    def write_obj(self, path: str) -> None:
        """Write mesh to Wavefront OBJ format."""
        import io as _io
        with open(path, "wb") as f:
            f.write(b"# Generated by pywolken\n")
            buf = _io.BytesIO()
            np.savetxt(buf, self.vertices, fmt="v %.6f %.6f %.6f")
            f.write(buf.getvalue())
            buf = _io.BytesIO()
            np.savetxt(buf, self.faces + 1, fmt="f %d %d %d")
            f.write(buf.getvalue())

    def write_stl(self, path: str) -> None:
        """Write mesh to binary STL format."""
        with open(path, "wb") as f:
            # 80-byte header
            f.write(b"pywolken STL" + b"\0" * 68)
            # Number of triangles
            f.write(np.array([self.num_faces], dtype="<u4").tobytes())

            # Vectorized: gather all face vertices
            v0 = self.vertices[self.faces[:, 0]]  # (M, 3)
            v1 = self.vertices[self.faces[:, 1]]
            v2 = self.vertices[self.faces[:, 2]]

            # Batch normal computation
            normals = np.cross(v1 - v0, v2 - v0)
            norms = np.linalg.norm(normals, axis=1, keepdims=True)
            norms[norms == 0] = 1.0
            normals /= norms

            # Build structured array for bulk binary write
            stl_dtype = np.dtype([
                ("normal", "<f4", (3,)),
                ("v0", "<f4", (3,)),
                ("v1", "<f4", (3,)),
                ("v2", "<f4", (3,)),
                ("attr", "<u2"),
            ])
            stl_data = np.empty(self.num_faces, dtype=stl_dtype)
            stl_data["normal"] = normals.astype(np.float32)
            stl_data["v0"] = v0.astype(np.float32)
            stl_data["v1"] = v1.astype(np.float32)
            stl_data["v2"] = v2.astype(np.float32)
            stl_data["attr"] = 0
            f.write(stl_data.tobytes())

    def write_ply(self, path: str) -> None:
        """Write mesh to PLY format (ASCII)."""
        import io as _io
        with open(path, "wb") as f:
            header_lines = [
                "ply",
                "format ascii 1.0",
                f"element vertex {self.num_vertices}",
                "property float x",
                "property float y",
                "property float z",
            ]
            if self.vertex_colors is not None:
                header_lines += [
                    "property uchar red",
                    "property uchar green",
                    "property uchar blue",
                ]
            header_lines += [
                f"element face {self.num_faces}",
                "property list uchar int vertex_indices",
                "end_header",
            ]
            f.write(("\n".join(header_lines) + "\n").encode("ascii"))

            # Bulk write vertices
            buf = _io.BytesIO()
            if self.vertex_colors is not None:
                vdata = np.column_stack([
                    self.vertices,
                    self.vertex_colors.astype(np.float64),
                ])
                np.savetxt(buf, vdata, fmt="%.6f %.6f %.6f %d %d %d")
            else:
                np.savetxt(buf, self.vertices, fmt="%.6f %.6f %.6f")
            f.write(buf.getvalue())

            # Bulk write faces
            buf = _io.BytesIO()
            fdata = np.column_stack([
                np.full(self.num_faces, 3, dtype=np.int64),
                self.faces,
            ])
            np.savetxt(buf, fdata, fmt="%d %d %d %d")
            f.write(buf.getvalue())

    def write(self, path: str, **options: Any) -> None:
        """Write mesh to file (auto-detect format from extension)."""
        ext = Path(path).suffix.lower()
        if ext == ".obj":
            self.write_obj(path)
        elif ext == ".stl":
            self.write_stl(path)
        elif ext == ".ply":
            self.write_ply(path)
        else:
            raise ValueError(f"Unsupported mesh format: {ext}. Use .obj, .stl, or .ply")

    def __repr__(self) -> str:
        return f"Mesh({self.num_vertices:,} vertices, {self.num_faces:,} faces)"


def triangulate_2d(
    pc: PointCloud,
    max_edge_length: float | None = None,
) -> Mesh:
    """Create a 2.5D triangle mesh from a point cloud.

    Projects points onto XY plane for Delaunay triangulation,
    then uses the Z values for the 3D mesh vertices.

    Args:
        pc: Point cloud with X, Y, Z dimensions.
        max_edge_length: Optional maximum triangle edge length.
            Triangles with any edge longer than this are removed
            (useful for removing large triangles at convex hull edges).

    Returns:
        Mesh with Delaunay triangulation.
    """
    x, y, z = pc["X"], pc["Y"], pc["Z"]
    points_2d = np.column_stack([x, y])

    tri = Delaunay(points_2d)
    faces = tri.simplices

    # Remove long-edge triangles if requested
    if max_edge_length is not None:
        vertices_3d = np.column_stack([x, y, z])
        # Vectorized edge length computation
        v0 = vertices_3d[faces[:, 0]]
        v1 = vertices_3d[faces[:, 1]]
        v2 = vertices_3d[faces[:, 2]]
        e1 = np.linalg.norm(v1 - v0, axis=1)
        e2 = np.linalg.norm(v2 - v0, axis=1)
        e3 = np.linalg.norm(v2 - v1, axis=1)
        max_edge = np.maximum(np.maximum(e1, e2), e3)
        faces = faces[max_edge <= max_edge_length]

    vertices = np.column_stack([x, y, z])

    # Add vertex colors if available
    colors = None
    if "Red" in pc and "Green" in pc and "Blue" in pc:
        r = pc["Red"].astype(np.float64)
        g = pc["Green"].astype(np.float64)
        b = pc["Blue"].astype(np.float64)
        # Normalize to 0-255 if values are in uint16 range
        if r.max() > 255:
            r = (r / 257).astype(np.uint8)
            g = (g / 257).astype(np.uint8)
            b = (b / 257).astype(np.uint8)
        else:
            r = r.astype(np.uint8)
            g = g.astype(np.uint8)
            b = b.astype(np.uint8)
        colors = np.column_stack([r, g, b])

    return Mesh(vertices, faces, colors)
