#!/usr/bin/env python

from __future__ import annotations

import gradio as gr

from constants import (HARDWARE_CHOICES, OWNER_CHOICES, SDK_CHOICES,
                       SLEEP_TIME_CHOICES, STATUS_CHOICES, VISIBILITY_CHOICES)
from demo_list import DemoList, get_df_from_yaml
from restart_scheduler import RestartScheduler
from settings import HF_TOKEN, INTERVAL_MINUTES, SPACE_ID

demo_list = DemoList(get_df_from_yaml('list.yaml'))

if SPACE_ID is not None and INTERVAL_MINUTES > 0:
    scheduler = RestartScheduler(space_id=SPACE_ID,
                                 interval_minutes=INTERVAL_MINUTES,
                                 hf_token=HF_TOKEN)
    scheduler.start()


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


def update_sleep_time_checkboxes(choices: list[str]) -> list[str]:
    if '(ALL)' in choices:
        return SLEEP_TIME_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=['(ALL)', '(NONE)'] + STATUS_CHOICES,
                                  value=STATUS_CHOICES,
                                  type='value')
        hardware = gr.CheckboxGroup(label='Hardware',
                                    choices=['(ALL)', '(NONE)'] +
                                    HARDWARE_CHOICES,
                                    value=HARDWARE_CHOICES,
                                    type='value')
        sleep_time = gr.CheckboxGroup(label='Sleep time',
                                      choices=['(ALL)', '(NONE)'] +
                                      SLEEP_TIME_CHOICES,
                                      value=SLEEP_TIME_CHOICES,
                                      type='value')
        multiple_replicas = gr.Checkbox(label='Multiple replicas', value=False)
        sdk = gr.CheckboxGroup(label='SDK',
                               choices=['(ALL)', '(NONE)'] + SDK_CHOICES,
                               value=SDK_CHOICES,
                               type='value')
        visibility = gr.CheckboxGroup(label='Visibility',
                                      choices=VISIBILITY_CHOICES,
                                      value=VISIBILITY_CHOICES,
                                      type='value')
        owner = gr.CheckboxGroup(label='Owner',
                                 choices=OWNER_CHOICES,
                                 value=OWNER_CHOICES,
                                 type='value')
        apply_button = gr.Button('Apply')
    df = gr.Dataframe(value=demo_list.df_prettified,
                      datatype=demo_list.column_datatype,
                      type='pandas',
                      interactive=False)

    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)
    sleep_time.change(fn=update_sleep_time_checkboxes,
                      inputs=sleep_time,
                      outputs=sleep_time,
                      queue=False,
                      show_progress=False,
                      api_name=False)
    apply_button.click(fn=demo_list.filter,
                       inputs=[
                           status,
                           hardware,
                           sleep_time,
                           multiple_replicas,
                           sdk,
                           visibility,
                           owner,
                       ],
                       outputs=df,
                       api_name=False)
demo.queue(api_open=False).launch()