import base64
import json
import os, shutil
import re
import time
import uuid

import cv2

import numpy as np
import streamlit as st
from PIL import Image
# from extract_video import extract_method_single_video

import shlex
import subprocess

import os

def main():
    st.markdown("###")    
    uploaded_file = st.file_uploader('Upload a picture', type=['mp4', 'jpg', 'jpeg', 'png'], accept_multiple_files=False)
    if uploaded_file: 
        random_id = uuid.uuid1()
        filename = "{}.{}".format(random_id, uploaded_file.type.split("/")[-1])
        file_type = uploaded_file.type.split("/")[0]

        if uploaded_file.type == 'video/mp4':
            with open(f"temps/{filename}", mode='wb') as f:
                f.write(uploaded_file.read())
            st.video(uploaded_file)
            pass
        else:
            img = Image.open(uploaded_file).convert('RGB')
            ext = uploaded_file.type.split("/")[-1]
            with open(f"temps/{filename}", mode='wb') as f:
                f.write(uploaded_file.getbuffer())  
            st.image(img)

        

        with st.spinner(f'Processing {file_type}...'):
            subprocess.run(shlex.split(f"python.exe extract_video.py --device cuda --max_frames 50 --bs 2 --frame_interval 5 --confidence_threshold 0.997 --data_path temps/{filename}"))
        st.text(f'1. Processing {file_type} ✅')
        with st.spinner(f'Analyzing {file_type}...'):
            pred = subprocess.run(shlex.split(f"python inference.py --weight weights/model_params_ffpp_c23.pickle --device cuda --image_folder temps/images/{filename}"), capture_output=True)
        st.text(f'2. Analyzing {file_type} ✅')
        
        print(pred)
        try:
            fake_probability = float(pred.stdout.decode('utf-8').split('Mean prediction: ')[-1])
            if fake_probability > 0.6:
                st.error('  FAKE! ', icon="🚨")
            else:
                st.success("  REAL FOOTAGE! ", icon="✅")
            st.text("fake probability {:.2f}".format(fake_probability))

            os.remove(f"temps/{filename}")
            folder_name = ".".join(filename.split(".")[:-1])
            shutil.rmtree(f"temps/images/{folder_name}")
        except:
            st.text(pred.stdout.decode('utf-8'))

            st.text("")
            st.text(pred)


def setup():
    if not os.path.isdir("temps"):
        os.makedirs("temps")

if __name__ == "__main__":
    st.set_page_config(
        page_title="Nodeflux Deepfake Detection", page_icon=":pencil2:"
    )
    st.title("Deepfake Detection")
    setup()
    main()