import gradio as gr
from yellow_tint_cleaner import auto_adjust
from PIL import Image
import numpy as np

def process_image(
    image,
    strength,
    brightness,
    contrast,
    saturation,
    red,
    green,
    blue,
    mode
):
    """Process the image with the given parameters."""
    if image is None:
        return None
    
    # Convert the image to PIL Image if it's not already
    if not isinstance(image, Image.Image):
        image = Image.fromarray(image)
    
    # Process the image
    result = auto_adjust(
        image,
        strength=strength,
        brightness=brightness,
        contrast=contrast,
        saturation=saturation,
        red=red,
        green=green,
        blue=blue,
        mode=mode
    )
    
    return result

# Create the Gradio interface
with gr.Blocks(title="Yellow Tint Cleaner") as demo:
    gr.Markdown("""
    # GPT-4o Yellow Tint Cleaner
    This app helps you clean yellow tints from your images using various adjustment parameters.
    Code adapted from [ComfyUI_LayerStyle](https://github.com/chflame163/ComfyUI_LayerStyle) by chflame163 
    """)
    
    with gr.Row():
        with gr.Column():
            input_image = gr.Image(label="Input Image", type="pil", height=600)
            
            with gr.Accordion("Adjustments", open=False):
                with gr.Row():
                    strength = gr.Slider(
                        minimum=0,
                        maximum=100,
                        value=75,
                        step=1,
                        label="Strength"
                    )
                    brightness = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Brightness"
                    )
                
                with gr.Row():
                    contrast = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Contrast"
                    )
                    saturation = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Saturation"
                    )
                
                with gr.Row():
                    red = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Red Channel"
                    )
                    green = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Green Channel"
                    )
                
                with gr.Row():
                    blue = gr.Slider(
                        minimum=-100,
                        maximum=100,
                        value=0,
                        step=1,
                        label="Blue Channel"
                    )
                    mode = gr.Dropdown(
                        choices=["RGB", "lum + sat", "luminance", "saturation", "mono"],
                        value="lum + sat",
                        label="Processing Mode"
                    )
            
            process_btn = gr.Button("Process Image")
        
        with gr.Column():
            output_image = gr.Image(label="Output Image", height=600)
    
    process_btn.click(
        fn=process_image,
        inputs=[
            input_image,
            strength,
            brightness,
            contrast,
            saturation,
            red,
            green,
            blue,
            mode
        ],
        outputs=output_image
    )

if __name__ == "__main__":
    demo.launch()