from __future__ import annotations

import time
from pathlib import Path

import numpy as np


class CADConverter:
    """Convert STEP CAD files to STL meshes."""

    def read_step_shape(self, input_step: str | Path):
        try:
            from OCC.Core.IFSelect import IFSelect_RetDone
            from OCC.Core.STEPControl import STEPControl_Reader
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "pythonocc-core is required for CAD conversion. "
                "Install it with conda-forge before running CAD workflows."
            ) from exc

        input_step = Path(input_step)
        step_reader = STEPControl_Reader()
        if step_reader.ReadFile(str(input_step)) != IFSelect_RetDone:
            raise ValueError(f"Failed to read STEP file: {input_step}")
        step_reader.TransferRoots()
        return step_reader.OneShape()

    def step_to_numpy_mesh(
        self,
        input_step: str | Path,
        mesh_deflection: float = 0.1,
    ) -> tuple[np.ndarray, np.ndarray]:
        try:
            from OCC.Core.BRep import BRep_Tool
            from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
            from OCC.Core.TopAbs import TopAbs_FACE
            from OCC.Core.TopExp import TopExp_Explorer

            try:
                from OCC.Core.TopoDS import topods
            except ImportError:
                from OCC.Core.TopoDS import TopoDS

                topods = None
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "pythonocc-core is required for CAD mesh extraction. "
                "Install it with conda-forge before running CAD workflows."
            ) from exc

        shape = self.read_step_shape(input_step)
        mesh = BRepMesh_IncrementalMesh(shape, mesh_deflection, False)
        mesh.Perform()
        if not mesh.IsDone():
            raise RuntimeError("CAD mesh discretization failed.")

        vertices: list[list[float]] = []
        faces: list[list[int]] = []
        explorer = TopExp_Explorer(shape, TopAbs_FACE)

        while explorer.More():
            current_shape = explorer.Current()
            face = topods.Face(current_shape) if topods is not None else TopoDS.Face(current_shape)
            location = face.Location()
            triangulation = BRep_Tool.Triangulation(face, location)
            if triangulation is None:
                explorer.Next()
                continue

            transform = location.Transformation()
            base_index = len(vertices)
            for index in range(1, triangulation.NbNodes() + 1):
                point = triangulation.Node(index)
                point.Transform(transform)
                vertices.append([point.X(), point.Y(), point.Z()])

            for index in range(1, triangulation.NbTriangles() + 1):
                n1, n2, n3 = triangulation.Triangle(index).Get()
                faces.append([base_index + n1 - 1, base_index + n2 - 1, base_index + n3 - 1])

            explorer.Next()

        vertices_np = np.asarray(vertices, dtype=np.float64)
        faces_np = np.asarray(faces, dtype=np.int64)
        if len(vertices_np) == 0 or len(faces_np) == 0:
            raise RuntimeError("CAD triangle mesh extraction failed; try a smaller mesh deflection.")
        return vertices_np, faces_np

    def convert_step_to_stl(self, input_step: str | Path, output_stl: str | Path, precision: float = 1.0) -> Path:
        try:
            from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
            from OCC.Core.StlAPI import StlAPI_Writer
        except ModuleNotFoundError as exc:
            raise RuntimeError(
                "pythonocc-core is required for CAD conversion. "
                "Install it with conda-forge before running cad2mesh."
            ) from exc

        input_step = Path(input_step)
        output_stl = Path(output_stl)
        output_stl.parent.mkdir(parents=True, exist_ok=True)

        start = time.time()
        print(f"[1/3] Reading STEP: {input_step}")
        shape = self.read_step_shape(input_step)

        print(f"[2/3] Meshing with precision {precision} mm")
        mesh = BRepMesh_IncrementalMesh(shape, precision, False)
        mesh.Perform()
        if not mesh.IsDone():
            raise RuntimeError("Mesh discretization failed.")

        print(f"[3/3] Writing STL: {output_stl}")
        writer = StlAPI_Writer()
        writer.Write(shape, str(output_stl))
        print(f"CAD conversion finished in {time.time() - start:.2f}s")
        return output_stl