"""Diffusion-inpaint restoration — redraw censored regions with a diffusion backend.

Unlike DeepMosaics (which *reconstructs* mosaic from its residual low-frequency data and
locates it itself), this engine *regenerates* the masked region with a diffusion inpaint
model: it builds a mask from the YOLO detections and hands ``(image, mask, params)`` to a
pluggable :class:`DiffusionBackend` (SwarmUI is the first, see ``swarmui.py``).

Consequences of that design:
  - It **needs detections** (``needs_detections = True``) — a frame with none comes back
    unchanged (no mask → nothing to regenerate). The caller feeds it the real detections.
  - It's **per-frame** (``temporal = False``): each frame is generated independently, so a
    video sequence will flicker. Best for stills / single frames, not coherent clips.
  - The backend runs in a **separate process/server** (e.g. SwarmUI over HTTP), so this
    path adds **no torch dependency** to the app and keeps the heavy model out-of-process.

The backend is abstract so other diffusion servers (ComfyUI/A1111) can be added later as
another :class:`DiffusionBackend`, without touching the restorer or the UI.
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from dataclasses import dataclass

import numpy as np

from ..detection.types import Detection
from .base import CancelCheck, Cancelled, Restorer
from .mask import detections_to_mask, mask_is_empty


@dataclass(frozen=True)
class InpaintParams:
    """Generation knobs handed to a :class:`DiffusionBackend`."""

    prompt: str = ""
    negative: str = ""
    model: str | None = None       # checkpoint name as the backend knows it (None = current)
    steps: int = 30
    cfg: float = 7.0
    denoise: float = 1.0           # 0..1 — how much to regenerate under the mask (1 = full)
    seed: int = -1                 # -1 = random each call
    mask_blur: int = 8             # px feather applied by the backend at its mask edge


class DiffusionBackend(ABC):
    """A diffusion inpaint engine reachable from our process (typically over HTTP)."""

    @property
    def name(self) -> str:
        return type(self).__name__

    @abstractmethod
    def inpaint(
        self,
        image_bgr: np.ndarray,
        mask: np.ndarray,
        params: InpaintParams,
        should_cancel: CancelCheck | None = None,
    ) -> np.ndarray:
        """Regenerate the white area of ``mask`` in ``image_bgr``; return a new BGR image."""
        raise NotImplementedError


class DiffusionRestorer(Restorer):
    """Restorer that masks the detected regions and inpaints them via a backend."""

    temporal = False
    needs_detections = True

    def __init__(
        self,
        backend: DiffusionBackend,
        params: InpaintParams,
        *,
        mask_dilate: int = 4,
        mask_blur: int = 8,
    ) -> None:
        self._backend = backend
        self._params = params
        self._dilate = mask_dilate
        self._blur = mask_blur

    @property
    def name(self) -> str:
        return f"Diffusion({self._backend.name})"

    def restore(
        self,
        image: np.ndarray,
        detections: list[Detection],
        should_cancel: CancelCheck | None = None,
    ) -> np.ndarray:
        if should_cancel is not None and should_cancel():
            raise Cancelled("Восстановление отменено")
        if not detections:
            return image.copy()  # no detections → no mask → nothing to regenerate
        mask = detections_to_mask(
            detections, image.shape, dilate=self._dilate, blur=self._blur
        )
        if mask_is_empty(mask):
            return image.copy()
        return self._backend.inpaint(image, mask, self._params, should_cancel)