import streamlit as st
from PIL import Image, ImageEnhance
import io

def apply_filters(image, grayscale, brightness, contrast):
    if grayscale:
        image = image.convert("L")
    
    enhancer = ImageEnhance.Brightness(image)
    image = enhancer.enhance(brightness)

    enhancer = ImageEnhance.Contrast(image)
    image = enhancer.enhance(contrast)

    return image

# Streamlit app setup
st.title("Image Editor")

# Image upload
uploaded_file = st.file_uploader("Upload an image", type=["jpg", "jpeg", "png"])
if uploaded_file is not None:
    # Load the image
    image = Image.open(uploaded_file)

    # Initialize session state for restoring the original image and filters
    if 'original_image' not in st.session_state:
        st.session_state.original_image = image.copy()
    if 'grayscale' not in st.session_state:
        st.session_state.grayscale = False
    if 'brightness' not in st.session_state:
        st.session_state.brightness = 1.0
    if 'contrast' not in st.session_state:
        st.session_state.contrast = 1.0

    # Layout with three columns: Filter Options, Edited Image, Original Image
    col1, col2, col3 = st.columns([1, 2, 2])

    with col1:
        st.sidebar.header("Filters")

        # Restore button
        if st.button("Restore to Original"):
            st.session_state.grayscale = False
            st.session_state.brightness = 1.0
            st.session_state.contrast = 1.0
            image = st.session_state.original_image.copy()

        # Filters
        grayscale = st.checkbox("Grayscale", value=st.session_state.grayscale, key="grayscale")
        brightness = st.slider("Brightness", 0.1, 2.0, st.session_state.brightness, key="brightness")
        contrast = st.slider("Contrast", 0.1, 2.0, st.session_state.contrast, key="contrast")

        # Update session state with current filter values (before widgets)
        if grayscale != st.session_state.grayscale:
            st.session_state.grayscale = grayscale
        if brightness != st.session_state.brightness:
            st.session_state.brightness = brightness
        if contrast != st.session_state.contrast:
            st.session_state.contrast = contrast

    with col2:
        st.subheader("Edited Image")

        # Apply filters
        edited_image = apply_filters(image, st.session_state.grayscale, st.session_state.brightness, st.session_state.contrast)

        st.image(edited_image, use_container_width=True)

        # Download button
        buf = io.BytesIO()
        edited_image.save(buf, format="PNG")
        byte_im = buf.getvalue()

        st.download_button(
            label="Download Edited Image",
            data=byte_im,
            file_name="edited_image.png",
            mime="image/png"
        )

    with col3:
        st.subheader("Original Image")
        st.image(image, use_container_width=True)  # Show the original image