"""Probe the diffusion-inpaint path against a LIVE SwarmUI server (no GUI).

Exercises exactly what the app does — builds an inpaint mask from a detection and runs
``DiffusionRestorer`` over a real SwarmUI backend — so you can verify the HTTP/API
plumbing (field names drift between SwarmUI versions) independently of the GUI.

Prereqs: a running SwarmUI server with an inpaint-capable checkpoint loaded.

Run (PowerShell)::

    .venv\\Scripts\\python.exe scripts\\swarmui_probe.py `
        --url http://localhost:7801 `
        --image "D:\\path\\to\\frame.jpg" `
        --out "D:\\path\\to\\out.png" `
        --prompt "clean skin" --steps 25 --denoise 1.0

With no --image a 512x512 synthetic frame is used (a grey box censored in the centre).
Exits non-zero on any error and prints the server's reply on failure.
"""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import numpy as np

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from hvideotool.core.detection.types import CensorType, Detection  # noqa: E402
from hvideotool.core.imageio import imread_unicode, imwrite_unicode  # noqa: E402
from hvideotool.core.restore.diffusion import DiffusionRestorer, InpaintParams  # noqa: E402
from hvideotool.core.restore.swarmui import SwarmUIBackend  # noqa: E402


def _synthetic() -> tuple[np.ndarray, Detection]:
    """A 512x512 frame with a 'censored' grey block in the middle + a matching detection."""
    img = np.full((512, 512, 3), 60, np.uint8)
    img[40:472, 40:472] = (120, 90, 70)            # some background
    img[180:332, 180:332] = 128                     # the "censored" block
    det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(180, 180, 152, 152), label="mosaic")
    return img, det


def main() -> int:
    ap = argparse.ArgumentParser(description="Probe diffusion-inpaint against a live SwarmUI")
    ap.add_argument("--url", default="http://localhost:7801")
    ap.add_argument("--image", default=None, help="frame to inpaint (default: synthetic)")
    ap.add_argument("--out", default="swarmui_probe_out.png")
    ap.add_argument("--model", default=None, help="checkpoint name as SwarmUI knows it")
    ap.add_argument("--prompt", default="")
    ap.add_argument("--negative", default="")
    ap.add_argument("--steps", type=int, default=25)
    ap.add_argument("--cfg", type=float, default=7.0)
    ap.add_argument("--denoise", type=float, default=1.0)
    ap.add_argument("--seed", type=int, default=-1)
    ap.add_argument("--dilate", type=int, default=4)
    ap.add_argument("--blur", type=int, default=8)
    args = ap.parse_args()

    if args.image:
        img = imread_unicode(args.image)
        if img is None:
            print(f"Не удалось прочитать {args.image}", file=sys.stderr)
            return 2
        h, w = img.shape[:2]
        # No detector here — mask the central third so there's something to inpaint.
        bx, by = w // 3, h // 3
        det = Detection(type=CensorType.MOSAIC, score=0.99, bbox=(bx, by, w // 3, h // 3), label="mosaic")
    else:
        img, det = _synthetic()

    backend = SwarmUIBackend(args.url)
    params = InpaintParams(
        prompt=args.prompt, negative=args.negative, model=args.model,
        steps=args.steps, cfg=args.cfg, denoise=args.denoise, seed=args.seed,
    )
    restorer = DiffusionRestorer(backend, params, mask_dilate=args.dilate, mask_blur=args.blur)

    print(f"→ SwarmUI {args.url}  frame={img.shape[1]}x{img.shape[0]}  steps={args.steps} denoise={args.denoise}")
    try:
        out = restorer.restore(img, [det])
    except Exception as e:  # noqa: BLE001 — surface the server error verbatim
        print(f"ОШИБКА: {e}", file=sys.stderr)
        return 1

    if not imwrite_unicode(args.out, out):
        print(f"Не удалось записать {args.out}", file=sys.stderr)
        return 3
    changed = int(np.count_nonzero(np.any(out.astype(int) - img.astype(int) != 0, axis=2)))
    print(f"✓ Готово → {args.out}  (изменено пикселей: {changed})")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())