import gradio as gr
import os

# 1. Define directory paths
IMG_DIR = "example_image"  # Directory for images
MODEL_DIR = "f3c"          # Directory for 3D models

# 2. Automatically scan and match files
examples = []

# Check if directories exist to avoid errors
if os.path.exists(IMG_DIR) and os.path.exists(MODEL_DIR):
    # Get all png images, sorted
    image_files = sorted([f for f in os.listdir(IMG_DIR) if f.endswith('.png')])
    
    for img_file in image_files:
        # Get filename without extension (e.g., "6" or "typical_creature_dragon")
        base_name = os.path.splitext(img_file)[0]
        
        # Construct the expected GLB filename: rule is "filename_demo_trellis.glb"
        expected_glb_name = f"{base_name}_demo_trellis.glb"
        
        # Construct full paths
        img_path = os.path.join(IMG_DIR, img_file)
        glb_path = os.path.join(MODEL_DIR, expected_glb_name)
        
        # Critical step: Only add to the list if the corresponding 3D file actually exists
        if os.path.exists(glb_path):
            examples.append({
                "image": img_path,
                "model": glb_path,
                "label": base_name  # Label displayed below the image
            })
        else:
            # Print to logs if a matching model is missing (for debugging)
            print(f"⚠️ Matching model not found: Image {img_file} -> Missing {expected_glb_name}")

print(f"✅ Successfully loaded {len(examples)} pairs of data")

# 3. Define click event: Return the corresponding model path
def display_model(evt: gr.SelectData):
    if evt.index < len(examples):
        return examples[evt.index]["model"]
    return None

# 4. Build the interface
with gr.Blocks(title="3D Model Viewer") as demo:
    gr.Markdown("## 3D Accelerated Instances Showcase")
    gr.Markdown(f"Click on an image on the left to view the 3D effect on the right. Loaded {len(examples)} instances in total.")
    
    with gr.Row():
        # Left column: Image Gallery
        with gr.Column(scale=1):
            gallery = gr.Gallery(
                value=[item["image"] for item in examples],
                label="Instance Gallery",
                columns=3,         # Number of images per row
                height=800,        # Set height (scrollable if many images)
                object_fit="contain",
                allow_preview=False # Disable click-to-zoom
            )
        
        # Right column: 3D Viewer
        with gr.Column(scale=2):
            model_3d = gr.Model3D(
                clear_color=[0.2, 0.2, 0.2, 1.0], # Dark gray background for better contrast
                label="3D Interactive View",
                camera_position=(0, 0, 2.5)       # Adjust initial camera distance
            )

    # Bind interaction: Gallery click -> Update Model3D
    gallery.select(fn=display_model, outputs=model_3d)

# Launch the app
if __name__ == "__main__":
    # ssr_mode=False is crucial to prevent "bool is not iterable" errors in some environments
    demo.launch()