import gradio as gr
from gradio_client import Client, handle_file

# Create a client for the external 3D generation API hosted on "tencent/Hunyuan3D-2"
api_client = Client("tencent/Hunyuan3D-2")

def generate_3d(caption: str, image, steps: float, guidance_scale: float, seed: float, octree_resolution: str, check_box_rembg: bool):
    """
    Calls the external /generation_all API endpoint to generate 3D outputs.
    Returns a tuple of:
      [0] White Mesh file (filepath),
      [1] Textured Mesh file (filepath),
      [2] HTML output string,
      [3] Additional HTML output string.
    """
    # If no image is provided, use a default image URL
    if image is None:
        image = "https://raw.githubusercontent.com/gradio-app/gradio/main/test/test_files/bus.png"
    
    image_path = handle_file(image)
    
    result = api_client.predict(
        caption=caption,
        image=image_path,
        steps=steps,
        guidance_scale=guidance_scale,
        seed=seed,
        octree_resolution=octree_resolution,
        check_box_rembg=check_box_rembg,
        api_name="/generation_all"
    )
    return result

# Build the UI using Gradio Blocks for a more organized layout.
with gr.Blocks(title="3D Generator") as demo:
    gr.Markdown("# 3D Generator")
    gr.Markdown("Generate 3D models using the external `/generation_all` API from `tencent/Hunyuan3D-2`.")
    
    with gr.Row():
        with gr.Column():
            caption_input = gr.Textbox(
                label="Text Prompt (Caption)",
                placeholder="Enter your caption here",
                lines=2
            )
            image_input = gr.Image(
                label="Input Image",
                type="filepath"
            )
        with gr.Column():
            steps_input = gr.Slider(
                minimum=1, maximum=100, step=1, value=50,
                label="Inference Steps"
            )
            guidance_scale_input = gr.Slider(
                minimum=1.0, maximum=10.0, step=0.1, value=5.5,
                label="Guidance Scale"
            )
            seed_input = gr.Slider(
                minimum=0, maximum=10000, step=1, value=1234,
                label="Seed"
            )
            octree_resolution_input = gr.Dropdown(
                choices=["256", "384", "512"],
                value="256",
                label="Octree Resolution"
            )
            check_box_rembg_input = gr.Checkbox(
                label="Remove Background",
                value=True
            )
    
    generate_button = gr.Button("Generate 3D Model")
    
    gr.Markdown("### Generated Outputs")
    with gr.Row():
        white_mesh_output = gr.File(label="Download White Mesh")
        textured_mesh_output = gr.File(label="Download Textured Mesh")
    output_html1 = gr.HTML(label="Output HTML 1")
    output_html2 = gr.HTML(label="Output HTML 2")
    
    # Connect the button click event to the generate_3d function.
    generate_button.click(
        fn=generate_3d,
        inputs=[
            caption_input, image_input, steps_input,
            guidance_scale_input, seed_input,
            octree_resolution_input, check_box_rembg_input
        ],
        outputs=[
            white_mesh_output, textured_mesh_output,
            output_html1, output_html2
        ]
    )
    
    gr.Markdown("## Enjoy your generated 3D models!")

if __name__ == "__main__":
    # Launch the app with a public share link (if desired)
    demo.launch(share=True)