import os
import trimesh
import numpy as np


def crop_mesh(mesh, desired_bounding_box_length=15, viz=False):
    print("Cropping mesh...")
    mesh_bounds = mesh.bounding_box.bounds
    center = (mesh_bounds[0] + mesh_bounds[1]) / 2
    bbox_half_size = desired_bounding_box_length / 2

    cropped_mesh = mesh
    for axis in [(1, 0, 0), (0, 1, 0), (0, 0, 1)]:  # Iterate through X, Y, Z axes
        plane_origin = center + bbox_half_size * np.array(axis)  # Convert axis to NumPy array for multiplication
        plane_normal = - np.array(axis)  # Negate the NumPy array version of axis
        cropped_mesh = cropped_mesh.slice_plane(plane_origin, plane_normal)

    if viz:
        print("Visualizing cropped mesh...")
        cropped_mesh.show()

    return cropped_mesh


def remesh_mesh(mesh, viz=False):
    print("Remeshing...")
    # Concat the uv with the vertices and subdivide
    vertices, faces = trimesh.remesh.subdivide(
        vertices=np.hstack((mesh.vertices, mesh.visual.uv.copy())),
        faces=mesh.faces
    )

    mesh.vertices = vertices[:, :3]
    mesh.faces = faces
    mesh.visual.uv = vertices[:, 3:]

    if viz:
        print("Visualizing remeshed mesh...")
        mesh.show()

    return mesh


def repair_mesh(mesh, viz=False):
    print("Repairing...")
    # Access repair functions from the 'repair' module:
    trimesh.repair.fix_winding(mesh)  # Ensure consistent winding order first
    trimesh.repair.fix_normals(mesh)  # Recalculate normals based on winding
    trimesh.repair.fix_inversion(mesh)  # Fix inverted faces
    trimesh.repair.fill_holes(mesh)  # Fill holes after fixing topology
    trimesh.repair.broken_faces(mesh)  # Handle broken faces after filling holes
    # trimesh.repair.stitch(mesh)  # Stitch disconnected components (optional)

    if viz:
        mesh.show()

    return mesh


def post_process_textured_mesh(mesh_path, out_path, crop=False, remesh=False, repair=True, viz=False, save=True):
    print("Loading mesh...")
    mesh = trimesh.load(mesh_path)

    if viz:
        mesh.show()

    # crop mesh
    if crop:
        mesh = crop_mesh(mesh)

    # remesh and repair and remesh
    if remesh:
        mesh = remesh_mesh(mesh)

    if repair:
        mesh = repair_mesh(mesh)

    # visualize
    if viz:
        print("Visualizing processed mesh...")
        mesh.show()

    # save the cropped mesh and texture
    if save:
        print("Saving mesh...")
        mesh.export(out_path)

    print("Finished!")


if __name__ == "__main__":
    mesh_dir = "output/refined_mesh/custom/"
    input_mesh_path = os.path.join(mesh_dir, "sugarfine_3Dgs7000_densityestim02_sdfnorm02_level03_decim1000000_normalconsistency01_gaussperface6.obj")

    out_dir = "output/final_mesh/custom/"
    os.makedirs(out_dir, exist_ok=True)
    output_mesh_path = os.path.join(out_dir, 'custom.glb')

    post_process_textured_mesh(input_mesh_path, output_mesh_path, crop=False, remesh=False, repair=True, viz=False, save=True)