import streamlit as st 
from PIL import Image, ImageEnhance, ImageFilter
import io
import numpy as np
import cv2

# Custom CSS for black background, non-draggable images, and removed three-dot menu
st.markdown(
    """
    <style>
        body, .stApp {
            background-color: black !important;
            color: white !important;
        }
        .stSidebar, .css-1d391kg, .css-1lcbmhc, .css-1l02zno {
            background-color: black !important; 
            color: white !important;
        }
        h1, h2, h3, h4, h5, h6, p, label {
            color: white !important;
        }
        .stButton>button {
            background-color: #4CAF50;
            color: white;
            border: none;
            padding: 10px 20px;
            text-align: center;
            text-decoration: none;
            display: inline-block;
            font-size: 14px;
            margin: 4px 2px;
            transition-duration: 0.4s;
            cursor: pointer;
        }
        .stButton>button:hover {
            background-color: #45a049;
        }
        .stDownloadButton>button {
            background-color: #008CBA;
            color: white;
        }
        .stDownloadButton>button:hover {
            background-color: #007bb5;
        }
        /* Make images non-draggable */
        img {
            pointer-events: none;
            -webkit-user-drag: none;
        }
        /* Remove three-dot menu */
        .stActionMenu, [data-testid="stActionButton"] {
            display: none !important;
        }
        /* Remove Streamlit footer */
        footer {visibility: hidden;}
        /* Remove Streamlit hamburger menu */
        #MainMenu {visibility: hidden;}
    </style>
    """,
    unsafe_allow_html=True
)

# Title and Description
st.title("Enhancer")

# 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_container_width=True)

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

    # Enhanced default values
    brightness = st.sidebar.slider("Brightness", 0.1, 3.0, 1.2)
    contrast = st.sidebar.slider("Contrast", 0.1, 3.0, 1.1)
    sharpness = st.sidebar.slider("Sharpness", 0.1, 3.0, 1.3)
    color = st.sidebar.slider("Color Saturation", 0.1, 3.0, 1.1)

    # Basic 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)

    # 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)

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

    # Display Enhanced Image
    st.image(image, caption="Enhanced Image", use_container_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"
    )