"""DeepMosaics restorer — real generative mosaic removal.

Rather than vendoring DeepMosaics' GPL network code (which must match the exact
checkpoint), we drive a **user-installed** DeepMosaics (https://github.com/HypoX64/DeepMosaics)
as a subprocess: write the frame to a temp file, run ``deepmosaic.py --mode clean``,
read the cleaned image back. This reuses their tested pipeline (incl. their own
mosaic locator ``mosaic_position.pth``) and respects the GPL boundary.

Setup the user must do once (see README → Восстановление):
1. ``git clone https://github.com/HypoX64/DeepMosaics`` and install its deps.
2. Download clean weights (e.g. ``clean_youknow_video.pth``) AND ``mosaic_position.pth``
   into one folder.
3. In the app: Восстановление… → engine "deepmosaics", set the DeepMosaics folder
   and the clean-model path (a CUDA GPU is strongly recommended).

NOTE: DeepMosaics finds the mosaic itself; our detections are used for navigation,
not passed to it.
"""

from __future__ import annotations

import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np

from ..detection.types import Detection
from ..imageio import imread_unicode, imwrite_unicode
from .base import Restorer

_IMG_EXTS = {".jpg", ".jpeg", ".png", ".bmp"}


class DeepMosaicsRestorer(Restorer):
    def __init__(
        self,
        deepmosaics_dir: str | None,
        model_path: str | None,
        python_exe: str | None = None,
        gpu_id: str = "0",
    ) -> None:
        if not deepmosaics_dir or not (Path(deepmosaics_dir) / "deepmosaic.py").is_file():
            raise ValueError(
                "Не указана папка DeepMosaics (с deepmosaic.py).\n"
                "Установите DeepMosaics и укажите её в «Восстановление…». См. README."
            )
        if not model_path or not Path(model_path).is_file():
            raise ValueError(
                "Не найдены веса DeepMosaics (clean_*.pth).\n"
                "Скачайте clean_youknow_video.pth + mosaic_position.pth в одну папку. См. README."
            )
        self._dir = Path(deepmosaics_dir)
        self._model = model_path
        self._python = python_exe or sys.executable
        self._gpu = gpu_id

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

    def restore(self, image: np.ndarray, detections: list[Detection]) -> np.ndarray:
        with tempfile.TemporaryDirectory(prefix="hvt_dm_") as tmp:
            tmpd = Path(tmp)
            src = tmpd / "frame.jpg"
            result_dir = tmpd / "result"
            result_dir.mkdir()
            imwrite_unicode(str(src), image)

            cmd = [
                self._python, "deepmosaic.py",
                "--media_path", str(src),
                "--model_path", str(self._model),
                "--mode", "clean",
                "--result_dir", str(result_dir),
                "--temp_dir", str(tmpd / "dmtmp"),
                "--gpu_id", str(self._gpu),
                "--no_preview",
            ]
            proc = subprocess.run(
                cmd, cwd=str(self._dir),
                stdin=subprocess.DEVNULL,            # so DeepMosaics' error input() can't hang
                capture_output=True, text=True,
            )
            outputs = [p for p in result_dir.iterdir() if p.suffix.lower() in _IMG_EXTS]
            if outputs:
                newest = max(outputs, key=lambda p: p.stat().st_mtime)
                restored = imread_unicode(str(newest))
                if restored is None:
                    raise RuntimeError("Не удалось прочитать результат DeepMosaics.")
                return restored

            # No output file — figure out why.
            log = (proc.stderr or "") + (proc.stdout or "")
            if "BVDNet.forward()" in log or "argument: 'previous'" in log:
                raise RuntimeError(
                    "Видеомодель (clean_*_video.pth) не работает покадрово — ей нужен "
                    "соседний кадр.\nУкажите картиночную модель clean_youknow_resnet_9blocks.pth."
                )
            if proc.returncode == 0:
                # DeepMosaics ran fine but found no mosaic to clean — keep the frame as is.
                return image.copy()
            tail = log.strip().splitlines()[-6:]
            raise RuntimeError(
                f"DeepMosaics не вернул результат (код {proc.returncode}).\n" + "\n".join(tail)
            )