import streamlit as st
import os
import base64
import atexit


def write_xyz_file(kvl_file, filename, keys):
    with open(filename, 'w') as f:
        nt = len(kvl_file['TIME'])
        dt = int(kvl_file['TIME'][1]) - int(kvl_file['TIME'][0])
        sampleno = len(kvl_file['TOP']) / nt

        for time in kvl_file['TIME']:
            start_index = int(time * sampleno)
            end_index = start_index + int(sampleno)
            for key in keys:
                if key == 'SEALEVEL':
                    index=int(time/dt)
                    sealevel = kvl_file[key][index]
                    f.write(f"0 0 {sealevel}\n")
                else:
                    matrix = kvl_file[key][start_index:end_index]
                    for i, row in enumerate(matrix):
                        for j, value in enumerate(row):
                            f.write(f"{i} {j} {value}\n")

def get_binary_file_downloader_html(bin_file, file_label='File'):
    with open(bin_file, 'rb') as f:
        data = f.read()
    bin_str = base64.b64encode(data).decode()
    href = f'<a href="data:application/octet-stream;base64,{bin_str}" download="{os.path.basename(bin_file)}">{file_label}</a>'
    return href

def cleanup(filename):
    os.remove('./' + filename)

# Write the XYZ file
def output_xyz(kvl_file):
    st.title('Convert KVL to XYZ')

    if not kvl_file:
        st.warning('No file has been uploaded')
    else:
        outkeys = st.multiselect('Select Keys to output, add SEALEVEL first if you want that key in as (0 0 z)', kvl_file.keys())
        filename = st.text_input('Enter output filename', 'output.xyz')
        write_xyz_file(kvl_file, filename, outkeys)

    # Function to create a download link

    if filename:
    # Use the function to create a download link for the XYZ file
        st.markdown(get_binary_file_downloader_html(filename, 'Download XYZ file'), unsafe_allow_html=True)
        atexit.register(cleanup, filename)