import trimesh
import os
import numpy as np
import argparse
import traceback


def convert_obj_glb(source_dir):
    """
    Convert all OBJ files in the given source directory to a single GLB file.

    Args:
        source_dir: Directory containing OBJ files
        output_visual_path: Path to the output directory
        output_file: Output GLB file name (default: base0.glb)

    Returns:
        bool: True if successful, False if an error occurs
    """
    try:
        texture_dir = os.path.join(source_dir, "textured_objs")
        visual_dir = os.path.join(source_dir, "visual")
        output_path = os.path.join(visual_dir, "base0.glb")
        if os.path.exists(output_path):
            print(f"File {output_path} already exists")
            return True
        if not os.path.exists(visual_dir):
            os.makedirs(visual_dir)
        # Create a scene to hold all meshes
        scene = trimesh.Scene()

        # Find all .obj files in the directory
        obj_files = [f for f in os.listdir(texture_dir) if f.endswith(".obj")]

        # Load each OBJ file and add it to the scene
        for obj_file in obj_files:
            file_path = os.path.join(texture_dir, obj_file)
            try:
                with open(file_path, "rb") as file_obj:
                    mesh = trimesh.load(file_obj, file_type="obj")
                scene.add_geometry(mesh)
                # print(f"Added mesh from {file_path}")
            except Exception as e:
                print(f"Error loading {file_path}: {e}")
                return False

        # Export the scene as GLB
        print(f"Exporting scene to {output_path}...")
        scene.export(output_path)
        print(f"Model successfully exported to {output_path}")
        return True
    except Exception as e:
        print(f"An error occurred in convert_to_glb: {e}" + traceback.format_exc())
        return False


def is_digital(name):
    """Check if a string contains only digits."""
    return name.isdigit()


def has_only_digital_subdirs(directory):
    """Check if a directory contains only subdirectories with digital names."""
    if not os.path.isdir(directory):
        return False

    subdirs = [item for item in os.listdir(directory) if os.path.isdir(os.path.join(directory, item))]

    # Return True if there are subdirs and all of them are digital
    return len(subdirs) > 0 and all(is_digital(subdir) for subdir in subdirs)


if __name__ == "__main__":
    # Set up argument parser
    parser = argparse.ArgumentParser(description="Convert OBJ files to GLB.")
    parser.add_argument(
        "--object_dir",
        type=str,
        help="Directory containing single object (e.g., assets/objects/060_kitchenpot)",
    )
    parser.add_argument(
        "--scan_all",
        action="store_true",
        help="Scan all objects in assets/objects directory",
    )
    args = parser.parse_args()

    total_conversions = 0

    assets_path = "../assets/objects"
    # Process each object directory in assets/objects
    for obj_dir in os.listdir(assets_path):
        obj_path = os.path.join(assets_path, obj_dir)

        # Check if it's a directory and has only digital subdirectories
        if os.path.isdir(obj_path) and has_only_digital_subdirs(obj_path):
            print(obj_path)
            # for final_path in os.listdir(obj_path):
            #     convert_obj_glb(os.path.join(obj_path, final_path))

    print(f"\nTotal completed GLB conversions: {total_conversions}")
