import gradio as gr
from PIL import Image
from film_simulation import process_images, load_default_profiles, load_film_profiles_from_json
from io import BytesIO
import tempfile
import os

# Load default profiles to display in the dropdown
default_profiles = load_default_profiles()
profile_names = list(default_profiles.keys())

def get_number_of_profiles_and_labels(profiles_json):
    if profiles_json:
        profiles = load_film_profiles_from_json(profiles_json.name)
    else:
        profiles = default_profiles
    return len(profiles), list(profiles.keys())

def resize_image(image, longer_side):
    # Calculate the resize dimensions while maintaining the aspect ratio
    width, height = image.size
    if width > height:
        new_width = longer_side
        new_height = int((longer_side / width) * height)
    else:
        new_height = longer_side
        new_width = int((longer_side / height) * width)
    
    return image.resize((new_width, new_height), Image.Resampling.LANCZOS)

def gradio_interface(image, profiles_json, selected_profile, chroma_override_flag, chroma_override_value, blur_override_flag, blur_override_value, color_temp_flag, color_temp_value, cross_process_flag, curve_type, resize_flag, resize_value):
    if image is None:
        return ["Error: No image provided. Please upload an image."]
    
    image = Image.open(image)
    
    # Resize the image if the resize option is enabled
    if resize_flag and resize_value:
        image = resize_image(image, resize_value)
    
    # Override parameters based on checkbox state
    chroma_override = chroma_override_value if chroma_override_flag else None
    blur_override = blur_override_value if blur_override_flag else None
    color_temp = color_temp_value if color_temp_flag else None
    
    profiles_json_path = profiles_json.name if profiles_json else None
    processed_images = process_images(image, profiles_json_path, selected_profile, chroma_override, blur_override, color_temp, cross_process_flag, curve_type)
    
    output_images = []
    for img in processed_images:
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmpfile:
            img = img.convert("RGB")  # Convert to RGB if the image has transparency
            img.save(tmpfile.name, format="JPEG", quality=90)  # Save as JPG with quality setting 90
            output_images.append(tmpfile.name)
    return output_images

def update_outputs(profiles_json):
    num_profiles, labels = get_number_of_profiles_and_labels(profiles_json)
    return {f"output_{i}": gr.Image(type="filepath", label=label) for i, label in enumerate(labels)}

with gr.Blocks() as iface:
    image_input = gr.Image(type="filepath", label="Input Image")
    profiles_json_input = gr.File(label="Profiles JSON")
    selected_profile = gr.Dropdown(choices=["All"] + profile_names, value="All", label="Select Profile")
    
    chroma_override_flag = gr.Checkbox(label="Override Chromatic Aberration")
    chroma_override_value = gr.Slider(0, 10, step=0.1, value=0, label="Chromatic Aberration Value")
    
    blur_override_flag = gr.Checkbox(label="Override Blur")
    blur_override_value = gr.Slider(0, 10, step=0.1, value=0, label="Blur Value")
    
    color_temp_flag = gr.Checkbox(True, label="Override Color Temperature")
    color_temp_value = gr.Slider(1000, 10000, step=100, value=6800, label="Color Temperature (K)")
    
    cross_process_flag = gr.Checkbox(label="Cross Process")
    curve_type = gr.Dropdown(choices=["color", "advanced", "both", "auto"], value="auto", label="Curve Type")
    
    resize_flag = gr.Checkbox(label="Resize Image (Longer Side)")
    resize_value = gr.Slider(256, 4096, step=16, value=2048, label="Resize Value")
    
    output_images = gr.Gallery(label="Processed Images")
    process_button = gr.Button("Process Image")

    # Update the number of outputs based on the loaded profiles
    profiles_json_input.change(fn=update_outputs, inputs=profiles_json_input, outputs=output_images)
    
    # Trigger image processing when the button is clicked
    process_button.click(fn=gradio_interface, 
                         inputs=[image_input, profiles_json_input, selected_profile, chroma_override_flag, chroma_override_value, blur_override_flag, blur_override_value, color_temp_flag, color_temp_value, cross_process_flag, curve_type, resize_flag, resize_value], 
                         outputs=output_images)

iface.launch(share=True)