import gradio as gr
import requests
import tempfile
import os
import trimesh

def download_model(url):
    """Download model from URL to temporary file."""
    try:
        response = requests.get(url, stream=True, timeout=10)
        response.raise_for_status()
        with tempfile.NamedTemporaryFile(suffix=".glb", delete=False) as tmp:
            for chunk in response.iter_content(chunk_size=8192):
                tmp.write(chunk)
            return tmp.name
    except Exception as e:
        print(f"Download error: {e}")
        return None

def load_and_convert_model(url, target_format):
    """Load model from URL and convert to target format."""
    if not url:
        return None
    
    # Download the GLB file
    glb_path = download_model(url)
    if not glb_path:
        return None
    
    try:
        # Load the model using trimesh
        mesh = trimesh.load(glb_path)
        
        # Save to temporary file in target format
        output_ext = target_format.lower()
        with tempfile.NamedTemporaryFile(suffix=f".{output_ext}", delete=False) as tmp:
            output_path = tmp.name
        
        # Export based on format
        if output_ext == "obj":
            mesh.export(output_path, file_type="obj")
        elif output_ext == "stl":
            mesh.export(output_path, file_type="stl")
        elif output_ext == "ply":
            mesh.export(output_path, file_type="ply")
        elif output_ext == "glb":
            mesh.export(output_path, file_type="glb")
        elif output_ext == "gltf":
            mesh.export(output_path, file_type="gltf")
        else:
            # For unsupported formats, return None
            return None
        
        return output_path
    except Exception as e:
        print(f"Conversion error: {e}")
        return None
    finally:
        # Clean up original downloaded file
        if os.path.exists(glb_path):
            os.unlink(glb_path)

def cleanup(file_path):
    if file_path and os.path.exists(file_path):
        os.unlink(file_path)

# Gradio interface
with gr.Blocks(title="3D Model Converter") as demo:
    gr.Markdown("## 🌐 3D Model Viewer & Converter")
    gr.Markdown("Enter a URL to a GLB file to view and download in different formats.")
    
    with gr.Row():
        url_input = gr.Textbox(
            label="GLB File URL",
            placeholder="https://example.com/model.glb",
            scale=4
        )
        submit_btn = gr.Button("Load Model", variant="primary", scale=1)
    
    model_output = gr.Model3D(
        label="3D Model Preview",
        clear_color=[0.2, 0.2, 0.2, 0.8],
        camera_position=(2, 2, 3),
        scale=1.0
    )
    
    gr.Markdown("### 💾 Download Options")
    with gr.Row():
        obj_btn = gr.Button("Download as OBJ", variant="secondary")
        stl_btn = gr.Button("Download as STL", variant="secondary")
        ply_btn = gr.Button("Download as PLY", variant="secondary")
        glb_btn = gr.Button("Download as GLB", variant="secondary")
    
    obj_file = gr.File(label="Download OBJ", visible=False)
    stl_file = gr.File(label="Download STL", visible=False)
    ply_file = gr.File(label="Download PLY", visible=False)
    glb_file = gr.File(label="Download GLB", visible=False)
    
    # State to store the current URL
    current_url = gr.State()
    
    # When loading, store URL and display model
    submit_btn.click(
        fn=lambda url: (url, download_model(url)),
        inputs=url_input,
        outputs=[current_url, model_output]
    )
    
    # Conversion functions
    def convert_and_download(url, format_type):
        if not url:
            return None
        output_path = load_and_convert_model(url, format_type)
        if output_path:
            return output_path
        return None
    
    # Wire up download buttons
    obj_btn.click(
        fn=lambda url: convert_and_download(url, "obj"),
        inputs=current_url,
        outputs=obj_file
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=obj_file
    )
    
    stl_btn.click(
        fn=lambda url: convert_and_download(url, "stl"),
        inputs=current_url,
        outputs=stl_file
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=stl_file
    )
    
    ply_btn.click(
        fn=lambda url: convert_and_download(url, "ply"),
        inputs=current_url,
        outputs=ply_file
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=ply_file
    )
    
    glb_btn.click(
        fn=lambda url: convert_and_download(url, "glb"),
        inputs=current_url,
        outputs=glb_file
    ).then(
        fn=lambda: gr.update(visible=True),
        outputs=glb_file
    )

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