import streamlit as st
import torch
from deepface import DeepFace
import cv2
import numpy as np
import random
import os
from moviepy.editor import *
from insightface.app import FaceAnalysis
import tempfile
from TTS.api import TTS

# Set Streamlit Page Config
st.set_page_config(page_title="DeepFlick: Live Deepfake & AI Voice Clone", layout="wide")

st.title("🎭 DeepFlick: Live Deepfake & AI Voice Clone Streaming")

# Load Face Swap Model
face_app = FaceAnalysis(name="buffalo_l")
face_app.prepare(ctx_id=0, det_size=(640, 640))

# Load AI Voice Model
tts = TTS("tts_models/en/ljspeech/glow-tts").to("cpu")

# Function to Swap Faces
def swap_faces(source_img_path, target_img_path, output_path):
    source_img = cv2.imread(source_img_path)
    target_img = cv2.imread(target_img_path)

    if source_img is None or target_img is None:
        return None

    swapped_img = DeepFace.swap(source_img, target_img, model_name="insightface")
    cv2.imwrite(output_path, swapped_img)
    
    return output_path

# Function to Generate AI Voice
def generate_voice(script_text):
    audio_path = f"output_{random.randint(1000,9999)}.wav"
    tts.tts_to_file(text=script_text, file_path=audio_path)
    return audio_path

# Sidebar for Uploads & Settings
st.sidebar.subheader("🎭 Upload Face for Deepfake")
source_img = st.sidebar.file_uploader("Upload Your Face (Source)", type=["jpg", "png"])
target_img = st.sidebar.file_uploader("Upload Target Face (Target)", type=["jpg", "png"])
st.sidebar.subheader("🎙️ AI Voice Options")
script_text = st.sidebar.text_area("Enter Your Script:", "Type something interesting...")
use_ai_voice = st.sidebar.checkbox("Use AI Voice Clone", value=True)

# Generate AI Video Button
if st.sidebar.button("🎬 Generate AI Video"):
    if source_img and target_img and script_text:
        with st.spinner("Processing AI Face Swap & Voice Clone... 🚀"):
            temp_source = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
            temp_target = tempfile.NamedTemporaryFile(delete=False, suffix=".jpg")
            temp_source.write(source_img.read())
            temp_target.write(target_img.read())

            # Perform Face Swap
            output_img_path = "swapped_face.jpg"
            swapped_face = swap_faces(temp_source.name, temp_target.name, output_img_path)

            if swapped_face:
                st.image(swapped_face, caption="AI-Swapped Face", use_column_width=True)

                # Generate AI Voice
                if use_ai_voice:
                    audio_path = generate_voice(script_text)
                    st.audio(audio_path, format="audio/wav")

                # Convert Swapped Image to Video
                video_clip = ImageSequenceClip([swapped_face] * 30, fps=10)
                if use_ai_voice:
                    video_audio = AudioFileClip(audio_path)
                    video_clip = video_clip.set_audio(video_audio)

                final_video_path = f"deepflick_output_{random.randint(1000,9999)}.mp4"
                video_clip.write_videofile(final_video_path, codec="libx264")

                st.video(final_video_path)

                # Cleanup temp files
                temp_source.close()
                temp_target.close()
                os.remove(temp_source.name)
                os.remove(temp_target.name)

    else:
        st.warning("⚠️ Please upload source & target images and enter a script.")