from functools import partial
from pathlib import Path
import os

import kangas as kg
import streamlit as st
import streamlit.components.v1 as components
from datasets import load_dataset

# If you duplicate this space, change this to match your space:
servername = 'comet-team-kangas-demo.hf.space'
st.set_page_config(layout="wide")
proj_dir = Path(__file__).parent

if 'iframe_source' not in st.session_state:
    st.session_state['iframe_source'] = f"https://{servername}/kangas/?datagrid=.%2Fdatagrids%2Fcoco-500.datagrid"

st.markdown("# HuggingFace Spaces + Streamlit + Kangas")
st.markdown("1. Enter the name of a HuggingFace dataset and Download, or select one below in Kangas")

def kangas_fn(dataset_repo):
    global src
    dg_file_name = dataset_repo.replace('/', '__') + '.datagrid'
    dg_full_path = str(proj_dir / 'datagrids' / dg_file_name)
    with st.spinner("2. Loading Dataset..."):
        dataset = load_dataset(dataset_repo, split="train")
    with st.spinner("3. Creating Kangas..."):
        if not os.path.exists(dg_full_path):
            dg = kg.DataGrid(dataset)
    with st.spinner("4. Saving Kangas..."):
        if not os.path.exists(dg_full_path):
            dg.save(dg_full_path)
    with st.spinner("5. Loading Kangas..."):
        st.session_state['iframe_source'] = f"https://{servername}/kangas/?datagrid=.%2Fdatagrids%2F" + dg_file_name
        
hf_dataset = st.text_input("HuggingFace Dataset", value='beans')
st.button("Download", on_click=partial(kangas_fn, hf_dataset))
st.markdown("""Click the dropdown in Kangas to select a previously-downloaded dataset""")
components.iframe(st.session_state['iframe_source'], None, 1000, scrolling=True)