from os import getenv
from textwrap import dedent
from torch import cuda

import gradio as gr
from src.automatikz.examples.webui.webui import build_ui, get_banner

PUBLIC_DEMO = getenv("SPACE_ID") == "nllg/AutomaTikZ"

if PUBLIC_DEMO and not cuda.is_available():
    with gr.Blocks(css=".gradio-container {text-align: center}", theme=gr.themes.Soft(), title="AutomaTikZ") as demo:
        badge = "https://huggingface.co/datasets/huggingface/badges/resolve/main/duplicate-this-space-xl.svg"
        link = "https://huggingface.co/spaces/nllg/AutomaTikZ?duplicate=true"
        html = f'<a style="display:inline-block" href="{link}"> <img src="{badge}" alt="Duplicate this Space"> </a>'
        message = dedent("""\
        The size of our models exceeds the resource constraints offered by the
        free tier of Hugging Face Spaces. For full functionality, we recommend
        duplicating this space on a paid private GPU runtime.
        """)
        gr.Markdown(f'{get_banner()}\n{message}\n{html}')
        demo.launch()
else:
    build_ui(lock=PUBLIC_DEMO).queue().launch(server_name="0.0.0.0", server_port=7860)