import gradio as gr
from main import main
from arguments import parse_args
import os
import shutil
import glob

def list_iter_images(save_dir):
    # Specify the image extensions you want to search for
    image_extensions = ['jpg', 'jpeg', 'png', 'gif', 'bmp']  # Add more if needed

    # Create a list to store the image file paths
    image_paths = []

    # Iterate through the specified image extensions and get the file paths
    for ext in image_extensions:
        # Use glob to find all image files with the given extension
        image_paths.extend(glob.glob(os.path.join(save_dir, f'*.{ext}')))

    # Now image_paths contains the list of all image file paths
    print(image_paths)

    return image_paths

def clean_dir(save_dir):
    # Check if the directory exists
    if os.path.exists(save_dir):
        # Check if the directory contains any files
        if len(os.listdir(save_dir)) > 0:
            # If it contains files, delete all files in the directory
            for filename in os.listdir(save_dir):
                file_path = os.path.join(save_dir, filename)
                try:
                    if os.path.isfile(file_path) or os.path.islink(file_path):
                        os.unlink(file_path)  # Remove file or symbolic link
                    elif os.path.isdir(file_path):
                        shutil.rmtree(file_path)  # Remove directory and its contents
                except Exception as e:
                    print(f"Failed to delete {file_path}. Reason: {e}")
            print(f"All files in {save_dir} have been deleted.")
        else:
            print(f"{save_dir} exists but is empty.")
    else:
        print(f"{save_dir} does not exist.")

def generate_image(prompt, model, num_iterations, learning_rate, progress=gr.Progress(track_tqdm=True)):
    # Set up arguments
    args = parse_args()
    args.task = "single"
    args.prompt = prompt
    args.model = model
    args.n_iters = num_iterations
    args.lr = learning_rate
    args.cache_dir = "./HF_model_cache"
    args.save_dir = "./outputs"
    args.save_all_images = True

    settings = (
        f"{args.model}{'_' + args.prompt if args.task == 't2i-compbench' else ''}"
        f"{'_no-optim' if args.no_optim else ''}_{args.seed if args.task != 'geneval' else ''}"
        f"_lr{args.lr}_gc{args.grad_clip}_iter{args.n_iters}"
        f"_reg{args.reg_weight if args.enable_reg else '0'}"
        f"{'_pickscore' + str(args.pickscore_weighting) if args.enable_pickscore else ''}"
        f"{'_clip' + str(args.clip_weighting) if args.enable_clip else ''}"
        f"{'_hps' + str(args.hps_weighting) if args.enable_hps else ''}"
        f"{'_imagereward' + str(args.imagereward_weighting) if args.enable_imagereward else ''}"
        f"{'_aesthetic' + str(args.aesthetic_weighting) if args.enable_aesthetic else ''}"
    )

    save_dir = f"{args.save_dir}/{args.task}/{settings}/{args.prompt}"
    clean_dir(save_dir)
    
    try:
        # Run the main function with progress tracking
        def progress_callback(step):
            progress(step / num_iterations, f"Iteration {step}/{num_iterations}")

        best_image, total_init_rewards, total_best_rewards = main(args, progress_callback)
        
        # Return the path to the generated image
        image_path = f"{save_dir}/best_image.png"
        
        if os.path.exists(image_path):
            iter_images = list_iter_images(save_dir)
            return image_path, f"Image generated successfully and saved at {image_path}", iter_images
        else:
            return None, "Image generation completed, but the file was not found.", None
    
    except Exception as e:
        return None, f"An error occurred: {str(e)}", None

# Create Gradio interface
title="# ReNO: Enhancing One-step Text-to-Image Models through Reward-based Noise Optimization"
description="Enter a prompt to generate an image using ReNO. Adjust the model and parameters as needed."

with gr.Blocks() as demo:
    with gr.Column():
        gr.Markdown(title)
        gr.Markdown(description)
        gr.HTML("""
        <div style="display:flex;column-gap:4px;">
            <a href='https://github.com/ExplainableML/ReNO'>
                <img src='https://img.shields.io/badge/GitHub-Repo-blue'>
            </a> 
            <a href='https://arxiv.org/abs/2406.04312v1'>
                <img src='https://img.shields.io/badge/Paper-Arxiv-red'>
            </a>
        </div>
        """)

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(label="Prompt")
                chosen_model = gr.Dropdown(["sd-turbo", "sdxl-turbo", "pixart", "hyper-sd"], label="Model", value="sd-turbo")
                
                with gr.Row():
                    n_iter = gr.Slider(minimum=10, maximum=100, step=10, value=50, label="Number of Iterations")
                    learning_rate = gr.Slider(minimum=0.1, maximum=10.0, step=0.1, value=5.0, label="Learning Rate")

                submit_btn = gr.Button("Submit")

                gr.Examples(
                    examples = [
                        "A minimalist logo design of a reindeer, fully rendered. The reindeer features distinct, complete shapes using bold and flat colors. The design emphasizes simplicity and clarity, suitable for logo use with a sharp outline and white background.",
                        "A blue scooter is parked near a curb in front of a green vintage car",
                        "A impressionistic oil painting: a lone figure walking on a misty beach, a weathered lighthouse on a cliff, seagulls above crashing waves",
                        "A bird with 8 legs",
                        "An orange chair to the right of a black airplane",
                        "A pink elephant and a grey cow",
                    ],
                    inputs = [prompt]     
                )
            
            with gr.Column():
                output_image = gr.Image(type="filepath", label="Generated Image")
                status = gr.Textbox(label="Status")
                iter_gallery = gr.Gallery(label="Iterations", columns=4)

    submit_btn.click(
        fn = generate_image,
        inputs = [prompt, chosen_model, n_iter, learning_rate],
        outputs = [output_image, status, iter_gallery]
    )

# Launch the app
demo.queue().launch(show_error=True, show_api=False)