import io
import cv2
import numpy as np
from PIL import Image
from filters import *
import streamlit as st


# Generating a button to download the image file.
def download_image_button(img, filename, text):
    buffered = io.BytesIO()
    img.save(buffered, format="JPEG")
    img_bytes = buffered.getvalue()

    # Using st.download_button to handle the download
    st.download_button(label=text, data=img_bytes, file_name=filename, mime="image/jpeg", use_container_width=True)


# Set title.
st.title("Artistic Image Filters")

# Upload image.
uploaded_file = st.file_uploader("Choose an image file:", type=["png", "jpg"])

if uploaded_file is not None:
    # Convert the file to an opencv image.
    raw_bytes = np.asarray(bytearray(uploaded_file.read()), dtype=np.uint8)
    img = cv2.imdecode(raw_bytes, cv2.IMREAD_COLOR)
    input_col, output_col = st.columns(2)
    with input_col:
        st.header("Original")
        # Display uploaded image.
        st.image(img, channels="BGR", use_column_width=True)

    st.header("Filter Examples:")

    # Display radio buttons for choosing the filter to apply.
    option = st.radio(
        "Select a filter:",
        (
            "None",
            "Black and White",
            "Sepia / Vintage",
            "Vignette Effect",
            "Pencil Sketch",
        ),
        horizontal=True,
    )

    # Define columns for thumbnail images.
    col1, col2, col3, col4 = st.columns(4)
    with col1:
        st.caption("Black and White")
        st.image("./filters/filter_bw.jpg")
    with col2:
        st.caption("Sepia / Vintage")
        st.image("./filters/filter_sepia.jpg")
    with col3:
        st.caption("Vignette Effect")
        st.image("./filters/filter_vignette.jpg")
    with col4:
        st.caption("Pencil Sketch")
        st.image("./filters/filter_pencil_sketch.jpg")

    # Flag for showing output image.
    output_flag = 1
    # Colorspace of output image.
    color = "BGR"

    # Generate filtered image based on the selected option.
    if option == "None":
        # Don't show output image.
        output_flag = 0
    elif option == "Black and White":
        output = bw_filter(img)
        color = "GRAY"
    elif option == "Sepia / Vintage":
        output = sepia(img)
    elif option == "Vignette Effect":
        level = st.slider("level", 0, 5, 2)
        output = vignette(img, level)
    elif option == "Pencil Sketch":
        ksize = st.slider("Blur kernel size", 1, 11, 5, step=2)
        output = pencil_sketch(img, ksize)
        color = "GRAY"

    with output_col:
        if output_flag == 1:
            st.header("Output")
            st.image(output, channels=color)
            # fromarray converts cv2 image into PIL format for saving it using download button.
            if color == "BGR":
                result = Image.fromarray(output[:, :, ::-1])
            else:
                result = Image.fromarray(output)
                
            # Display the download button with the text "Download Output"
            download_image_button(result, "output.jpg", "Download Output")
        else:
            st.header("Output")
            st.image(img, channels=color)