# Import standard libs.
import tempfile
import os
from pathlib import Path

# Import primary libs.
import streamlit as st

# Import custom libs.
import fbx_handler


def process_file(file: Path) -> int:
    fbx_content = fbx_handler.FBXContainer(file)
    return 1


# Initialize session state variables if they don't exist
if "uploaded_files" not in st.session_state:
    st.session_state.uploaded_files = {}
if "processed_files" not in st.session_state:
    st.session_state.processed_files = {}

st.title('Optical MoCap AI Processing')

st.write('Select FBX files to upload and process. This will extract all marker animation data and turn it into a csv.')

new_uploaded_files = st.file_uploader('Select FBX files', accept_multiple_files=True, type='fbx', label_visibility='collapsed')

for uploaded_file in new_uploaded_files:
    if uploaded_file.name not in st.session_state.processed_files.keys():
        st.session_state.uploaded_files[uploaded_file.name] = uploaded_file

if st.session_state.uploaded_files and st.button("Process Files"):
    progress_bar = st.progress(0)
    # Create a temporary directory to store the newly uploaded files
    with tempfile.TemporaryDirectory() as temp_dir:
        incr = 1. / len(st.session_state.uploaded_files)

        for idx, (name, uploaded_file) in enumerate(st.session_state.uploaded_files.items()):
            # Save the uploaded file to the temporary directory
            temp_path = Path(os.path.join(temp_dir, name))
            with open(temp_path, "wb") as f:
                f.write(uploaded_file.getbuffer())
                print(f'[LOAD FBX] Finished uploading {temp_path}.')
                # Process the file and append the resulting DataFrame to dataframes
                st.session_state.processed_files[name] = process_file(temp_path)
                progress_bar.progress((idx+1) * incr, f'Processing {name}')

        st.session_state.uploaded_files = {}
        st.experimental_rerun()

for name in list(st.session_state.processed_files.keys()):
    new_file_name = name.replace('.fbx', '.csv')
    if st.download_button(
            label=f"Download {new_file_name}",
            data=st.session_state.processed_files[name],
            file_name=new_file_name,
            mime='text/csv'):
        del st.session_state.processed_files[name]
        st.experimental_rerun()

if st.button('Delete cache', type='primary'):
    st.session_state.uploaded_files = {}
    st.session_state.processed_files = {}
    st.experimental_rerun()