import time
from datetime import UTC, datetime

import gradio as gr
import spaces
import torch

GENERATE_SECONDS = 5
VISUALIZE_SECONDS = 3


def _now() -> str:
    return datetime.now(UTC).strftime("%H:%M:%S")


@spaces.GPU(duration=10)
def generate(text: str) -> str:
    start = _now()
    gpu_name = torch.cuda.get_device_name()
    print(f"[generate] start={start} input={text!r} gpu={gpu_name}")
    time.sleep(GENERATE_SECONDS)
    end = _now()
    print(f"[generate] end={end} input={text!r}")
    return f"generate | start {start} | end {end} | gpu {gpu_name} | input {text!r}"


def visualize(text: str) -> str:
    start = _now()
    print(f"[visualize] start={start} input={text!r}")
    time.sleep(VISUALIZE_SECONDS)
    end = _now()
    print(f"[visualize] end={end} input={text!r}")
    return f"visualize | start {start} | end {end}\n{text}"


demo = gr.Workflow(bind={"generate": generate, "visualize": visualize})

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