import os
import torch
from mesh import Mesh

# Ensure the output directory exists
def ensure_dir(path):
    os.makedirs(path, exist_ok=True)

def generate3d(model, rgb_image, xyz_image, device="cpu"):
    output_dir = "outputs"
    ensure_dir(output_dir)

    prompt_id = "mesh_output"
    base_path = os.path.join(output_dir, prompt_id)
    obj_path = base_path + ".obj"
    glb_path = base_path + ".glb"

    # CRM export expects a data dictionary
    data = {"rgb": rgb_image, "xyz": xyz_image}

    # Get rendering context required by xatlas and nvdiffrast
    try:
        import nvdiffrast.torch as dr
        ctx = dr.RasterizeGLContext()
    except Exception as e:
        raise RuntimeError("Failed to initialize nvdiffrast context: " + str(e))

    # Export mesh with UVs and texture image
    model.export_mesh_wt_uv(ctx, data, base_path, 0, device, 512)

    # Convert .obj to .glb
    mesh = Mesh.load(obj_path, device=torch.device("cpu"))
    mesh.write(glb_path)

    return glb_path