import streamlit as st
from huggingface_hub import list_models
from datetime import date, timedelta, datetime, timezone
import pandas as pd
import asyncio

# os.environ["KAGGLE_USERNAME"] = "hahunavth"
# os.environ["KAGGLE_KEY"] = ""
from kaggle.api.kaggle_api_extended import KaggleApi

st.set_page_config(layout="wide")

st.header("Experiments")

@st.cache_resource
def init_kaggle():
    api = KaggleApi()
    api.authenticate()
    return api

api = init_kaggle()

# @st.cache_resource
# def get_model_list():
#     return [model for model in list(list_models(author="hahunavth", filter="emofs2"))]

def get_ckpt_list(model):
    files = model.siblings
    file_names = [file.rfilename for file in files if "pytorch_model" in file.rfilename]
    return file_names

with st.expander("Filter:"):
    with st.form("my_form"):
        _author = st.text_input("Author", "hahunavth")
        _filter = st.text_input("Slug", "emofs2")
        submit = st.form_submit_button("Submit")

models = list(
    list_models(author=_author, filter=_filter, sort="last_modified", direction=-1)
)

_models = []
now = datetime.now(timezone.utc)
for model in models:
    ckpts = get_ckpt_list(model)
    ckpts_step = [ckpt.replace("pytorch_model.", "").replace(".bin", "") for ckpt in ckpts]
    ckpts_step = [int(ckpt) for ckpt in ckpts_step if ckpt.isdigit()]
    ckpts_step.sort()
    last_ckpt = "{:,}".format(ckpts_step[-1]) if ckpts_step else None
    diff = None
    if model.last_modified:
        now = now.replace(microsecond=0)
        last_modified = model.last_modified.replace(microsecond=0)
        diff = now - last_modified
        diff = str(diff)

    step = model.id.replace('hahunavth/emofs2-exp', '').split('_')[0]

    _models.append(
        {
            "id": model.id,
            "last_modified": diff,
            "ckpt_list": last_ckpt,
            "status": None,
            "huggingface": f'https://huggingface.co/{model.id}/tree/main' if model.id else None,
            "notebook": f"https://www.kaggle.com/code/hahunavth/ess-vlsp2023-train-{step}" if model.id else None,
            "config": f"/kaggle/repo/vlsp2023-ess/config/exp{step}",
            # "command": "st.balloons",
            # "is_widget": True
        }
    )

df = pd.DataFrame(_models)

@st.cache_data
def get_kernel_status(step):
    kaggle_slug = f"ess-vlsp2023-train-{step}"
    try:
        kernel_status = api.kernel_status(user_name="hahunavth", kernel_slug=kaggle_slug)["status"]
    except:
        kernel_status = "not found"
    return kernel_status

async def set_df_kernel_status(df, n=12):
    for i in range(min(len(df), n)):
        step = df.iloc[i]["id"].replace('hahunavth/emofs2-exp', '').split('_')[0]
        status = get_kernel_status(step)
        df.at[i, "status"] = status
        print(i, status)

def on_update_kernel_status_click():
    st.session_state["last_update"] = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    get_kernel_status.clear()

st.button("Update kernel status", on_click=on_update_kernel_status_click)
st.text(f"Last updated: {st.session_state.get('last_update', 'not yet')}")

asyncio.run(set_df_kernel_status(df))

edit_df = st.data_editor(
    df,
    column_config={
        "huggingface": st.column_config.LinkColumn(),
        "notebook": st.column_config.LinkColumn(),
        # "a": st.button("Click me"),
    }, 
)

# st.dataframe(
#     edit_df, 
#     column_config={
#         "huggingface": st.column_config.LinkColumn(),
#         "notebook": st.column_config.LinkColumn(),
#     }, 
#     use_container_width=True
# )