import gradio as gr
import subprocess
import os
import uuid
import shutil
import tempfile

DESCRIPTION = """# SHARP · 3D from a Single Photo
Upload any photo and get a **3D Gaussian Splat (.ply)** in seconds.
Powered by Apple's [SHARP](https://github.com/apple/ml-sharp) monocular view synthesis model.

**How to use:**
1. Upload or drag & drop an image
2. Click **Generate 3D Splat**
3. Download the `.ply` file
4. View it at [SuperSplat](https://playcanvas.com/supersplat/editor)

> Running on CPU - first run downloads the 2.6GB model and may take 5-10 min."""

def generate_splat(image_path):
    if image_path is None:
        raise gr.Error("Please upload an image first.")
    job_id = str(uuid.uuid4())
    input_dir = os.path.join(tempfile.gettempdir(), f"sharp_in_{job_id}")
    output_dir = os.path.join(tempfile.gettempdir(), f"sharp_out_{job_id}")
    os.makedirs(input_dir, exist_ok=True)
    os.makedirs(output_dir, exist_ok=True)
    try:
        ext = os.path.splitext(image_path)[1] or ".jpg"
        shutil.copy(image_path, os.path.join(input_dir, f"input{ext}"))
        result = subprocess.run(
            ["sharp", "predict", "-i", input_dir, "-o", output_dir],
            capture_output=True, text=True, timeout=600
        )
        if result.returncode != 0:
            raise gr.Error(f"SHARP failed: {result.stderr[-500:]}")
        ply_files = [f for f in os.listdir(output_dir) if f.endswith(".ply")]
        if not ply_files:
            raise gr.Error("No .ply file generated. Try a different image.")
        out_path = os.path.join(tempfile.gettempdir(), f"output_{job_id}.ply")
        shutil.copy(os.path.join(output_dir, ply_files[0]), out_path)
        return out_path, "Done! Download your .ply file above, then open it in SuperSplat."
    except subprocess.TimeoutExpired:
        raise gr.Error("Timed out after 10 minutes.")
    finally:
        shutil.rmtree(input_dir, ignore_errors=True)
        shutil.rmtree(output_dir, ignore_errors=True)

with gr.Blocks(title="SHARP 3D") as demo:
    gr.Markdown(DESCRIPTION)
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="filepath", label="Input Image")
            run_btn = gr.Button("Generate 3D Splat", variant="primary")
        with gr.Column():
            file_output = gr.File(label="Download .ply file")
            status_output = gr.Markdown("")
    run_btn.click(fn=generate_splat, inputs=image_input, outputs=[file_output, status_output])

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