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

import cv2

import numpy as np
import streamlit as st
from pydub import AudioSegment
import torch
import yaml
# from extract_video import extract_method_single_video

from utils import st_file_selector, img2base64
from evaluate_models import inference, load_model
from src import commons

import os

DEBUG = True

def main():
    st.markdown("###")    
    uploaded_file = st.file_uploader('Upload an audio file', type=['wav', 'mp3'], accept_multiple_files=False)

    with st.spinner(f'Loading samples...'):
        while not os.path.isdir("sample_files"):
            time.sleep(1)
    st.markdown("### or")
    selected_file = st_file_selector(st, path='sample_files', key = 'selected', label = 'Choose a sample image/video')

    if uploaded_file: 
        random_id = uuid.uuid1()
        ext = uploaded_file.name.split('.')[-1]
        
        base_folder = "temps"
        filename = "{}.{}".format(random_id, ext)
        file_type = uploaded_file.type.split("/")[0]
        filepath = f"{base_folder}/{filename}"

        uploaded_file_length = len(uploaded_file.getvalue())
        if uploaded_file_length > 0:
            with open(filepath, 'wb') as f:
                f.write(uploaded_file.read())
            st.audio(uploaded_file, format=ext)
    elif selected_file:
        base_folder = "sample_files"
        file_type = selected_file.split(".")[-1]
        filename = selected_file.split("/")[-1]
        filepath = f"{base_folder}/{selected_file}"

        st.write('file_type', file_type)
        with open(filepath, 'rb') as f:
            audio_bytes = f.read()
        st.audio(audio_bytes, format=file_type)
    else:
        return


    

    with st.spinner(f'Analyzing {file_type}...'):
        

        seed = config["data"].get("seed", 42)
        # fix all seeds - this should not actually change anything
        commons.set_seed(seed)

        result = inference(
            model,
            datasets_path=filepath,
            device=device,
        )
        result = result[0]

    if 'Real' == result[0]:
        st.success(f'Audio is real! \nprob: {result[1]}', icon="✅")
    else:
        st.error(f'Audio is fake! \nprob: {result[1]}', icon="🚨")
    
    st.divider()
    st.write('## Response JSON')
    st.write(result)


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

    

if __name__ == "__main__":
    # from download_whisper import download_whisper, extract_and_save_encoder
    # model = download_whisper()
    # extract_and_save_encoder(model)\

    # if torch.cuda.is_available():
    #     device = "cuda"
    # else:
    #     device = "cpu"
    device = 'cpu'

    with open('config.yaml', "r") as f:
        config = yaml.safe_load(f)

    model = load_model(config, device)

    st.title("Face Fake Detection")
    setup()
    main()