"""
Gradio app: robust .vox → .glb conversion + built-in Model3D viewer.
Handles small and large voxel files, centers and colors mesh to avoid 'paper icon'.
"""

import gradio as gr
import tempfile
import trimesh
import numpy as np
from pathlib import Path
from pyvox.parser import VoxParser


def vox_to_glb(file_path):
    vox = VoxParser(file_path).parse()
    model = vox.models[0]
    voxels = model.voxels
    size_x, size_y, size_z = model.size.x, model.size.y, model.size.z

    if not voxels:
        raise ValueError("No voxels found")

    # Small grid: cube per voxel, Large grid: marching cubes
    if max(size_x, size_y, size_z) <= 16:
        cubes = []
        for v in voxels:
            cube = trimesh.creation.box(extents=(1,1,1))
            cube.apply_translation([v.x, v.y, v.z])
            cubes.append(cube)
        mesh = trimesh.util.concatenate(cubes)
    else:
        grid = np.zeros((size_x, size_y, size_z), dtype=bool)
        for v in voxels:
            grid[v.x, v.y, v.z] = True
        mesh = trimesh.voxel.ops.matrix_to_marching_cubes(grid)

    # Center mesh, scale, add color
    mesh.apply_translation(-mesh.centroid)
    mesh.apply_scale(1.0)
    mesh.visual.vertex_colors = [200, 100, 50, 255]

    tmp = tempfile.NamedTemporaryFile(suffix=".glb", delete=False)
    mesh.export(tmp.name)
    return tmp.name


def handle_upload(file):
    if file is None:
        return None
    ext = Path(file.name).suffix.lower()
    if ext == '.vox':
        try:
            return vox_to_glb(file.name)
        except Exception as e:
            print(f"Failed to convert vox: {e}")
            return None
    elif ext in ['.glb', '.gltf']:
        return file.name
    else:
        return None


with gr.Blocks() as demo:
    gr.Markdown("# VOX / GLB Viewer\nUpload a `.vox` or `.glb/.gltf` file to preview.")
    upload = gr.File(file_types=['.vox','.glb','.gltf'], label="Upload File")
    viewer = gr.Model3D(label="Preview")

    upload.change(handle_upload, upload, viewer)


if __name__ == '__main__':
    demo.launch()