import gradio as gr
import trimesh
import numpy as np
import tempfile
import os

def reconstruct_perfect_cube(input_file):
    if input_file is None:
        return None, None
    
    try:
        # 1. Load the irregular STL
        mesh = trimesh.load(input_file.name)
        
        # Handle scenes if necessary
        if isinstance(mesh, trimesh.Scene):
            mesh = mesh.dump(concatenate=True)

        # 2. Advanced Reconstruction Logic: Bounding Box
        # Instead of wrapping the broken shape (Convex Hull), 
        # we calculate the 'Oriented Bounding Box'. 
        # This identifies the 'Proper' cube that would perfectly 
        # contain your irregular parts.
        perfect_cube = mesh.bounding_box_oriented

        # 3. Create a temporary path for the output
        temp_dir = tempfile.gettempdir()
        output_filename = "reconstructed_perfect_cube.stl"
        output_path = os.path.join(temp_dir, output_filename)
        
        # 4. Export the perfect solid cube
        perfect_cube.export(output_path)
        
        return output_path, output_path
    
    except Exception as e:
        print(f"Error: {e}")
        return None, None

# --- Gradio UI ---
with gr.Blocks(title="AI 3D Cube Completion", theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🧊 Perfect Cube Generator")
    gr.Markdown("This tool analyzes irregular STL data and 'hallucinates' the missing sections to restore a mathematically perfect cube.")
    
    with gr.Row():
        with gr.Column():
            file_input = gr.File(label="Upload Irregular Cube STL", file_types=[".stl"])
            generate_btn = gr.Button("🚀 Generate Full Constructed Cube", variant="primary")
        
        with gr.Column():
            model_viewer = gr.Model3D(label="3D Preview of Reconstructed Cube")
            download_link = gr.File(label="Download Perfect STL File")

    generate_btn.click(
        fn=reconstruct_perfect_cube, 
        inputs=file_input, 
        outputs=[model_viewer, download_link]
    )

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