#!/usr/bin/env python

from __future__ import annotations

import os

import gradio as gr
import pandas as pd
from apscheduler.schedulers.background import BackgroundScheduler
from huggingface_hub import HfApi

from demo_list import DemoList

INTERVAL_MIN = int(os.getenv('INTERVAL_MIN', '10'))

demo_list = DemoList()

api = HfApi()
scheduler = BackgroundScheduler()
scheduler.add_job(func=lambda: api.restart_space(os.getenv('SPACE_ID')),
                  trigger='interval',
                  seconds=60 * INTERVAL_MIN)
scheduler.start()

STATUS_CHOICES = [
    'RUNNING',
    'PAUSED',
    'STOPPED',
    'RUNTIME_ERROR',
    'BUILD_ERROR',
    'BUILDING',
]
HARDWARE_CHOICES = [
    'cpu-basic',
    'cpu-upgrade',
    't4-small',
    't4-medium',
    'zero-a10g',
    'a10g-small',
    'a10g-large',
    'a100-large',
]
SDK_CHOICES = [
    'gradio',
    'streamlit',
    'docker',
]


def update_df(status: list[str], hardware: list[str],
              sdk: list[str]) -> pd.DataFrame:
    df_raw = demo_list.df_raw
    df = demo_list.df
    return df[(df_raw.status.isin(status)) & (df_raw.hardware.isin(hardware)) &
              (df_raw.sdk.isin(sdk))]


def update_status_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return STATUS_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


def update_hardware_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return HARDWARE_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


def update_sdk_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return SDK_CHOICES
    elif '(NONE)' in choices:
        return []
    else:
        return choices


with gr.Blocks(css='style.css') as demo:
    with gr.Accordion(label='Filter', open=False):
        status = gr.CheckboxGroup(label='Status',
                                  choices=STATUS_CHOICES + ['(ALL)', '(NONE)'],
                                  value=STATUS_CHOICES,
                                  type='value')
        hardware = gr.CheckboxGroup(label='Hardware',
                                    choices=HARDWARE_CHOICES +
                                    ['(ALL)', '(NONE)'],
                                    value=HARDWARE_CHOICES,
                                    type='value')
        sdk = gr.CheckboxGroup(label='SDK',
                               choices=SDK_CHOICES + ['(ALL)', '(NONE)'],
                               value=SDK_CHOICES,
                               type='value')
        apply_button = gr.Button('Apply')
    df = gr.Dataframe(value=demo_list.df,
                      datatype=demo_list.column_datatype,
                      type='pandas')

    status.change(fn=update_status_checkboxes,
                  inputs=status,
                  outputs=status,
                  queue=False,
                  show_progress=False,
                  api_name=False)
    hardware.change(fn=update_hardware_checkboxes,
                    inputs=hardware,
                    outputs=hardware,
                    queue=False,
                    show_progress=False,
                    api_name=False)
    sdk.change(fn=update_sdk_checkboxes,
               inputs=sdk,
               outputs=sdk,
               queue=False,
               show_progress=False,
               api_name=False)
    apply_button.click(fn=update_df,
                       inputs=[status, hardware, sdk],
                       outputs=df,
                       api_name=False)
demo.queue(api_open=False).launch()