"""Restorer interface — "un-censor" detected regions of an image.

A Restorer takes an image plus the detected censored regions and returns a new
image with those regions reconstructed. This mirrors the ``Detector`` abstraction
so different DeepMosaics engines (per-frame and temporal/BVDNet; LADA later) plug
in behind the same interface. Note DeepMosaics locates the mosaic itself, so the
``detections`` argument is currently advisory (unused by the DeepMosaics engines).
"""

from __future__ import annotations

from abc import ABC, abstractmethod
from collections.abc import Callable

import numpy as np

from ..detection.types import Detection

# Optional cooperative-cancel hook: returns True to abort. Engines that loop or
# drive a subprocess should poll it; instant engines may ignore it.
CancelCheck = Callable[[], bool]


class Cancelled(Exception):
    """Raised by a Restorer when ``should_cancel`` asked it to stop."""


# Sequence (batch) restore callbacks — see ``Restorer.restore_sequence``.
FrameGetter = Callable[[int], np.ndarray]            # index -> BGR image
DetGetter = Callable[[int], list[Detection]]         # index -> that frame's detections
ResultSink = Callable[[int, np.ndarray], None]       # (index, restored image) -> None


class Restorer(ABC):
    #: Whether this engine uses *neighbouring* frames (so a batch run must feed it a
    #: contiguous, ordered sequence — see :meth:`restore_sequence`). Per-frame engines
    #: leave this False; the temporal DeepMosaics (BVDNet) sets it True.
    temporal: bool = False

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

    @abstractmethod
    def restore(
        self,
        image: np.ndarray,
        detections: list[Detection],
        should_cancel: CancelCheck | None = None,
    ) -> np.ndarray:
        """Return a copy of ``image`` with the detected regions reconstructed.

        ``should_cancel`` (if given) is polled periodically; when it returns
        True the engine should abort and raise :class:`Cancelled`.
        """
        raise NotImplementedError

    def restore_sequence(
        self,
        count: int,
        get_frame: FrameGetter,
        get_dets: DetGetter,
        emit: ResultSink,
        should_cancel: CancelCheck | None = None,
    ) -> None:
        """Restore ``count`` frames *in order*, calling ``emit(i, restored)`` for each.

        The default treats every frame independently (just loops :meth:`restore`).
        Temporal engines override this to pull neighbouring frames via ``get_frame``
        and carry recurrent state across the sequence. ``get_frame``/``get_dets`` are
        lazy so the engine only reads the frames it needs; ``emit`` lets the caller
        stream results to disk instead of holding them all in memory.
        """
        for i in range(count):
            if should_cancel is not None and should_cancel():
                raise Cancelled("Восстановление отменено")
            emit(i, self.restore(get_frame(i), get_dets(i), should_cancel))