"""Inference backend for the DiffusionDB SD 1.5 LoRA.

A Gradio app whose only job is to expose generation as an API. The UI Space
(spaces/whosouravsharma/diffusiondb-sd15-lora) calls it with gradio_client;
the minimal interface here is also usable directly for debugging.

This is the Space that needs the GPU. Weights are pulled from the model repo
at first request, not baked into the image.
"""

import os

import gradio as gr

from src.pipeline import LORA_CHECKPOINT, LORA_REPO, ModelRunner, pick_device
from src.tracing import (
    TRACING_ACTIVE,
    configure_logging,
    continue_trace,
    logger,
)

configure_logging()

# Loaded on first request rather than at import, so the Space becomes
# reachable immediately and a slow weight download cannot kill startup.
_runner = None


def get_runner() -> ModelRunner:
    global _runner
    if _runner is None:
        _runner = ModelRunner()
    return _runner


def model_is_warm() -> bool:
    return _runner is not None


# Status strings for the visible UI. The cold path is genuinely slow -- the
# container has to wake, then torch/diffusers import, then ~9 s of weight
# loading before a single step runs -- so say so rather than showing a
# spinner that looks indistinguishable from a hang.
STATUS_IDLE = "Ready."
STATUS_COLD = (
    "⏳ **Starting server…** waking the GPU and loading the model. "
    "The first image takes around 40 seconds; later ones take a few."
)
STATUS_WARM = "⏳ **Generating…**"


def status_on_click() -> str:
    """Fires the instant the button is pressed, before the GPU job queues."""
    return STATUS_WARM if model_is_warm() else STATUS_COLD


def status_on_load() -> str:
    """Fires when someone opens the page."""
    if model_is_warm():
        return STATUS_IDLE
    return (
        "Model not loaded yet — the first generation also starts the server "
        "and takes around 40 seconds."
    )


def generate(prompt, negative_prompt, steps, guidance, lora_scale, seed,
             request: gr.Request = None):
    """One image per call. Returns (image, seed_used).

    Deliberately single-image: gradio_client handles a plain Image return
    far more predictably than a Gallery, and the UI can call again for more.

    `request` is supplied by Gradio from the type annotation and is not a UI
    input, so the published /generate signature stays at six arguments. It
    exists only to pick up the caller's trace headers.
    """
    try:
        # When the UI Space sent trace headers this run becomes a child of
        # its run, so one LangSmith trace covers UI and backend -- and the
        # visitor's location, carried in baggage, lands on this run too.
        with continue_trace(request):
            result = get_runner().generate(
                prompt=prompt,
                negative_prompt=negative_prompt,
                steps=int(steps),
                guidance=float(guidance),
                lora_scale=float(lora_scale),
                seed=None if int(seed) < 0 else int(seed),
                count=1,
            )
    except ValueError as error:
        raise gr.Error(str(error))
    except Exception as error:
        # The gr.Error string is all the caller sees; without this the
        # traceback is lost and the Space logs show nothing useful.
        logger.exception("request failed")
        raise gr.Error(f"Generation failed: {error}")
    return result.images[0], result.seed

with gr.Blocks(title="DiffusionDB SD 1.5 LoRA — inference API") as demo:
    gr.Markdown(
        f"## Inference API\n"
        f"Backend for [`diffusiondb-sd15-lora`]"
        f"(https://huggingface.co/spaces/whosouravsharma/diffusiondb-sd15-lora).\n\n"
        f"Model `{LORA_REPO}` @ `{LORA_CHECKPOINT}` · device `{pick_device()}` · "
        f"tracing `{'on' if TRACING_ACTIVE else 'off'}`\n\n"
        f"Call it with `gradio_client`:\n"
        f"```python\n"
        f"from gradio_client import Client\n"
        f"c = Client('whosouravsharma/diffusiondb-sd15-lora-inference')\n"
        f"image, seed = c.predict('a steampunk owl', '', 25, 7.5, 1.0, 42,\n"
        f"                        api_name='/generate')\n"
        f"```"
    )
    with gr.Row():
        with gr.Column():
            prompt = gr.Textbox(label="prompt", lines=2)
            negative = gr.Textbox(label="negative_prompt", lines=1)
            steps = gr.Number(value=25, label="steps", precision=0)
            guidance = gr.Number(value=7.5, label="guidance")
            lora_scale = gr.Number(value=1.0, label="lora_scale")
            seed = gr.Number(value=-1, label="seed", precision=0)
            run = gr.Button("Generate", variant="primary")
        with gr.Column():
            # Initial value assumes cold, because at import time it always
            # is; demo.load corrects it to "Ready." on a warm server. The
            # reverse default would flash a false "Ready." at every cold
            # visitor before the load event lands.
            status = gr.Markdown(status_on_load())
            image = gr.Image(label="image", type="pil")
            used_seed = gr.Number(label="seed_used", interactive=False)

    # Two independent listeners on the same button. The cosmetic one uses
    # queue=False so it paints immediately instead of waiting behind the GPU
    # job it is announcing. The API contract of /generate is untouched: same
    # six inputs, same two outputs, so gradio_client callers are unaffected.
    # api_name=False keeps these cosmetic events out of the published API,
    # so gradio_client sees /generate and nothing else.
    run.click(status_on_click, outputs=status, queue=False, api_name=False)

    generate_event = run.click(
        generate,
        inputs=[prompt, negative, steps, guidance, lora_scale, seed],
        outputs=[image, used_seed],
        api_name="generate",
    )
    generate_event.then(
        lambda: STATUS_IDLE, outputs=status, queue=False, api_name=False
    )

    demo.load(status_on_load, outputs=status, api_name=False)

if __name__ == "__main__":
    demo.queue(max_size=12).launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", 7860)),
    )