import subprocess
import sys
import os

# Clona TripoSG se non presente (non installabile via pip)
TRIPOSG_PATH = "/app/TripoSG"
if not os.path.exists(TRIPOSG_PATH):
    subprocess.run(
        ["git", "clone", "--depth=1",
         "https://github.com/VAST-AI-Research/TripoSG.git",
         TRIPOSG_PATH],
        check=True
    )
# Aggiungi al path Python
sys.path.insert(0, TRIPOSG_PATH)

import gradio as gr
import torch
import spaces
import tempfile
from PIL import Image
from triposg.pipelines.pipeline_triposg import TripoSGPipeline

print("Caricamento del modello VAST-AI/TripoSG...")
pipe = TripoSGPipeline.from_pretrained(
    "VAST-AI/TripoSG",
    torch_dtype=torch.float16,
)

@spaces.GPU(duration=120)
def generate_3d_mesh(input_image: Image.Image):
    if input_image is None:
        raise gr.Error("Carica un'immagine valida prima di generare il modello.")

    pipe.to("cuda")
    input_image = input_image.convert("RGB")

    generator = torch.Generator(device="cuda").manual_seed(42)
    output = pipe(input_image, generator=generator)

    if hasattr(output, "meshes"):
        mesh = output.meshes[0]
    elif hasattr(output, "mesh"):
        mesh = output.mesh[0]
    else:
        raise ValueError("Nessuna mesh trovata nell'output della pipeline.")

    temp_dir = tempfile.mkdtemp()
    glb_path = os.path.join(temp_dir, "output.glb")
    mesh.export(glb_path)

    return glb_path

with gr.Blocks(title="Generatore 3D - TripoSG", theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
        # 🚀 TripoSG: Da Immagine a Modello 3D
        Utilizza **VAST-AI/TripoSG** per generare modelli 3D ad altissima fedeltà.  
        *Consiglio: usa immagini con sfondo bianco o trasparente.*
    """)
    with gr.Row():
        with gr.Column():
            img_input = gr.Image(type="pil", label="Carica la tua Immagine")
            gen_btn = gr.Button("🎨 Genera Modello 3D", variant="primary")
        with gr.Column():
            model_output = gr.Model3D(
                label="Modello Generato",
                clear_color=[0.0, 0.0, 0.0, 0.0]
            )

    gen_btn.click(fn=generate_3d_mesh, inputs=[img_input], outputs=[model_output])

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