import numpy as np
import time
from math import cos, sin, pi, sqrt
from typing import List, Optional

# --- OCC Core imports ---
from OCC.Core.gp import (
    gp_Pnt, gp_Dir, gp_Ax2, gp_Ax3, gp_Vec, gp_Trsf, gp_Circ
)
from OCC.Core.BRepPrimAPI import (
    BRepPrimAPI_MakeBox,
    BRepPrimAPI_MakeCylinder,
    BRepPrimAPI_MakeSphere
)
from OCC.Core.BRepBuilderAPI import (
    BRepBuilderAPI_MakeEdge, BRepBuilderAPI_MakeWire,
    BRepBuilderAPI_MakeFace, BRepBuilderAPI_Transform
)
from OCC.Core.BRepOffsetAPI import BRepOffsetAPI_MakePipe
from OCC.Core.BRepAlgoAPI import (
    BRepAlgoAPI_Fuse, BRepAlgoAPI_Cut, BRepAlgoAPI_Common
)
from OCC.Core.BRepMesh import BRepMesh_IncrementalMesh
from OCC.Core.BRep import BRep_Tool
from OCC.Core.TopExp import TopExp_Explorer
from OCC.Core.TopAbs import TopAbs_FACE, TopAbs_REVERSED
from OCC.Core.TopoDS import topods_Face
from OCC.Core.TopLoc import TopLoc_Location
from OCC.Core.Bnd import Bnd_Box
from OCC.Core.BRepBndLib import brepbndlib_Add
from OCC.Core.TColgp import TColgp_Array1OfPnt
from OCC.Core.GeomAPI import GeomAPI_PointsToBSpline
from OCC.Core.BRepCheck import BRepCheck_Analyzer


# ============================================================
# 1. 基础几何体构造
# ============================================================
def make_box(x: float, y: float, z: float,
             dx: float, dy: float, dz: float):
    axes = gp_Ax2(gp_Pnt(x, y, z), gp_Dir(0, 0, 1))
    return BRepPrimAPI_MakeBox(axes, dx, dy, dz).Shape()


def make_cylinder(cx: float, cy: float, cz: float,
                  radius: float, height: float,
                  axis_dir=None):
    if axis_dir is None:
        axis_dir = gp_Dir(0, 0, 1)
    axes = gp_Ax2(gp_Pnt(cx, cy, cz), axis_dir)
    return BRepPrimAPI_MakeCylinder(axes, radius, height).Shape()


def make_sphere(cx: float, cy: float, cz: float, radius: float):
    center = gp_Pnt(cx, cy, cz)
    return BRepPrimAPI_MakeSphere(center, radius).Shape()


# ============================================================
# 2. 螺旋线扫掠
# ============================================================
def make_bspline_helix(radius: float, pitch: float,
                       t0: float, t1: float, n_points: int = 128):
    points = TColgp_Array1OfPnt(1, n_points)
    for i in range(n_points):
        t = t0 + (t1 - t0) * i / (n_points - 1)
        points.SetValue(i + 1, gp_Pnt(
            radius * cos(t),
            radius * sin(t),
            pitch * t / (2.0 * pi)
        ))
    approx = GeomAPI_PointsToBSpline(points)
    return approx.Curve()


def make_helix_sweep(helix_r: float, pitch: float, turns: float,
                     section_r: float, n_points: int = 128):
    t0, t1 = 0.0, 2.0 * pi * turns
    helix_curve = make_bspline_helix(helix_r, pitch, t0, t1, n_points)

    helix_edge = BRepBuilderAPI_MakeEdge(helix_curve).Edge()
    helix_wire = BRepBuilderAPI_MakeWire(helix_edge).Wire()

    sec_axes = gp_Ax2(gp_Pnt(helix_r, 0.0, 0.0), gp_Dir(0, 0, 1))
    sec_circ = gp_Circ(sec_axes, section_r)
    sec_edge = BRepBuilderAPI_MakeEdge(sec_circ).Edge()
    sec_wire = BRepBuilderAPI_MakeWire(sec_edge).Wire()
    sec_face = BRepBuilderAPI_MakeFace(sec_wire).Face()

    pipe = BRepOffsetAPI_MakePipe(helix_wire, sec_face)
    pipe.Build()
    if not pipe.IsDone():
        raise RuntimeError("Helix sweep failed")
    return pipe.Shape()


# ============================================================
# 3. 布尔运算
# ============================================================
def boolean_op(a, b, op: str = "fuse"):
    if op == "fuse":
        api = BRepAlgoAPI_Fuse(a, b)
    elif op == "cut":
        api = BRepAlgoAPI_Cut(a, b)
    elif op == "common":
        api = BRepAlgoAPI_Common(a, b)
    else:
        raise ValueError(f"Unknown op: {op}")

    if not api.IsDone():
        raise RuntimeError(f"Boolean {op} failed")

    result = api.Shape()
    analyzer = BRepCheck_Analyzer(result)
    if not analyzer.IsValid():
        print("[Warning] Boolean result has invalid topology")
    return result


def cascade_boolean(shapes: List, op: str = "fuse") -> Optional:
    if not shapes:
        return None
    if len(shapes) == 1:
        return shapes[0]

    current = shapes[:]
    while len(current) > 1:
        nxt = []
        for i in range(0, len(current) - 1, 2):
            nxt.append(boolean_op(current[i], current[i + 1], op))
        if len(current) % 2 == 1:
            nxt.append(current[-1])
        current = nxt
    return current[0]


