"""Configure the restoration ("расцензурить") engine.

Restoration is DeepMosaics-only — pick the per-frame ("картинка") or temporal ("видео")
engine. The network code is vendored (built-in); the user picks a clean model from a
dropdown of the bundled ``models/deepmosaics`` weights (or browses to another
``clean_*.pth``) — for the video engine only ``clean_*_video.pth`` is offered.
``mosaic_position.pth`` must sit beside the chosen model. A CUDA GPU is strongly
recommended (GPU id, -1 = CPU/slow).
"""

from __future__ import annotations

from pathlib import Path

from PySide6.QtWidgets import (
    QCheckBox,
    QComboBox,
    QDialog,
    QDialogButtonBox,
    QFileDialog,
    QFormLayout,
    QHBoxLayout,
    QLabel,
    QLineEdit,
    QPushButton,
    QWidget,
)

from ..config import AppConfig
from ..core.restore.deepmosaics import discover_models


class RestoreDialog(QDialog):
    def __init__(self, config: AppConfig, parent: QWidget | None = None) -> None:
        super().__init__(parent)
        self._cfg = config
        self.setWindowTitle("Движок восстановления")
        self.setMinimumWidth(560)

        self.engine = QComboBox()
        self.engine.addItem("DeepMosaics — картинка (покадрово, нужна модель+GPU)", "deepmosaics")
        self.engine.addItem("DeepMosaics — видео (соседние кадры, лучше для роликов)", "deepmosaics_video")
        self.engine.setCurrentIndex(max(0, self.engine.findData(config.restorer)))
        self.engine.currentIndexChanged.connect(self._sync)

        # Model dropdown — bundled clean models, plus the configured one if external.
        self.model_combo = QComboBox()
        self._populate_models(config.dm_model)

        self.dm_gpu = QLineEdit(config.dm_gpu or "0")
        self.dm_gpu.setPlaceholderText("0 = первая CUDA-карта, -1 = CPU (медленно)")

        # Video engine only: feed already-restored past frames into the temporal window.
        self.feed_restored = QCheckBox(
            "Подавать уже расцензуренные прошлые кадры в окно (эксперим.)"
        )
        self.feed_restored.setChecked(bool(getattr(config, "dm_feed_restored", True)))
        self.feed_restored.setToolTip(
            "Только для видеодвижка: прошлые соседние кадры в окне берутся из уже\n"
            "восстановленных результатов, а не из оригинала с мозаикой — больше\n"
            "временной связности. Сеть обучалась на мозаичных окнах, так что эффект\n"
            "не гарантирован; выключите для точной реализации DeepMosaics."
        )

        form = QFormLayout(self)
        form.addRow("Движок:", self.engine)
        form.addRow("Модель:", self._with_browse(self.model_combo, self._browse_model))
        form.addRow("GPU id:", self.dm_gpu)
        form.addRow("", self.feed_restored)
        hint = QLabel(
            "Модели берутся из models/deepmosaics (рядом нужен mosaic_position.pth).\n"
            "• Картинка: clean_youknow_resnet_9blocks.pth — покадрово.\n"
            "• Видео: clean_youknow_video.pth — использует соседние кадры (когерентнее на "
            "роликах), требует прогона по диапазону («Расцензурить все»).\n"
            "На CPU медленно — лучше GPU. См. README."
        )
        hint.setWordWrap(True)
        form.addRow(hint)

        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        form.addRow(buttons)
        self._sync()

    def _populate_models(self, current: str | None) -> None:
        self.model_combo.clear()
        is_video = self.engine.currentData() == "deepmosaics_video"
        if is_video:  # temporal engine: only the video weights (clean_*_video.pth)
            models = [(n, p) for n, p in discover_models(include_video=True) if "video" in n.lower()]
        else:
            models = discover_models()  # per-frame engine: image clean models only
        for name, path in models:
            self.model_combo.addItem(name, path)
        # Keep an externally-configured model selectable even if it's outside the folder.
        if current and self.model_combo.findData(current) < 0:
            self.model_combo.addItem(Path(current).stem + " (внешняя)", current)
        if self.model_combo.count() == 0:
            self.model_combo.addItem("(модели не найдены — положите в models/deepmosaics)", None)
        idx = self.model_combo.findData(current) if current else 0
        self.model_combo.setCurrentIndex(max(0, idx))

    def _with_browse(self, widget: QWidget, slot) -> QWidget:
        w = QWidget()
        h = QHBoxLayout(w)
        h.setContentsMargins(0, 0, 0, 0)
        h.addWidget(widget, 1)
        btn = QPushButton("Обзор…")
        btn.clicked.connect(slot)
        h.addWidget(btn)
        return w

    def _sync(self) -> None:
        is_dm = self.engine.currentData() in ("deepmosaics", "deepmosaics_video")
        is_video = self.engine.currentData() == "deepmosaics_video"
        # The model list differs per engine (image vs video weights) — repopulate.
        self._populate_models(self._cfg.dm_model)
        self.model_combo.setEnabled(is_dm)
        self.dm_gpu.setEnabled(is_dm)
        self.feed_restored.setEnabled(is_video)  # only the temporal engine has a window

    def _browse_model(self) -> None:
        p, _ = QFileDialog.getOpenFileName(self, "Веса DeepMosaics", "", "Веса (*.pth);;Все файлы (*.*)")
        if p:
            if self.model_combo.findData(p) < 0:
                self.model_combo.addItem(Path(p).stem, p)
            self.model_combo.setCurrentIndex(self.model_combo.findData(p))

    def apply_to_config(self) -> None:
        self._cfg.restorer = self.engine.currentData()
        self._cfg.dm_model = self.model_combo.currentData()
        self._cfg.dm_gpu = self.dm_gpu.text().strip() or "0"
        self._cfg.dm_feed_restored = self.feed_restored.isChecked()