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

st.set_page_config(page_title="AI Image Enhancer", page_icon="🖼️")

st.title("🖼️ Image Quality Enhancer (HD + Deblur)")

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

def enhance_image(image, scale=2):
    img = np.array(image)

    # Convert RGB → BGR (OpenCV fix)
    img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

    # Upscale
    width = int(img.shape[1] * scale)
    height = int(img.shape[0] * scale)
    resized = cv2.resize(img, (width, height), interpolation=cv2.INTER_CUBIC)

    # Denoise
    denoised = cv2.fastNlMeansDenoisingColored(resized, None, 10, 10, 7, 21)

    # Sharpen
    kernel = np.array([[0, -1, 0],
                       [-1, 5, -1],
                       [0, -1, 0]])
    sharpened = cv2.filter2D(denoised, -1, kernel)

    # Convert back RGB
    final = cv2.cvtColor(sharpened, cv2.COLOR_BGR2RGB)

    return final

if uploaded_file:
    image = Image.open(uploaded_file)

    st.subheader("Original Image")
    st.image(image, use_container_width=True)

    scale = st.slider("Upscale Level", 1, 4, 2)

    if st.button("Enhance Image"):
        enhanced = enhance_image(image, scale)

        st.subheader("Enhanced Image")
        st.image(enhanced, use_container_width=True)

        # FIXED DOWNLOAD
        result = Image.fromarray(enhanced)

        buf = io.BytesIO()
        result.save(buf, format="PNG")
        byte_im = buf.getvalue()

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

st.markdown("---")
st.caption("Built with ❤️ using Streamlit")