import concurrent
import copyreg
import hashlib
import io
import re
from concurrent.futures import ProcessPoolExecutor, Executor
from dataclasses import dataclass
from typing import Tuple, Callable, Generator

import OCP
import numpy as np
from OCP.BRep import BRep_Tool
from OCP.BRepAdaptor import BRepAdaptor_Curve
from OCP.GCPnts import GCPnts_TangentialDeflection
from OCP.TopExp import TopExp
from OCP.TopLoc import TopLoc_Location
from OCP.TopTools import TopTools_IndexedMapOfShape
from OCP.TopoDS import TopoDS_Face, TopoDS_Edge, TopoDS_Shape, TopoDS_Vertex
from build123d import Face, Vector, Shape, Vertex
from partcad.wrappers import cq_serialize
from pygltflib import LINE_STRIP, GLTF2, Material, PbrMetallicRoughness, TRIANGLES, POINTS, TextureInfo

from gltf import create_gltf, _checkerboard_image


@dataclass
class TessellationUpdate:
    """Tessellation update"""
    progress: float
    """Progress in percent"""

    # Current shape
    shape: TopoDS_Shape
    """(Sub)shape that was tessellated"""
    gltf: GLTF2
    """The valid GLTF containing only the current shape"""

    @property
    def kind(self) -> str:
        """The kind of the shape"""
        if isinstance(self.shape, TopoDS_Face):
            return "face"
        elif isinstance(self.shape, TopoDS_Edge):
            return "edge"
        elif isinstance(self.shape, TopoDS_Vertex):
            return "vertex"
        else:
            raise ValueError(f"Unknown shape type: {self.shape}")


progress_callback_t = Callable[[TessellationUpdate], None]


def _inflate_vec(*values: float):
    pnt = OCP.gp.gp_Vec(values[0], values[1], values[2])
    return pnt


def _reduce_vec(pnt: OCP.gp.gp_Vec):
    return _inflate_vec, (pnt.X(), pnt.Y(), pnt.Z())


def tessellate_count(ocp_shape: TopoDS_Shape) -> int:
    """Count the number of elements that will be tessellated"""
    shape = Shape(ocp_shape)
    return len(shape.faces()) + len(shape.edges()) + len(shape.vertices())


def tessellate(
        ocp_shape: TopoDS_Shape,
        tolerance: float = 0.1,
        angular_tolerance: float = 0.1,
        executor: Executor = ProcessPoolExecutor(),  # Set to ThreadPoolExecutor if pickling fails...
) -> Generator[TessellationUpdate, None, None]:
    """Tessellate a whole shape into a list of triangle vertices and a list of triangle indices.

    It uses multiprocessing to speed up the process, and publishes progress updates to the callback.
    """
    shape = Shape(ocp_shape)
    _register_pickle_if_needed()
    with executor:
        futures = []

        # Submit tessellation tasks
        for face in shape.faces():
            futures.append(executor.submit(_tessellate_element, face.wrapped, tolerance, angular_tolerance))
        for edge in shape.edges():
            futures.append(executor.submit(_tessellate_element, edge.wrapped, tolerance, angular_tolerance))
        for vertex in shape.vertices():
            futures.append(executor.submit(_tessellate_element, vertex.wrapped, tolerance, angular_tolerance))

        # Collect results as they come in
        for i, future in enumerate(concurrent.futures.as_completed(futures)):
            sub_shape, gltf = future.result()
            yield TessellationUpdate(
                progress=(i + 1) / len(futures),
                shape=sub_shape,
                gltf=gltf,
            )


_pickle_registered = False


def _register_pickle_if_needed():
    global _pickle_registered
    if _pickle_registered:
        return
    cq_serialize.register()
    copyreg.pickle(OCP.gp.gp_Vec, _reduce_vec)


# Define the function that will tessellate each element in parallel
def _tessellate_element(
        element: TopoDS_Shape, tolerance: float, angular_tolerance: float) -> Tuple[TopoDS_Shape, GLTF2]:
    if isinstance(element, TopoDS_Face):
        return element, _tessellate_face(element, tolerance, angular_tolerance)
    elif isinstance(element, TopoDS_Edge):
        return element, _tessellate_edge(element, angular_tolerance, angular_tolerance)
    elif isinstance(element, TopoDS_Vertex):
        return element, _tessellate_vertex(element)
    else:
        raise ValueError(f"Unknown element type: {element}")