# ============================================================
# 4. 网格提取
# ============================================================
def extract_triangulation(shape, merge_tol: float = 1e-6):
    vertices = []
    faces = []
    vmap = {}

    def quantize(v: float) -> int:
        return int(round(v / merge_tol))

    explorer = TopExp_Explorer(shape, TopAbs_FACE)
    while explorer.More():
        face = topods_Face(explorer.Current())

        loc = TopLoc_Location()
        tri = BRep_Tool.Triangulation(face, loc)
        if tri is None:
            explorer.Next()
            continue

        trsf = loc.Transformation()
        n_nodes = tri.NbNodes()
        n_tri = tri.NbTriangles()

        local2global = [0] * (n_nodes + 1)

        for i in range(1, n_nodes + 1):
            p = tri.Node(i).Transformed(trsf)
            key = (quantize(p.X()), quantize(p.Y()), quantize(p.Z()))
            if key not in vmap:
                idx = len(vertices)
                vmap[key] = idx
                vertices.append((p.X(), p.Y(), p.Z()))
                local2global[i] = idx
            else:
                local2global[i] = vmap[key]

        for i in range(1, n_tri + 1):
            t = tri.Triangle(i)
            n1, n2, n3 = t.Get()
            if face.Orientation() == TopAbs_REVERSED:
                n1, n2 = n2, n1
            faces.append((
                local2global[n1],
                local2global[n2],
                local2global[n3]
            ))

        explorer.Next()

    return (
        np.array(vertices, dtype=np.float64),
        np.array(faces, dtype=np.int32)
    )


# ============================================================
# 5. 自适应目标面数网格化
# ============================================================
def bounding_diagonal(shape) -> float:
    bbox = Bnd_Box()
    brepbndlib_Add(shape, bbox)
    x1, y1, z1, x2, y2, z2 = bbox.Get()
    return sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2 + (z2 - z1) ** 2)


def mesh_to_target(shape, target_faces: int,
                   ang_deflection: float = 0.5,
                   tol_percent: float = 0.05,
                   max_iter: int = 25):
    diag = bounding_diagonal(shape)
    lo = 1e-4 * diag
    hi = diag

    best = None
    best_diff = float('inf')

    for _ in range(max_iter):
        mid = (lo + hi) * 0.5

        BRepMesh_IncrementalMesh(shape, mid,
                                 False,
                                 ang_deflection,
                                 True)

        verts, faces = extract_triangulation(shape)
        diff = faces.shape[0] - target_faces
        abs_diff = abs(diff)

        if abs_diff < best_diff:
            best_diff = abs_diff
            best = (verts.copy(), faces.copy())

        if abs_diff <= target_faces * tol_percent:
            break

        if diff > 0:
            lo = mid
        else:
            hi = mid

    return best


# ============================================================
# 6. 导出为 Wavefront OBJ
# ============================================================
def write_obj(path: str, vertices: np.ndarray, faces: np.ndarray):
    """
    导出 Wavefront OBJ 格式。
    面索引使用 1-based（OBJ 标准）。
    """
    with open(path, 'w') as f:
        f.write("# Generated by pythonOCC\n")
        for v in vertices:
            f.write(f"v {v[0]:.17g} {v[1]:.17g} {v[2]:.17g}\n")
        for tri in faces:
            # OBJ 为 1-based 索引
            f.write(f"f {tri[0]+1} {tri[1]+1} {tri[2]+1}\n")


# ============================================================
# 7. 主流程（含计时，仅输出 OBJ）
# ============================================================
if __name__ == "__main__":
    t_total_0 = time.perf_counter()

    # --------------------------------------------------------
    # 7.1 几何体构造
    # --------------------------------------------------------
    t_geom_0 = time.perf_counter()

    # base = make_box(0, 0, 0, 100, 80, 20)
    # boss = make_cylinder(50, 40, 20, 15, 30)
    # hole = make_cylinder(50, 40, 0, 8, 60)
    # sphere = make_sphere(50, 40, 55, 12)

    # spring = make_helix_sweep(20, 10, 3, 3)
    # trsf = gp_Trsf()
    # trsf.SetTranslation(gp_Vec(20, 20, 20))
    # spring = BRepBuilderAPI_Transform(spring, trsf).Shape()

    cylinder = make_cylinder(0, 0, 0, 1, 1)
    sphere = make_sphere(1.9, 0, 0, 1)

    t_geom_1 = time.perf_counter()
    print(f"[Timing] Geometry construction: {t_geom_1 - t_geom_0:.4f} s")

    # --------------------------------------------------------
    # 7.2 布尔运算
    # --------------------------------------------------------
    t_bool_0 = time.perf_counter()

    # body = cascade_boolean([base, boss, sphere, spring], "fuse")
    part = boolean_op(cylinder, sphere, "common")

    t_bool_1 = time.perf_counter()
    print(f"[Timing] Boolean operations: {t_bool_1 - t_bool_0:.4f} s")

    # --------------------------------------------------------
    # 7.3 网格化
    # --------------------------------------------------------
    t_mesh_0 = time.perf_counter()

    verts, faces = mesh_to_target(part, target_faces=10000,
                                   ang_deflection=0.5,
                                   tol_percent=0.03)

    t_mesh_1 = time.perf_counter()
    print(f"[Timing] Meshing (target 50k faces): {t_mesh_1 - t_mesh_0:.4f} s")
    print(f"[Result] {verts.shape[0]} vertices, {faces.shape[0]} faces")

    # --------------------------------------------------------
    # 7.4 导出 OBJ
    # --------------------------------------------------------
    t_export_0 = time.perf_counter()

    write_obj("common.obj", verts, faces)

    t_export_1 = time.perf_counter()
    print(f"[Timing] OBJ export: {t_export_1 - t_export_0:.4f} s")

    # --------------------------------------------------------
    # 总计时
    # --------------------------------------------------------
    t_total_1 = time.perf_counter()
    print(f"[Timing] Total elapsed: {t_total_1 - t_total_0:.4f} s")