from functools import partial
from pathlib import Path

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

proj_dir = Path(__file__).parent

servername = 'hellofriend22-kangas-demo.hf.space'
src = f"https://{servername}:443"

st.set_page_config(layout="wide")

st.markdown("1. Select dataset of your choice")


def kangas_fn(dataset_repo):
    with st.spinner("Loading Dataset..."):
        dataset = load_dataset(dataset_repo, split="train")
    with st.spinner("Creating Kangas..."):
        dg = kg.DataGrid(dataset)
        dg_file_name = dataset_repo.replace('/', '__') + '.datagrid'
    with st.spinner("Saving Kangas..."):
        dg.save(str(proj_dir / 'datagrids' / dg_file_name))


kg.launch(host=servername)
height = st.sidebar.slider("iFrame Height", 200, 1500, 900, 100)
scrolling = st.sidebar.checkbox("iFrame Scrolling")

hf_dataset = st.text_input("HuggingFace Dataset", value='beans')
st.button("Download and Run", on_click=partial(kangas_fn, hf_dataset))
st.markdown("""Click the dropdown in Kangas to see pre-loaded datasets""")
st.components.v1.iframe(src, None, height, scrolling=True)