import gradio as gr
from PIL import Image
import os
from datetime import datetime
from generate_image import generate_and_wait_for_result
import concurrent.futures

def create_markdown_with_images(prompt, image_paths, batch_folder, reference_image_paths=None):
    """
    Create a markdown file with the prompt, reference images, and generated image links.
    """
    markdown_content = f"# Image Generation Results\n\n"
    markdown_content += f"**Prompt:** {prompt}\n\n"
    markdown_content += f"**Generated on:** {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n"
    markdown_content += f"**Batch:** {batch_folder}\n\n"
    
    if reference_image_paths:
        markdown_content += "## Reference Images\n\n"
        markdown_content += "| Character | Scene | Style |\n"
        markdown_content += "|-----------|-------|-------|\n"
        ref_cells = []
        ref_labels = ["Character", "Scene", "Style"]
        for i, ref_path in enumerate(reference_image_paths):
            if ref_path and os.path.exists(ref_path):
                ref_filename = os.path.basename(ref_path)
                import shutil
                ref_dest = os.path.join("output", batch_folder, f"ref_{ref_filename}")
                shutil.copy2(ref_path, ref_dest)
                ref_cells.append(f"![{ref_labels[i] if i < len(ref_labels) else 'Reference'}](ref_{ref_filename})<br>*{ref_filename}*")
            else:
                ref_cells.append("*No image*")
        while len(ref_cells) < 3:
            ref_cells.append("*No image*")
        markdown_content += f"| {ref_cells[0]} | {ref_cells[1]} | {ref_cells[2]} |\n\n"
    
    markdown_content += "## Generated Images\n\n"
    for i, image_path in enumerate(image_paths, 1):
        if image_path and os.path.exists(image_path):
            filename = os.path.basename(image_path)
            markdown_content += f"- **Image {i}:** `![Generated Image {i}]({filename})`\n"
        else:
            markdown_content += f"- **Image {i}:** *Generation failed*\n"
    
    output_dir = os.path.join("output", batch_folder)
    os.makedirs(output_dir, exist_ok=True)
    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    markdown_filename = f"generation_report_{timestamp}.md"
    markdown_path = os.path.join(output_dir, markdown_filename)
    
    with open(markdown_path, 'w', encoding='utf-8') as f:
        f.write(markdown_content)
    
    return markdown_path

def save_uploaded_image(image, folder_name):
    """
    Save uploaded image to the appropriate assets folder.
    """
    if image is None:
        return image
    assets_path = os.path.join("assets", folder_name)
    os.makedirs(assets_path, exist_ok=True)
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    filename = f"uploaded_{timestamp}.jpg"
    filepath = os.path.join(assets_path, filename)
    image.save(filepath, "JPEG", quality=95)
    print(f"Saved uploaded image to: {filepath}")
    return image

