import streamlit as st
from utils import transcribe_audio_segments, create_video_segments, generate_speech
from PIL import Image
import os
import cv2
import numpy as np
from io import BytesIO
import tempfile
from moviepy.editor import ImageSequenceClip
from rembg import remove
from basicsr.archs.rrdbnet_arch import RRDBNet
from realesrgan import RealESRGANer

# Helper function to apply cartoon effect
def cartoonize(img):
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    gray = cv2.medianBlur(gray, 5)
    edges = cv2.adaptiveThreshold(
        gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C,
        cv2.THRESH_BINARY, 9, 9
    )
    color = cv2.bilateralFilter(img, 9, 300, 300)
    cartoon = cv2.bitwise_and(color, color, mask=edges)
    return cartoon

# Helper function to apply edge detection
def edge_detect(img):
    return cv2.Canny(img, 100, 200)

# Helper function to enhance face (upscale image)
def enhance_image(img):
    model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64,
                   num_block=23, num_grow_ch=32, scale=4)
    upsampler = RealESRGANer(
        scale=4, model_path='realesrgan-x4plus.pth', model=model,
        tile=0, tile_pad=10, pre_pad=0, half=False
    )
    _, _, output = upsampler.enhance(np.array(img), outscale=4)
    return output

# Save processed images
def save_image(img, filename):
    cv2.imwrite(filename, img)

# Convert images to slideshow video
def create_slideshow(images, out_path="slideshow.mp4"):
    clip = ImageSequenceClip(images, fps=1)
    clip.write_videofile(out_path, codec='libx264')

# Streamlit UI
st.set_page_config(layout="wide")
st.title("🎨 Interactive Image Effects Studio")

uploaded_files = st.file_uploader("Upload images", type=["jpg", "jpeg", "png"], accept_multiple_files=True)
effects = st.multiselect("Select effects to apply", ["Cartoonize", "Edge Detection", "Background Removal", "Enhance/ Upscale"])

effect_images = []

if uploaded_files:
    for uploaded_file in uploaded_files:
        image = Image.open(uploaded_file).convert("RGB")
        img_array = np.array(image)
        original_img = cv2.cvtColor(img_array, cv2.COLOR_RGB2BGR)

        processed_img = original_img.copy()
        if "Cartoonize" in effects:
            processed_img = cartoonize(processed_img)
        if "Edge Detection" in effects:
            processed_img = cv2.cvtColor(edge_detect(processed_img), cv2.COLOR_GRAY2BGR)
        if "Background Removal" in effects:
            processed_img = remove(Image.fromarray(cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB)))
            processed_img = cv2.cvtColor(np.array(processed_img), cv2.COLOR_RGB2BGR)
        if "Enhance/ Upscale" in effects:
            processed_img = enhance_image(processed_img)

        st.image(cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB), caption="Processed Image", use_column_width=True)

        # Save processed
        effect_images.append(cv2.cvtColor(processed_img, cv2.COLOR_BGR2RGB))
        with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as tmp_file:
            save_image(processed_img, tmp_file.name)
            with open(tmp_file.name, "rb") as file:
                btn = st.download_button(
                    label="Download Processed Image",
                    data=file,
                    file_name=os.path.basename(uploaded_file.name),
                    mime="image/jpeg"
                )

    if st.button("Create Slideshow from Images"):
        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as tmp_video:
            create_slideshow(effect_images, tmp_video.name)
            with open(tmp_video.name, "rb") as file:
                st.video(file)
                st.download_button("Download Slideshow Video", file, file_name="slideshow.mp4")