import streamlit as st
import os
import mimetypes
from deface.deface import anonymize_frame
from deface.centerface import CenterFace
from PIL import Image
import imageio.v2 as imageio
import cv2
import tempfile
import os

def streamlit_video_deface_cli(uploaded_file, keep_audio):
    temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    processed_video_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name

    with open(temp_video_path, 'wb') as f:
        f.write(uploaded_file.read())

    if keep_audio:
        command = f"deface {temp_video_path} --keep-audio -o {processed_video_temp_path}"
    else:
        command = f"deface {temp_video_path} -o {processed_video_temp_path}"
    
    os.system(command)

    return processed_video_temp_path


def process_video_frame(frame, centerface, threshold, replacewith, mask_scale, ellipse, draw_scores, replaceimg=None, mosaicsize=20):
    dets, _ = centerface(frame, threshold=threshold)
    anonymize_frame(dets, frame, mask_scale=mask_scale, replacewith=replacewith, ellipse=ellipse, draw_scores=draw_scores, replaceimg=replaceimg, mosaicsize=mosaicsize)
    return frame


def streamlit_video_detect(uploaded_file, centerface, threshold, replacewith, mask_scale, ellipse, draw_scores, enable_preview, keep_audio, replaceimg=None, mosaicsize=20):
    temp_video_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    processed_video_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    output_with_audio_temp_path = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
    with open(temp_video_path, 'wb') as f:
        f.write(uploaded_file.read())
    
    vidcap = cv2.VideoCapture(temp_video_path)
    
    # Get video properties
    width  = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
    height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
    fps    = vidcap.get(cv2.CAP_PROP_FPS)
    total_frames = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT))

    
    # Create a video writer object
    out = cv2.VideoWriter(processed_video_temp_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
    
    # Progress Bar setup
    progress_text = "Processing video. Please wait."
    my_bar = st.progress(0, text=progress_text)
    frames_processed = 0
    
    while True:
        ret, frame = vidcap.read()
        if not ret:
            break
        
        processed_frame = process_video_frame(frame, centerface, threshold, replacewith, mask_scale, ellipse, draw_scores, replaceimg, mosaicsize)
        out.write(processed_frame)

        # Progress Bar update
        frames_processed += 1
        progress_percent = int((frames_processed / total_frames) * 100)
        my_bar.progress(progress_percent, text=progress_text)
    
    vidcap.release()
    out.release()
    
    # Empty the progress bar after completion
    my_bar.empty()
    st.write("Re-encoding video. Please wait..."")
    os.system(f"ffmpeg -y -i {processed_video_temp_path} -c:v libx264 {processed_video_temp_path.split('.')[0]}_processed_.mp4")
    # If keep_audio is checked, use FFmpeg to overlay the original audio on the processed video
    if keep_audio:
        try:
            st.write("Overlaying audio. Please wait...")
            command = f"ffmpeg -y -i {processed_video_temp_path.split('.')[0]}_processed_.mp4 -i {temp_video_path} -c:v copy -c:a aac -strict experimental {output_with_audio_temp_path}"
            os.system(command)
            return output_with_audio_temp_path
        except:
            return f"{processed_video_temp_path.split('.')[0]}_processed_.mp4"
    else:
        return f"{processed_video_temp_path.split('.')[0]}_processed_.mp4"


def streamlit_image_detect(uploaded_file, centerface, threshold, replacewith, mask_scale, ellipse, draw_scores, enable_preview, keep_metadata, replaceimg=None, mosaicsize=20):
    # Read the uploaded image into a numpy array
    frame = imageio.imread(uploaded_file)

    # For the sake of this example, I'm skipping the metadata part
    
    dets, _ = centerface(frame, threshold=threshold)
    anonymize_frame(dets, frame, mask_scale=mask_scale, replacewith=replacewith, ellipse=ellipse, draw_scores=draw_scores, replaceimg=replaceimg, mosaicsize=mosaicsize)

    # Convert numpy array back to a PIL.Image so that we can display/save it easily in Streamlit
    result_img = Image.fromarray(frame)
    return result_img

def get_file_type(uploaded_file):
    mime = mimetypes.guess_type(uploaded_file.name)[0]
    if mime is None:
        return None
    if mime.startswith('video'):
        return 'video'
    if mime.startswith('image'):
        return 'image'
    return mime

def main():
    st.title("Media Anonymizer")
    
    uploaded_files = st.file_uploader("Upload media files (images,videos)", accept_multiple_files=True)
    keep_audio = st.checkbox("Keep Audio")

    for uploaded_file in uploaded_files:
        file_type = get_file_type(uploaded_file)
        
        if file_type == 'image':
            # Initialize centerface or other configurations if needed
            centerface = CenterFace() 
            result_img = streamlit_image_detect(uploaded_file, centerface, threshold=0.2, replacewith='blur', mask_scale=1.0, ellipse=True, draw_scores=False, enable_preview=False, keep_metadata=False)
            st.image(result_img, caption=f'Anonymized {uploaded_file.name}', use_column_width=True)
        elif file_type == 'video':
            centerface = CenterFace()
            processed_video_path = streamlit_video_detect(uploaded_file, centerface, threshold=0.2, replacewith='blur', mask_scale=1.0, ellipse=True, draw_scores=False, enable_preview=False, keep_audio=keep_audio)
            st.video(processed_video_path, format='video/mp4', start_time=0)
        else:
            st.write(f"Unsupported file type for {uploaded_file.name}")

if __name__ == "__main__":
    main()