import spaces
import datetime
import os
import subprocess
import gradio as gr


CUSTOM_CSS = """
#output_box textarea {
    font-family: IBM Plex Mono, ui-monospace, SFMono-Regular, Menlo, Monaco, Consolas, "Liberation Mono", "Courier New", monospace;
}
"""




@spaces.GPU
def run():
    output: str = ""
    try:
        output = subprocess.check_output(["nvidia-smi"], text=True)
    except FileNotFoundError:
        output = subprocess.check_output(["ls", "-alh"], text=True)
    comment = (
        datetime.datetime.now().replace(microsecond=0).isoformat().replace("T", " ")
    )
    return f"# {comment}\n\n{output}"


def run_custom_command(custom_command: str, secret: str):
    if secret != os.environ.get("SECRET"):
        return "You can't access this"

    print("custom_command", custom_command)
    try:
        return subprocess.check_output(custom_command.split(), text=True)
    except Exception as e:
        return f"{e}"


output = gr.Textbox(
    label="Command Output", max_lines=32, elem_id="output_box", value=run()
)

with gr.Blocks(css=CUSTOM_CSS) as demo:
    gr.Markdown("#### `nvidia-smi`: How is my GPU Space running right now 🔥")

    with gr.Accordion(label="Power user mode", open=False):
        custom_command = gr.Textbox(label="Input command", value="pwd")
        secret = gr.Textbox(
            label="Secret",
        )
        custom_command_btn = gr.Button("Run")
        custom_command_btn.click(
            fn=run_custom_command,
            inputs=[custom_command, secret],
            outputs=output,
        )

    output.render()

    demo.load(fn=run, inputs=None, outputs=output, every=1)


demo.queue().launch()