import os

# cudaMallocAsync bypasses NVML memory queries that fail on MIG GPU instances
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "backend:cudaMallocAsync")

import gradio as gr

from examples_ui import EXAMPLE_CARDS_HTML, SUGGESTION_CHIPS_HTML, load_example_data
from inference import MAX_SEED, infer

with open("static/app.css") as _f:
    css = _f.read()

with open("static/gallery.js") as _f:
    gallery_js = _f.read()

with open("static/wire_outputs.js") as _f:
    wire_outputs_js = _f.read()

with open("static/run_preprocess.js") as _f:
    run_preprocess_js = _f.read()

with open("static/mode_toggle.js") as _f:
    mode_toggle_js = _f.read()

with open("templates/app.html") as _f:
    app_html = _f.read().format(
        example_cards_html=EXAMPLE_CARDS_HTML,
        suggestion_chips_html=SUGGESTION_CHIPS_HTML,
    )

with gr.Blocks() as demo:

    hidden_images_b64 = gr.Textbox(value="[]", elem_id="hidden-images-b64", elem_classes="hidden-input", container=False)
    prompt = gr.Textbox(value="", elem_id="prompt-gradio-input", elem_classes="hidden-input", container=False)
    seed = gr.Slider(minimum=0, maximum=MAX_SEED, step=1, value=0, elem_id="gradio-seed", elem_classes="hidden-input", container=False)
    randomize_seed = gr.Checkbox(value=True, elem_id="gradio-randomize", elem_classes="hidden-input", container=False)
    guidance_scale = gr.Slider(minimum=1.0, maximum=10.0, step=0.1, value=1.2, elem_id="gradio-guidance", elem_classes="hidden-input", container=False)
    steps = gr.Slider(minimum=1, maximum=50, step=1, value=3, elem_id="gradio-steps", elem_classes="hidden-input", container=False)
    mode = gr.Textbox(value="fast", elem_id="gradio-mode", elem_classes="hidden-input", container=False)
    gpu_duration = gr.Slider(minimum=10, maximum=120, step=5, value=30, elem_id="gradio-gpu-duration", elem_classes="hidden-input", container=False)
    result = gr.Image(elem_id="gradio-result", elem_classes="hidden-input", container=False, format="png")

    example_idx = gr.Textbox(value="", elem_id="example-idx-input", elem_classes="hidden-input", container=False)
    example_result = gr.Textbox(value="", elem_id="example-result-data", elem_classes="hidden-input", container=False)
    example_load_btn = gr.Button("Load Example", elem_id="example-load-btn")

    gr.HTML(app_html)

    run_btn = gr.Button("Run", elem_id="gradio-run-btn")

    demo.load(fn=None, js=gallery_js)
    demo.load(fn=None, js=wire_outputs_js)
    demo.load(fn=None, js=mode_toggle_js)

    run_btn.click(
        fn=infer,
        inputs=[hidden_images_b64, prompt, seed, randomize_seed, guidance_scale, steps, mode, gpu_duration],
        outputs=[result, seed],
        js=run_preprocess_js,
    )

    example_load_btn.click(
        fn=load_example_data,
        inputs=[example_idx],
        outputs=[example_result],
        queue=False,
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(
        css=css,
        mcp_server=True,
        ssr_mode=False,
        show_error=True,
        allowed_paths=["examples"],
    )