import streamlit as st
from PIL import Image, ImageEnhance, ImageFilter, ImageDraw, ImageFont
import io
import numpy as np
import cv2
from rembg import remove  # For background removal

# Title and Description
st.title("Advanced Image Enhancer and Editor")
st.write("Enhance, edit, and apply AI-powered enhancements to your images!")

# File Upload
uploaded_file = st.file_uploader("Upload an Image", type=["jpg", "jpeg", "png"])

if uploaded_file:
    # Load image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_column_width=True)

    # Sidebar Tools
    st.sidebar.title("Editing Tools")

    # Basic Adjustments
    brightness = st.sidebar.slider("Brightness", 0.1, 3.0, 1.0)
    contrast = st.sidebar.slider("Contrast", 0.1, 3.0, 1.0)
    sharpness = st.sidebar.slider("Sharpness", 0.1, 3.0, 1.0)
    color = st.sidebar.slider("Color Saturation", 0.1, 3.0, 1.0)

    # Debugging print for color value
    st.write(f"Color Value: {color}")

    # Filters
    if st.sidebar.checkbox("Apply Filters"):
        filter_option = st.sidebar.selectbox("Select a Filter", ["None", "BLUR", "DETAIL", "EDGE_ENHANCE", "SMOOTH"])
        if filter_option == "BLUR":
            image = image.filter(ImageFilter.BLUR)
        elif filter_option == "DETAIL":
            image = image.filter(ImageFilter.DETAIL)
        elif filter_option == "EDGE_ENHANCE":
            image = image.filter(ImageFilter.EDGE_ENHANCE)
        elif filter_option == "SMOOTH":
            image = image.filter(ImageFilter.SMOOTH)

    # AI Filters
    if st.sidebar.checkbox("Apply AI Filters"):
        ai_filter = st.sidebar.selectbox("Select AI Filter", ["Cartoonize", "Sketch"])
        image_array = np.array(image)

        if ai_filter == "Cartoonize":
            gray = cv2.cvtColor(image_array, cv2.COLOR_BGR2GRAY)
            edges = cv2.adaptiveThreshold(gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 9, 9)
            color = cv2.bilateralFilter(image_array, 9, 250, 250)
            cartoon = cv2.bitwise_and(color, color, mask=edges)
            image = Image.fromarray(cartoon)

        elif ai_filter == "Sketch":
            gray = cv2.cvtColor(image_array, cv2.COLOR_BGR2GRAY)
            inv = 255 - gray
            blur = cv2.GaussianBlur(inv, (21, 21), 0)
            sketch = cv2.divide(gray, 255 - blur, scale=256)
            image = Image.fromarray(sketch)

    # Background Removal
    if st.sidebar.checkbox("Remove Background"):
        image = Image.open(io.BytesIO(remove(uploaded_file.read())))

    # Watermarking
    if st.sidebar.checkbox("Add Watermark"):
        watermark_text = st.sidebar.text_input("Watermark Text", "My Watermark")
        draw = ImageDraw.Draw(image)
        font = ImageFont.load_default()  # Use default font if arial.ttf doesn't work
        width, height = image.size
        text_width, text_height = draw.textsize(watermark_text, font)
        x, y = width - text_width - 10, height - text_height - 10
        draw.text((x, y), watermark_text, font=font, fill=(255, 255, 255, 128))

    # Enhance Image
    try:
        enhancer = ImageEnhance.Brightness(image)
        image = enhancer.enhance(brightness)
        enhancer = ImageEnhance.Contrast(image)
        image = enhancer.enhance(contrast)
        enhancer = ImageEnhance.Sharpness(image)
        image = enhancer.enhance(sharpness)
        enhancer = ImageEnhance.Color(image)
        image = enhancer.enhance(color)  # Enhance color saturation

    except Exception as e:
        st.write(f"Error during enhancement: {e}")

    # Display Enhanced Image
    st.image(image, caption="Enhanced Image", use_column_width=True)

    # Download Button
    buf = io.BytesIO()
    image.save(buf, format="PNG")
    byte_im = buf.getvalue()
    st.download_button("Download Enhanced Image", data=byte_im, file_name="enhanced_image.png", mime="image/png")