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

# Set Page Config
st.set_page_config(page_title="AI Ultra Enhancer", layout="wide")

# --- IMPROVED CSS ---
st.markdown(
    """
    <style>
        .stApp { background-color: #050505 !important; color: white !important; }
        [data-testid="stSidebar"] { background-color: #000000 !important; border-right: 1px solid #222; }
        .stDownloadButton button {
            width: 100%;
            background-color: #007bff !important;
            border-radius: 10px;
            padding: 15px;
            font-weight: bold;
        }
        #MainMenu, footer, .stActionMenu { visibility: hidden !important; }
        img { pointer-events: none; -webkit-user-drag: none; border-radius: 8px; border: 1px solid #333; }
    </style>
    """,
    unsafe_allow_html=True
)

# --- THE "FIX" LOGIC (CLAHE + SMART SHARPEN) ---

def smart_enhance(image):
    # Convert PIL to OpenCV BGR
    img_array = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR)
    
    # 1. CLAHE (Adaptive Equalization) - Fixes the 'blackish' tint issue
    lab = cv2.cvtColor(img_array, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    
    # clipLimit 2.0 keeps it natural; tileGridSize 8x8 is standard
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8,8))
    cl = clahe.apply(l)
    
    limg = cv2.merge((cl,a,b))
    enhanced_bgr = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    
    # 2. Smart Sharpen (Unsharp Mask) - Makes text HD without artifacts
    # We blur a copy and subtract it from the original to find edges
    blurred = cv2.GaussianBlur(enhanced_bgr, (0, 0), 3)
    sharpened = cv2.addWeighted(enhanced_bgr, 1.5, blurred, -0.5, 0)
    
    return Image.fromarray(cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB))

def reduce_noise(image):
    img_array = np.array(image)
    denoised = cv2.fastNlMeansDenoisingColored(img_array, None, 10, 10, 7, 21)
    return Image.fromarray(denoised)

# --- MAIN APP ---
st.title("📸 AI Ultra Enhancer")
uploaded_file = st.file_uploader("Drop an image here", type=["jpg", "png", "jpeg"])

if uploaded_file:
    original = Image.open(uploaded_file).convert("RGB")
    
    st.sidebar.header("🛠️ Magic Tools")
    
    # Use the new Smart Enhance instead of the old Auto-Fix
    do_smart = st.sidebar.toggle("Smart HD Enhancement", value=True)
    do_denoise = st.sidebar.toggle("Remove Noise (HD)", value=False)
    
    st.sidebar.divider()
    
    with st.sidebar.expander("Manual Fine-Tune"):
        br = st.slider("Brightness", 0.5, 2.0, 1.0)
        ct = st.slider("Contrast", 0.5, 2.0, 1.0)
        sh = st.slider("Manual Sharpness", 0.0, 3.0, 1.0)
        sa = st.slider("Color Pop", 0.0, 2.0, 1.0)

    preset = st.sidebar.radio("Quick Presets", ["Original", "Vibrant", "Document/Text", "B&W", "Sketch"])

    # --- PIPELINE ---
    processed = original.copy()

    # 1. AI Logic
    if do_smart:
        processed = smart_enhance(processed)
    if do_denoise:
        processed = reduce_noise(processed)

    # 2. Preset Logic
    if preset == "Vibrant":
        processed = ImageEnhance.Color(processed).enhance(1.5)
    elif preset == "Document/Text":
        # Specifically for screenshots/Allens students notes!
        processed = ImageEnhance.Contrast(processed).enhance(1.5)
        processed = ImageEnhance.Sharpness(processed).enhance(2.0)
    elif preset == "B&W":
        processed = processed.convert("L").convert("RGB")
    elif preset == "Sketch":
        gray = cv2.cvtColor(np.array(processed), cv2.COLOR_RGB2GRAY)
        inv = 255 - gray
        blur = cv2.GaussianBlur(inv, (21, 21), 0)
        sketch = cv2.divide(gray, 255 - blur, scale=256)
        processed = Image.fromarray(sketch).convert("RGB")

    # 3. Final Sliders
    processed = ImageEnhance.Brightness(processed).enhance(br)
    processed = ImageEnhance.Contrast(processed).enhance(ct)
    processed = ImageEnhance.Sharpness(processed).enhance(sh)
    processed = ImageEnhance.Color(processed).enhance(sa)

    # --- DISPLAY ---
    col1, col2 = st.columns(2)
    with col1:
        st.caption("Before")
        st.image(original, use_container_width=True)
    with col2:
        st.caption("After (Enhanced)")
        st.image(processed, use_container_width=True)

    # --- DOWNLOAD ---
    st.divider()
    buf = io.BytesIO()
    processed.save(buf, format="PNG")
    st.download_button("✨ Download Enhanced Image", data=buf.getvalue(), file_name="AI_Enhanced.png", mime="image/png")