TriMesh = Tuple[list[Vector], list[Tuple[int, int, int]]]


def _tessellate_face(
        ocp_face: TopoDS_Face,
        tolerance: float = 0.1,
        angular_tolerance: float = 0.1
) -> GLTF2:
    """Tessellate a face into a list of triangle vertices and a list of triangle indices"""
    face = Face(ocp_face)
    tri_mesh = face.tessellate(tolerance, angular_tolerance)

    # Get UV of each face from the parameters
    loc = TopLoc_Location()
    poly = BRep_Tool.Triangulation_s(face.wrapped, loc)
    uv = [
        [v.X(), v.Y()]
        for v in (poly.UVNode(i) for i in range(1, poly.NbNodes() + 1))
    ]

    vertices = np.array(list(map(lambda v: [v.X, v.Y, v.Z], tri_mesh[0])))
    indices = np.array(tri_mesh[1])
    tex_coord = np.array(uv)
    mode = TRIANGLES
    material = Material(pbrMetallicRoughness=PbrMetallicRoughness(
        baseColorFactor=[0.3, 1.0, 0.2, 1.0], roughnessFactor=0.1, baseColorTexture=TextureInfo(index=0)),
        alphaCutoff=None)
    return create_gltf(vertices, indices, tex_coord, mode, material, images=[_checkerboard_image])


def _tessellate_edge(
        ocp_edge: TopoDS_Edge,
        angular_deflection: float = 0.1,
        curvature_deflection: float = 0.1,
) -> GLTF2:
    """Tessellate a wire or edge into a list of ordered vertices"""
    curve = BRepAdaptor_Curve(ocp_edge)
    discretizer = GCPnts_TangentialDeflection(curve, angular_deflection, curvature_deflection)
    assert discretizer.NbPoints() > 1, "Edge is too small??"

    # TODO: get and apply transformation??

    # add vertices
    vertices: list[list[float]] = [
        [v.X(), v.Y(), v.Z()]
        for v in (
            discretizer.Value(i)  # .Transformed(transformation)
            for i in range(1, discretizer.NbPoints() + 1)
        )
    ]
    indices = np.array(list(map(lambda i: [i, i + 1], range(len(vertices) - 1))), dtype=np.uint8)
    tex_coord = np.array([], dtype=np.float32)
    mode = LINE_STRIP
    material = Material(
        pbrMetallicRoughness=PbrMetallicRoughness(baseColorFactor=[1.0, 1.0, 0.5, 1.0]),
        alphaCutoff=None)
    return create_gltf(np.array(vertices), indices, tex_coord, mode, material)


def _tessellate_vertex(ocp_vertex: TopoDS_Vertex) -> GLTF2:
    """Tessellate a vertex into a list of triangle vertices and a list of triangle indices"""
    c = Vertex(ocp_vertex).center()
    vertices = np.array([[c.X, c.Y, c.Z]])
    indices = np.array([0])
    tex_coord = np.array([], dtype=np.float32)
    mode = POINTS
    material = Material(
        pbrMetallicRoughness=PbrMetallicRoughness(baseColorFactor=[1.0, 0.5, 0.5, 1.0]),
        alphaCutoff=None)
    return create_gltf(vertices, indices, tex_coord, mode, material)


def _hashcode(obj: TopoDS_Shape) -> str:
    """Utility to compute the hash code of a shape recursively without the need to tessellate it"""
    # NOTE: obj.HashCode(MAX_HASH_CODE) is not stable across different runs of the same program
    # This is best-effort and not guaranteed to be unique
    data = io.BytesIO()
    map_of_shapes = TopTools_IndexedMapOfShape()
    TopExp.MapShapes_s(obj, map_of_shapes)
    for i in range(1, map_of_shapes.Extent() + 1):
        sub_shape = map_of_shapes.FindKey(i)
        sub_data = io.BytesIO()
        TopoDS_Shape.DumpJson(sub_shape, sub_data)
        val = sub_data.getvalue()
        val = re.sub(b'"this": "[^"]*"', b'', val)  # Remove memory address
        data.write(val)
    to_hash = data.getvalue()
    return hashlib.md5(to_hash, usedforsecurity=False).hexdigest()