import logging
import os
import shlex
import subprocess
import tempfile
import time
import sys
import types
import torch
import numpy as np
import rembg
import spaces
import gradio as gr
from PIL import Image
from functools import partial

# --- PARCHE DE CPU (DEBE IR ANTES DE IMPORTAR TSR) ---
try:
    import mcubes
    mock_torchmcubes = types.ModuleType("torchmcubes")
    def marching_cubes_cpu(vertices, threshold):
        v, f = mcubes.marching_cubes(vertices.detach().cpu().numpy(), threshold)
        return torch.from_numpy(v.astype("float32")), torch.from_numpy(f.astype("int64"))
    mock_torchmcubes.marching_cubes = marching_cubes_cpu
    sys.modules["torchmcubes"] = mock_torchmcubes
except ImportError:
    print("Error: PyMCubes no está en requirements.txt")

# --- IMPORTS DE TSR ---
from tsr.system import TSR
from tsr.utils import remove_background, resize_foreground, to_gradio_3d_orientation

HEADER = """# TripoSR Demo
<table bgcolor="#1E2432" cellspacing="0" cellpadding="0"  width="450"><tr style="height:50px;"><td style="text-align: center;"><a href="https://stability.ai"><img src="https://images.squarespace-cdn.com/content/v1/6213c340453c3f502425776e/6c9c4c25-5410-4547-bc26-dc621cdacb25/Stability+AI+logo.png" width="200" height="40" /></a></td><td style="text-align: center;"><a href="https://www.tripo3d.ai"><img src="https://tripo-public.cdn.bcebos.com/logo.png" width="40" height="40" /></a></td></tr></table>
<table bgcolor="#1E2432" cellspacing="0" cellpadding="0"  width="450"><tr style="height:30px;"><td style="text-align: center;"><a href="https://huggingface.co/stabilityai/TripoSR"><img src="https://img.shields.io/badge/%F0%9F%A4%97%20Model_Card-Huggingface-orange" height="20"></a></td><td style="text-align: center;"><a href="https://github.com/VAST-AI-Research/TripoSR"><img src="https://postimage.me/images/2024/03/04/GitHub_Logo_White.png" width="100" height="20"></a></td><td style="text-align: center; color: white;"><a href="https://arxiv.org/abs/2403.02151"><img src="https://img.shields.io/badge/arXiv-2403.02151-b31b1b.svg" height="20"></a></td></tr></table>
> Try our new model: **SF3D** with several improvements such as faster generation and more game-ready assets.
> The model is available [here](https://huggingface.co/stabilityai/stable-fast-3d) and we also have a [demo](https://huggingface.co/spaces/stabilityai/stable-fast-3d). 
**TripoSR** is a state-of-the-art open-source model for **fast** feedforward 3D reconstruction from a single image.
"""

if torch.cuda.is_available():
    device = "cuda:0"
else:
    device = "cpu"

model = TSR.from_pretrained(
    "stabilityai/TripoSR",
    config_name="config.yaml",
    weight_name="model.ckpt",
)
model.renderer.set_chunk_size(131072)
model.to(device)
rembg_session = rembg.new_session()

def check_input_image(input_image):
    if input_image is None:
        raise gr.Error("No image uploaded!")

def preprocess(input_image, do_remove_background, foreground_ratio):
    def fill_background(image):
        image = np.array(image).astype(np.float32) / 255.0
        image = image[:, :, :3] * image[:, :, 3:4] + (1 - image[:, :, 3:4]) * 0.5
        image = Image.fromarray((image * 255.0).astype(np.uint8))
        return image
    if do_remove_background:
        image = input_image.convert("RGB")
        image = remove_background(image, rembg_session)
        image = resize_foreground(image, foreground_ratio)
        image = fill_background(image)
    else:
        image = input_image
        if image.mode == "RGBA":
            image = fill_background(image)
    return image

@spaces.GPU
def generate(image, mc_resolution, formats=["obj", "glb"]):
    scene_codes = model(image, device=device)
    mesh = model.extract_mesh(scene_codes, resolution=mc_resolution)[0]
    mesh = to_gradio_3d_orientation(mesh)
    mesh_path_glb = tempfile.NamedTemporaryFile(suffix=f".glb", delete=False)
    mesh.export(mesh_path_glb.name)
    mesh_path_obj = tempfile.NamedTemporaryFile(suffix=f".obj", delete=False)
    mesh.apply_scale([-1, 1, 1])
    mesh.export(mesh_path_obj.name)
    return mesh_path_obj.name, mesh_path_glb.name

def run_example(image_pil):
    preprocessed = preprocess(image_pil, False, 0.9)
    mesh_name_obj, mesh_name_glb = generate(preprocessed, 256, ["obj", "glb"])
    return preprocessed, mesh_name_obj, mesh_name_glb

with gr.Blocks() as demo:
    gr.Markdown(HEADER)
    with gr.Row(variant="panel"):
        with gr.Column():
            with gr.Row():
                input_image = gr.Image(
                    label="Input Image",
                    image_mode="RGBA",
                    sources="upload",
                    type="pil",
                    elem_id="content_image",
                )
                processed_image = gr.Image(label="Processed Image", interactive=False)
            with gr.Row():
                with gr.Group():
                    do_remove_background = gr.Checkbox(
                        label="Remove Background", value=True
                    )
                    foreground_ratio = gr.Slider(
                        label="Foreground Ratio",
                        minimum=0.5,
                        maximum=1.0,
                        value=0.85,
                        step=0.05,
                    )
                    mc_resolution = gr.Slider(
                        label="Marching Cubes Resolution",
                        minimum=32,
                        maximum=320,
                        value=256,
                        step=32 
                    )
            with gr.Row():
                submit = gr.Button("Generate", elem_id="generate", variant="primary")
        with gr.Column():
            with gr.Tab("OBJ"):
                output_model_obj = gr.Model3D(
                    label="Output Model (OBJ Format)",
                    interactive=False,
                )
            with gr.Tab("GLB"):
                output_model_glb = gr.Model3D(
                    label="Output Model (GLB Format)",
                    interactive=False,
                )
    
    if os.path.exists("examples"):
        with gr.Row(variant="panel"):
            gr.Examples(
                examples=[
                    os.path.join("examples", img_name) for img_name in sorted(os.listdir("examples"))
                ] if os.path.exists("examples") else [],
                inputs=[input_image],
                outputs=[processed_image, output_model_obj, output_model_glb],
                cache_examples=False,
                fn=partial(run_example),
                label="Examples",
                examples_per_page=20
            )

    submit.click(fn=check_input_image, inputs=[input_image]).success(
        fn=preprocess,
        inputs=[input_image, do_remove_background, foreground_ratio],
        outputs=[processed_image],
    ).success(
        fn=generate,
        inputs=[processed_image, mc_resolution],
        outputs=[output_model_obj, output_model_glb],
    )

demo.queue(max_size=10)
demo.launch()