import gradio as gr
import os
from pathlib import Path

def process_selection(image_choice, video_choice):
    if image_choice and video_choice:
        img_number = Path(image_choice).name.split(".png")[0]
        output_video = os.path.join(
            "./data/Generated_Results", 
            img_number,
            f"{video_choice}.mp4"
        )
        
        if os.path.exists(output_video):
            return (f"Selected Image (Imitator): {image_choice}\nSelected Video (Actor): {video_choice}\nTo avoid GPU outrage, we cache the generated videos in advance.\nOutput: Generated video loaded from cache.", output_video)
        else:
            return f"Error: Generated video not found at {output_video}", None
            
    return "Please select both an image and a video.", None

# Define the reference images and videos
image_folder = "./data/Reference_Images"
video_folder = "./data/Source_Videos"

def get_sorted_files(folder, extension):
    files = [f for f in os.listdir(folder) if f.endswith(extension)]
    files.sort()
    return [os.path.join(folder, f) for f in files]

# Sort the files
image_paths = get_sorted_files(image_folder, ".png")
video_paths = get_sorted_files(video_folder, ".mp4")

def create_gallery_items(paths, prefix):
    return [(path, os.path.splitext(os.path.basename(path))[0]) for path in paths]

def update_gallery(choice, paths, prefix):
    items = []
    for path in paths:
        name = os.path.splitext(os.path.basename(path))[0]
        if choice and name == choice:
            # Highlight Selected item
            items.insert(0, (path, f"✓ {name} (Selected)"))
        else:
            items.append((path, name))
    return items

with gr.Blocks() as app:
    gr.Markdown("## 3DHM Results🤸‍♂️")
    gr.Markdown("We showcase more results of 3DHM. The interactive demo will be available soon - stay tuned!")
    # Gallery
    gr.Markdown("---")
    gr.Markdown("## Select input image and motion.")

    with gr.Row():
        with gr.Column():
            gr.Markdown("### Reference Images")
            gallery_images = gr.Gallery(
                value=create_gallery_items(image_paths, "Image"),
                label="Preview Images",
                show_label=True,
                columns=4,
                height=400,
                interactive=False,
                object_fit="contain",
                preview=True
            )

        with gr.Column():
            gr.Markdown("### Sources Videos")
            gallery_videos = gr.Gallery(
                value=create_gallery_items(video_paths, "Video"),
                label="Preview Videos",
                show_label=True,
                columns=2,
                height=400,
                interactive=False,
                object_fit="contain",
                preview=True
            )

    
    with gr.Row():
        # Output Panel
        with gr.Column(scale=2):
            output_video = gr.Video(
                label="Generated Video", 
                interactive=False,
                height=500,
                width=700
            )
        
        # Select input Panel
        with gr.Column(scale=1):
            output = gr.Textbox(
                label="Generation Status",
                lines=4
            )
            gr.Markdown("---")
            # Select image and video
            image_choice = gr.Dropdown(
                choices=[os.path.splitext(os.path.basename(p))[0] for p in image_paths],
                label="Select Reference Image",
                type="value"
            )
            video_choice = gr.Dropdown(
                choices=[os.path.splitext(os.path.basename(p))[0] for p in video_paths],
                label="Select Source Video",
                type="value"
            )
            current_selection = gr.Textbox(
                label="Current Selection",
                lines=2
            )
            submit_button = gr.Button("Generate Video", variant="primary", size="lg")



    def update_selection(img, vid):
        if img and vid:
            return f"Input Image: {img}\nSource motion:{vid}"
        elif img:
            return f"Selected image: {img}, please select a video"
        elif vid:
            return f"Selected video: {vid}, please select an image"
        return "Please select both an image and a video"

    def update_image_gallery(img_choice):
        return update_gallery(img_choice, image_paths, "Image")
    
    def update_video_gallery(vid_choice):
        return update_gallery(vid_choice, video_paths, "Video")

    image_choice.change(
        update_image_gallery,
        inputs=[image_choice],
        outputs=gallery_images
    )
    video_choice.change(
        update_video_gallery,
        inputs=[video_choice],
        outputs=gallery_videos
    )

    # Update current selection when either choice changes
    image_choice.change(
        update_selection,
        inputs=[image_choice, video_choice],
        outputs=current_selection
    )
    video_choice.change(
        update_selection,
        inputs=[image_choice, video_choice],
        outputs=current_selection
    )

    submit_button.click(
        process_selection,
        inputs=[image_choice, video_choice],
        outputs=[output, output_video]
    )

if __name__ == "__main__":
    app.launch(inline=True)