import gradio as gr
import spaces

from src.config import EXAMPLE_ROOT
from src.downloads import ensure_runtime_dirs
from src.gpu_timing import start_gpu_callback_timing
from src.jobs import JOB_ROOT
from src.preparation import estimate_prepared_gpu_seconds
from src.ui import build_demo, prepare_separation_with_progress, run_prepared_separation


ensure_runtime_dirs()
JOB_ROOT.mkdir(parents=True, exist_ok=True)


@spaces.GPU(duration=estimate_prepared_gpu_seconds)
def gpu_separate(
    prepared_state_value,
    execution_mode,
    progress=gr.Progress(track_tqdm=True),
):
    callback_timing = start_gpu_callback_timing(
        dispatch_mode="prepared-job-button-direct-to-spaces-gpu-callback"
    )
    return run_prepared_separation(
        prepared_state_value,
        execution_mode=execution_mode,
        progress=progress,
        callback_timing=callback_timing,
    )


demo = build_demo(
    gpu_separate,
    prepare_function=prepare_separation_with_progress,
)

if __name__ == "__main__":
    demo.queue(
        default_concurrency_limit=1,
        max_size=12,
        status_update_rate=0.5,
    ).launch(
        show_error=True,
        allowed_paths=[str(JOB_ROOT), str(EXAMPLE_ROOT)],
    )