import streamlit as st
from azure.storage.blob import BlobServiceClient, BlobClient, ContainerClient
import base64
import os

# Session State workaround
class SessionState(object):
    def __init__(self, **kwargs):
        for key, val in kwargs.items():
            setattr(self, key, val)


def get_state(**kwargs):
    if 'session_state' not in st.session_state:
        st.session_state['session_state'] = SessionState(**kwargs)
    return st.session_state['session_state']


def save_to_azure(files, connect_str, container_name):
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    for file in files:
        blob_client = blob_service_client.get_blob_client(container_name, file.name)
        blob_client.upload_blob(file, overwrite=True)


def list_blobs(connect_str, container_name):
    blob_service_client = BlobServiceClient.from_connection_string(connect_str)
    container_client = blob_service_client.get_container_client(container_name)

    blob_list = container_client.list_blobs()
    for blob in blob_list:
        blob_url = f"https://{blob_service_client.account_name}.blob.core.windows.net/{container_name}/{blob.name}"
        b64 = base64.b64encode(blob_url.encode()).decode()  
        href = f'<a href="data:file/octet-stream;base64,{b64}" download="{blob.name}">Download {blob.name}</a>'
        st.markdown(href, unsafe_allow_html=True)


def app():
    st.title('Azure Blob Storage App 💾')

    state = get_state(connect_str='', container_name='')

    with st.sidebar:
        st.subheader("Azure Settings 🔧")
        state.connect_str = st.text_input('Connection String', value=state.connect_str)
        state.container_name = st.text_input('Container Name', value=state.container_name)

        st.subheader("Your documents 📑")
        docs = st.file_uploader("Import documents", accept_multiple_files=True)

        if st.button('Save Files'):
            with st.spinner("Saving..."):
                save_to_azure(docs, state.connect_str, state.container_name)

        if st.button('Retrieve Files'):
            with st.spinner("Retrieving..."):
                list_blobs(state.connect_str, state.container_name)


if __name__ == "__main__":
    app()