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


st.set_page_config(layout="wide")
# @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.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,
            "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,
            "kaggle": f"/kaggle/repo/vlsp2023-ess/config/exp{step}",
        }
    )


df = pd.DataFrame(_models)

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