import re
import subprocess
import gradio as gr
import trimesh
import tempfile
import time
import pathlib
import os

os.system("pip install -e ./simple-knn")
os.system("pip install -e ./diff-gaussian-rasterization")

TEMP_DIR = "/tmp/dreamgaussian"
os.makedirs(TEMP_DIR, exist_ok=True)


def create_from_text(prompt):
    temp_dir = tempfile.mkdtemp(dir=TEMP_DIR)
    sanitized_prompt = re.sub("[^0-9a-zA-Z]+", "_", prompt)
    cmd1 = [
        "python",
        "main.py",
        "--config",
        "configs/text.yaml",
        f"prompt={prompt}",
        f"outdir={temp_dir}",
        f"save_path={sanitized_prompt}",
        "force_cuda_rast=True",
        "mesh_format=glb",
    ]
    cmd2 = [
        "python",
        "main2.py",
        "--config",
        "configs/text.yaml",
        f"prompt={prompt}",
        f"outdir={temp_dir}",
        f"save_path={sanitized_prompt}",
        "force_cuda_rast=True",
        "mesh_format=glb",
    ]
    subprocess.run(cmd1)
    subprocess.run(cmd2)
    glb_path = f"{temp_dir}/{sanitized_prompt}.glb"
    return [get_html_model(glb_path), glb_path]


def create_from_image(image):
    temp_dir = tempfile.mkdtemp(dir=TEMP_DIR)
    sanitized_prompt = "image"
    image.save(f"{temp_dir}/{sanitized_prompt}.png")
    cmd1 = [
        "python",
        "process.py",
        f"{temp_dir}/{sanitized_prompt}.png",
        "--size",
        "512",
    ]  # <<< 256 may be a good default
    cmd2 = [
        "python",
        "main.py",
        "--config",
        "configs/image.yaml",
        f"outdir={temp_dir}",
        f"input={temp_dir}/{sanitized_prompt}_rgba.png",
        f"save_path={sanitized_prompt}",
        "force_cuda_rast=True",
        "mesh_format=glb",
    ]
    cmd3 = [
        "python",
        "main2.py",
        "--config",
        "configs/image.yaml",
        f"outdir={temp_dir}",
        f"input={temp_dir}/{sanitized_prompt}_rgba.png",
        f"save_path={sanitized_prompt}",
        "force_cuda_rast=True",
        "mesh_format=glb",
    ]
    print(cmd1)
    print(cmd2)
    print(cmd3)
    subprocess.run(cmd1)
    subprocess.run(cmd2)
    subprocess.run(cmd3)
    glb_path = f"{temp_dir}/{sanitized_prompt}.glb"
    return [get_html_model(glb_path), glb_path]


def get_html_model(f_path):
    iframe = f"""<iframe src="file=model.html" model-url="file={f_path}" width="100%" height="500px"></iframe>"""

    return iframe


def generate(prompt, image):
    if prompt:
        gr.Info("Generating from prompt")
        return create_from_text(prompt)
    elif image:
        gr.Info("Generating from image")
        return create_from_image(image)
    raise gr.Error("Please enter a prompt or upload an image.")


with gr.Blocks() as demo:
    gr.Markdown(
        """
# DreamGaussian
source: https://github.com/dreamgaussian/dreamgaussian
"""
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(lines=2, placeholder="Enter Prompt...")
            image = gr.Image(type="pil", label="Image")
            with gr.Row():
                clear = gr.Button("Clear")
                btn = gr.Button("Generate")
        with gr.Column():
            with gr.Box():
                model_3d = gr.HTML(
                    value="<div style='min-height: 240px;'>",
                    label="3D Model Viewer",
                    show_label=True,
                )
            download = gr.File(label="Download 3D file")
    with gr.Row():
        with gr.Column():
          with gr.Row():
            gr.Examples(
                [["pikachu pokemon"]],
                fn=lambda x: generate(x, None),
                inputs=[prompt],
                outputs=[model_3d, download],
                cache_examples=False,
            )
            gr.Examples(
                [["./examples/b.jpg"]],
                fn=lambda x: generate(None, x),
                inputs=[image],
                outputs=[model_3d, download],
                cache_examples=False,
            )
        with gr.Column():  
          pass
    btn.click(generate, inputs=[prompt, image], outputs=[model_3d, download])
    clear.click(
        lambda x: (gr.update(value=None), gr.update(value=None)),
        None,
        [prompt, image],
        queue=False,
    )

demo.queue(api_open=False, concurrency_count=1)
demo.launch(
    debug=True, show_api=False, inline=False, share=True, allowed_paths=[TEMP_DIR]
)