import gradio as gr
from huggingface.user import HFUser, GR_CONF

Theme = gr.Theme.load(GR_CONF["theme"])
GR_CONF["theme"] = Theme

def login(token):
    u = HFUser.from_token(token)
    return u, u.name, gr.Column(visible=False)

def show_time(u, name):
    return u.ping(name), gr.Column(visible=True)

def list_dataset(u, repo):
    files = u.list_dataset(repo)
    return gr.Dropdown(value=files[0], choices=files), gr.Column(visible=True), gr.Column(visible=False)

def fetch_parquet(u, fname):
    _cache = u.fetch_file(fname)
    return _cache

def split_parquet(u, file, batch_size):
    batch_size = int(batch_size)
    file_slice = u.split_parquet(file, batch_size)
    return file_slice[0][0], file_slice, gr.Slider(value=0, maximum=batch_size-1), gr.Column(visible=True)

def select_video(chunks, epoch_idx , batch_idx):
    epoch_idx = int(epoch_idx)
    batch_idx = int(batch_idx)
    return chunks[epoch_idx][batch_idx]

def show_lables():
    return gr.Column(visible=True)

def next_chunks(video_chunks, epoch_idx):
    length = len(video_chunks)
    return (epoch_idx+1)%length, gr.Slider(value=0)

with gr.Blocks(**GR_CONF) as Core:
    user = gr.State()
    epoch_idx = gr.State(0)
    video_chunks = gr.State()

    with gr.Row(variant="panel"):
        with gr.Column(scale=6):
            _video = gr.Video(height=720)
        with gr.Column(scale=2):
            with gr.Column() as Auth:
                _token = gr.Textbox(label="Huggingface Token")
                _auth = gr.Button("Auth", variant="primary", size="lg")

            with gr.Row() as UUID:
                name= gr.Textbox(label="Name", interactive=False, scale=1)
                time= gr.Textbox(label="Time", interactive=False, scale=1)

            with gr.Column(visible=False) as Repo:
                raw_dataset= gr.Textbox("OpenVideo/pexels-raw", label="Raw Dataset")
                _list = gr.Button("List", variant='secondary', size='sm')

            with gr.Column(visible=False) as Batch:
                file = gr.Dropdown(label="Parquet")
                with gr.Row():
                    _cache= gr.Textbox("Downloading", label="Cache")
                    batch_size= gr.Textbox("8", label="Batch")
                _fetch = gr.Button("Fetch", variant='primary', size='sm')

            with gr.Column(visible=False) as Pick:
                _pick = gr.Slider(0, 7, value=0, step=1, label="Batch", info="Choose between 1 and $BATCH")
                gr.Label()

    with gr.Row(variant="panel", visible=False) as Tag:
        _human_tag = gr.Textbox(label="Tag", scale=2)
        with gr.Column():
            submit = gr.Button("Submit", variant="primary", size="sm", scale=1)
            with gr.Row():
                rst = gr.Button("Reset", variant="stop", size="sm", scale=1)
                nxt = gr.Button("Next Batch", variant="secondary", size="sm", scale=1)

    _auth.click(fn=login, inputs=_token, outputs=[user, name,  Auth])
    name.change(fn=show_time, inputs=[user, name], outputs=[time, Repo])
    _list.click(fn=list_dataset, inputs=[user, raw_dataset], outputs=[file, Batch,  Repo])
    _fetch.click(fn=fetch_parquet, inputs=[user, file], outputs=[_cache] )
    file.change(fn=fetch_parquet, inputs=[user, file], outputs=[_cache] )
    _cache.change(fn=split_parquet, inputs=[user, _cache, batch_size], outputs=[_video, video_chunks, _pick, Pick])
    _pick.change(fn=select_video, inputs=[video_chunks, epoch_idx, _pick], outputs=_video)
    _video.change(fn=show_lables, outputs=Tag)
    nxt.click(fn=next_chunks, inputs=[video_chunks, epoch_idx], outputs=[epoch_idx, _pick])

if __name__ == "__main__":
    Core.launch()