"""DeepMosaics restorer — real generative mosaic removal, in-process.

The DeepMosaics network code (GPL-3.0) is vendored under ``_deepmosaics/`` (see
its NOTICE/LICENSE). We load the models **once** and run the per-frame clean path
in-process — far faster than spawning a subprocess per frame (which reloaded the
models every time). Only the model *weights* are user-supplied.

Per-frame clean = DeepMosaics' ``cleanmosaic_img_server`` logic, reimplemented
here (so we don't pull in their video/ffmpeg modules):
  locate mosaic (BiSeNet ``mosaic_position.pth``) → run the clean generator on the
  crop → feather it back. DeepMosaics finds the mosaic itself; our detections are
  used for navigation, not passed to it.

Setup (see README → Восстановление): download the **image** clean weights
``clean_youknow_resnet_9blocks.pth`` + ``mosaic_position.pth`` into one folder and
point the app at the clean-model file. The video model ``clean_youknow_video.pth``
(BVDNet) needs neighbour frames and does NOT work per-frame.
"""

from __future__ import annotations

import sys
from pathlib import Path
from types import SimpleNamespace

import numpy as np

from ..detection.types import Detection
from .base import CancelCheck, Cancelled, Restorer

_VENDOR = Path(__file__).parent / "_deepmosaics"
# Default place to drop DeepMosaics clean weights (gitignored — see models/).
DEFAULT_WEIGHTS_DIR = Path(__file__).resolve().parents[3] / "models" / "deepmosaics"


def discover_models(extra_dir: str | None = None) -> list[tuple[str, str]]:
    """Find usable per-frame clean models: (display_name, full_path).

    Scans the bundled ``models/deepmosaics`` folder (plus ``extra_dir`` if given)
    for ``clean_*.pth``. The video model is skipped — it can't run per-frame.
    """
    dirs = [DEFAULT_WEIGHTS_DIR]
    if extra_dir:
        dirs.insert(0, Path(extra_dir))
    out: list[tuple[str, str]] = []
    seen: set[str] = set()
    for d in dirs:
        if not d.is_dir():
            continue
        for p in sorted(d.glob("clean_*.pth")):
            if "video" in p.name.lower() or p.name in seen:
                continue
            seen.add(p.name)
            out.append((p.stem, str(p)))
    return out


def _netg_kind(model_name: str) -> str:
    """Pick DeepMosaics' netG type from the weights filename (see their options.py)."""
    n = model_name.lower()
    if "video" in n:
        raise ValueError(
            "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен соседний "
            "кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
        )
    if "unet_128" in n:
        return "unet_128"
    if "hd" in n:
        return "HD"
    return "resnet_9blocks"


class DeepMosaicsRestorer(Restorer):
    def __init__(
        self,
        deepmosaics_dir: str | None,   # kept for factory/config compatibility (weights hint)
        model_path: str | None,
        python_exe: str | None = None,  # unused now (in-process)
        gpu_id: str = "0",
    ) -> None:
        if not model_path or not Path(model_path).is_file():
            discovered = discover_models()  # fall back to a bundled model
            if discovered:
                model_path = discovered[0][1]
            else:
                raise ValueError(
                    "Не найдены веса DeepMosaics (clean_*.pth).\n"
                    "Положите clean_youknow_resnet_9blocks.pth + mosaic_position.pth в "
                    "models/deepmosaics (или выберите в «Восстановление…»). См. README."
                )
        model = Path(model_path)
        self._netg = _netg_kind(model.name)  # raises on a video model
        pos = self._find_mosaic_position(model, deepmosaics_dir)
        if pos is None:
            raise ValueError(
                "Рядом с clean-моделью не найден mosaic_position.pth.\n"
                "Положите mosaic_position.pth в ту же папку, что и clean_*.pth. См. README."
            )
        self._model = str(model)
        self._pos = str(pos)
        self._gpu = gpu_id
        self._loaded = False  # models loaded lazily on first restore

    @staticmethod
    def _find_mosaic_position(model: Path, dm_dir: str | None) -> Path | None:
        candidates = [model.parent / "mosaic_position.pth"]
        if dm_dir:
            candidates.append(Path(dm_dir) / "pretrained_models" / "mosaic" / "mosaic_position.pth")
        return next((p for p in candidates if p.is_file()), None)

    @property
    def name(self) -> str:
        return f"DeepMosaics(gpu={self._gpu})"

    # ------------------------------------------------------------------ engine
    def _ensure_loaded(self) -> None:
        if self._loaded:
            return
        if str(_VENDOR) not in sys.path:
            sys.path.insert(0, str(_VENDOR))  # so vendored `from models/util import …` resolve
        from models import loadmodel, runmodel  # type: ignore  # noqa: E402
        import util.image_processing as impro    # type: ignore  # noqa: E402

        self._runmodel = runmodel
        self._impro = impro
        self._opt = SimpleNamespace(
            gpu_id=self._gpu,
            netG=self._netg,
            model_path=self._model,
            mosaic_position_model_path=self._pos,
            mask_threshold=64,
            all_mosaic_area=False,
            ex_mult=1.5,
            no_feather=False,
            traditional=False,
        )
        self._netM = loadmodel.bisenet(self._opt, "mosaic")
        self._netG = loadmodel.pix2pix(self._opt)
        self._loaded = True

    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("Восстановление отменено")
        self._ensure_loaded()
        rm, impro, opt = self._runmodel, self._impro, self._opt

        # DeepMosaics' cleanmosaic_img_server, faithfully reproduced.
        x, y, size, mask = rm.get_mosaic_position(image, self._netM, opt)
        if size <= 100:
            return image.copy()  # no mosaic located — leave the frame untouched
        if should_cancel is not None and should_cancel():
            raise Cancelled("Восстановление отменено")
        work = image.copy()
        img_mosaic = work[y - size:y + size, x - size:x + size]
        img_fake = rm.run_pix2pix(img_mosaic, self._netG, opt)
        return impro.replace_mosaic(work, img_fake, mask, x, y, size, opt.no_feather)