def process_images(prompt, character_image, scene_image, style_image, num_images):
    """
    Process text prompt and input images, generate specified number of images using RunwayML.
    """
    num_images = int(num_images)
    try:
        processed_prompt = prompt
        
        reference_images = []
        temp_paths = []
        
        if character_image:
            char_path = os.path.join("assets", "characters", f"temp_char_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg")
            os.makedirs(os.path.dirname(char_path), exist_ok=True)
            character_image.save(char_path, "JPEG", quality=95)
            reference_images.append(char_path)
            temp_paths.append(char_path)
        
        if scene_image:
            scene_path = os.path.join("assets", "scenes", f"temp_scene_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg")
            os.makedirs(os.path.dirname(scene_path), exist_ok=True)
            scene_image.save(scene_path, "JPEG", quality=95)
            reference_images.append(scene_path)
            temp_paths.append(scene_path)
        
        if style_image:
            style_path = os.path.join("assets", "styles", f"temp_style_{datetime.now().strftime('%Y%m%d_%H%M%S')}.jpg")
            os.makedirs(os.path.dirname(style_path), exist_ok=True)
            style_image.save(style_path, "JPEG", quality=95)
            reference_images.append(style_path)
            temp_paths.append(style_path)
        
        if not reference_images:
            return "No reference images provided.", []
        
        timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
        batch_folder = f"batch_{timestamp}"
        batch_folder_path = os.path.join("output", batch_folder)
        os.makedirs(batch_folder_path, exist_ok=True)
        
        image_paths = [os.path.join(batch_folder_path, f"generated_{i+1}_{timestamp}.jpg") for i in range(num_images)]
        
        def generate_single_image(index):
            filename = os.path.basename(image_paths[index])
            task_id, generated_image_path = generate_and_wait_for_result(
                prompt_text=processed_prompt,
                reference_image_paths=reference_images,
                auto_tag_prompt=False,
                filename=filename,
                batch_folder=batch_folder
            )
            return index, task_id, generated_image_path
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=num_images) as executor:
            futures = [executor.submit(generate_single_image, i) for i in range(num_images)]
            completed_image_paths = [None] * num_images
            for future in concurrent.futures.as_completed(futures):
                index, task_id, generated_image_path = future.result()
                completed_image_paths[index] = generated_image_path
        
        markdown_path = create_markdown_with_images(
            prompt=prompt,
            image_paths=completed_image_paths,
            batch_folder=batch_folder,
            reference_image_paths=reference_images
        )
        
        for temp_path in temp_paths:
            if os.path.exists(temp_path):
                os.remove(temp_path)
        
        existing_image_paths = [path for path in completed_image_paths if path and os.path.exists(path)]
        status_msg = f"{processed_prompt}\n\nGenerated {len(existing_image_paths)}/{num_images} images in batch: {batch_folder}\nMarkdown report: {markdown_path}"
        
        return status_msg, existing_image_paths
    
    except Exception as e:
        return f"Error: {str(e)}", []

# Gradio Interface
with gr.Blocks(title="RunwayML Image Generation") as demo:
    gr.Markdown("# RunwayML Image Generation")
    gr.Markdown("""Generate images using RunwayML with reference images. Choose how many images to generate (1-4). 
    Duplicate the Space to your own Hugging Face account (via the "Duplicate Space" option). 
    Make sure to set your own RunwayAPI key as RUNWAYML_API_SECRET.""")
    
    with gr.Tab("RunwayML with References"):
        with gr.Row():
            num_images = gr.Dropdown(
                choices=["1", "2", "3", "4"],
                value="4",
                label="Number of Images to Generate"
            )
        runway_prompt_input = gr.Textbox(
            label="Text Prompt",
            placeholder="Enter your prompt here...",
            value="@character is in @scene with @style style"
        )
        
        with gr.Row():
            character_input = gr.Image(label="Character", type="pil")
            scene_input = gr.Image(label="Scene", type="pil")
            style_input = gr.Image(label="Style", type="pil")
        
        runway_process_btn = gr.Button("Generate Images", variant="primary")
        
        runway_output_text = gr.Textbox(label="Status", interactive=False)
        runway_gallery = gr.Gallery(
            label="Generated Images",
            show_label=True,
            elem_id="gallery",
            columns=2,
            rows=2
        )
    
    # Event handlers
    character_input.upload(
        fn=lambda img: save_uploaded_image(img, "characters"),
        inputs=[character_input],
        outputs=[character_input]
    )
    scene_input.upload(
        fn=lambda img: save_uploaded_image(img, "scenes"),
        inputs=[scene_input],
        outputs=[scene_input]
    )
    style_input.upload(
        fn=lambda img: save_uploaded_image(img, "styles"),
        inputs=[style_input],
        outputs=[style_input]
    )
    
    runway_process_btn.click(
        fn=process_images,
        inputs=[runway_prompt_input, character_input, scene_input, style_input, num_images],
        outputs=[runway_output_text, runway_gallery]
    )

if __name__ == "__main__":
    demo.queue(default_concurrency_limit=2, max_size=20)
    demo.launch()