"""Gradio demo: PCB defect + golden reference -> enriched image (A11_CA prebackbone)."""

from __future__ import annotations

import os
from pathlib import Path

import gradio as gr
import numpy as np
from PIL import Image

from prebackbone_infer import enrich_pair, get_enricher

EXAMPLES_DIR = Path(__file__).resolve().parent / "examples"
TITLE = "RefDiffNet — PCB Reference–Defect Enrichment"
DESCRIPTION = """
Upload a **defect PCB image** and its **golden reference** (same board, no defect).
The trained A11_CA prebackbone fuses them into an **enriched** image used by YOLO detection.

Formula: `enriched = defect + α · gate · delta`
"""


def _examples() -> list[list[str]]:
    if not EXAMPLES_DIR.exists():
        return []
    inputs = sorted(EXAMPLES_DIR.glob("*_input.jpg"))
    out: list[list[str]] = []
    for inp in inputs:
        ref = EXAMPLES_DIR / inp.name.replace("_input.", "_reference.")
        if ref.exists():
            out.append([str(inp), str(ref)])
    return out


def run_demo(defect_img, reference_img):
    if defect_img is None or reference_img is None:
        raise gr.Error("Please upload both defect and reference images.")

    defect_rgb, ref_rgb, enriched_rgb = enrich_pair(defect_img, reference_img)

    # Side-by-side comparison for the demo panel
    h = max(defect_rgb.shape[0], ref_rgb.shape[0], enriched_rgb.shape[0])

    def _pad(im: np.ndarray) -> np.ndarray:
        if im.shape[0] == h:
            return im
        pad = h - im.shape[0]
        return np.pad(im, ((0, pad), (0, 0), (0, 0)), mode="constant", constant_values=114)

    row = np.concatenate([_pad(defect_rgb), _pad(ref_rgb), _pad(enriched_rgb)], axis=1)
    return enriched_rgb, row


def _preload_model():
    try:
        get_enricher()
    except FileNotFoundError as e:
        print(f"[warn] Model not loaded yet: {e}")


if __name__ == "__main__":
    _preload_model()

    with gr.Blocks(title=TITLE) as demo:
        gr.Markdown(f"# {TITLE}")
        gr.Markdown(DESCRIPTION)

        with gr.Row():
            defect_in = gr.Image(label="Defect image (input)", type="numpy", image_mode="RGB")
            ref_in = gr.Image(label="Golden reference", type="numpy", image_mode="RGB")

        run_btn = gr.Button("Generate enriched image", variant="primary")
        with gr.Row():
            enriched_out = gr.Image(label="Enriched output", type="numpy", image_mode="RGB")
            compare_out = gr.Image(
                label="Defect | Reference | Enriched",
                type="numpy",
                image_mode="RGB",
            )

        run_btn.click(fn=run_demo, inputs=[defect_in, ref_in], outputs=[enriched_out, compare_out])

        ex = _examples()
        if ex:
            gr.Examples(examples=ex, inputs=[defect_in, ref_in], label="Example pairs")

    demo.launch(
        server_name="0.0.0.0",
        server_port=int(os.environ.get("PORT", "7860")),
        share=os.environ.get("GRADIO_SHARE", "").lower() in {"1", "true", "yes"},
    )