import gradio as gr
from huggingface_hub import HfApi
import os


HF_TOKEN = os.environ.get("HF_TOKEN")
HF_USER = os.environ.get("HF_USER")


def wakeup_spaces(username: str, hf_token: str="", is_private=True, progress=gr.Progress(track_tqdm=True)):
    try:
        api = HfApi(token=hf_token if hf_token else HF_TOKEN)
        space_infos = api.list_spaces(author=username)
        spaces = []
        logs = []
        for s in space_infos:
            if s.private and not is_private: continue
            spaces.append(s.id)
        for s in spaces:
            stage = api.get_space_runtime(repo_id=s).stage
            if stage in ("RUNNING", "APP_STARTING", "BUILDING"):
                print(f"'{s}' is running.")
                progress(0, desc=f"'{s}' is running.")
            else:
                print(f"'{s}' isn't running. Restarting...")
                progress(0, desc=f"'{s}' isn't running. Restarting...")
                api.restart_space(repo_id=s)
                logs.append(f"'{s}' is in '{stage}'.")
                logs.append(f"'{s}' is restarted.")
        return "<br>\n".join(logs)
    except Exception as e:
        print(f"Error: {e}")
        raise gr.Error(f"Error: {e}")


with gr.Blocks() as demo:
    with gr.Row():
        username = gr.Textbox(label="HF Username", value=HF_USER, lines=1)
        token = gr.Textbox(label="HF Write Token", value="", lines=1)
        is_private = gr.Checkbox(label="Apply to private space", value=True)
    run_button = gr.Button("Wake up spaces", variant="primary")
    info_md = gr.Markdown("<br><br><br>")

    run_button.click(wakeup_spaces, [username, token, is_private], [info_md])

demo.launch()