"""Dialog for selecting and downloading SAM2 ONNX model weights."""

from __future__ import annotations

import urllib.request
from pathlib import Path

import structlog

from qgis.PyQt.QtCore import QThread, pyqtSignal
from qgis.PyQt.QtWidgets import (
    QComboBox,
    QDialog,
    QDialogButtonBox,
    QLabel,
    QMessageBox,
    QProgressDialog,
    QVBoxLayout,
    QWidget,
)

from .model_registry import (
    MODELS,
    DEFAULT_MODEL,
    get_model_dir,
    is_model_downloaded,
    verify_download,
)

logger = structlog.get_logger(__name__)


class _DownloadWorker(QThread):
    """Background thread that downloads model weight files."""

    progress = pyqtSignal(int, str)  # (overall_pct, label_text)
    finished = pyqtSignal(bool, str)  # (success, error_message)

    _CANCEL_MSG = "Download cancelled by user"

    def __init__(
        self,
        files: list[tuple[str, str, Path]],
        parent: QWidget | None = None,
        request_id: str | None = None,
    ) -> None:
        """files: list of (url, local_name, target_path) tuples."""
        super().__init__(parent)
        self._files = files
        self._request_id = request_id

    def run(self) -> None:
        if self._request_id:
            from samgis_core.utilities.request_context import bind_request_id

            bind_request_id(self._request_id)
        total = len(self._files)
        try:
            for idx, (url, local_name, target_path) in enumerate(self._files):
                if self.isInterruptionRequested():
                    self.finished.emit(False, self._CANCEL_MSG)
                    return

                self.progress.emit(
                    idx * 100 // total,
                    f"Downloading {local_name} ({idx + 1}/{total})...",
                )

                def reporthook(
                    block_num: int, block_size: int, total_size: int, _idx: int = idx
                ) -> None:
                    if self.isInterruptionRequested():
                        raise InterruptedError("cancelled")
                    if total_size > 0:
                        file_pct = block_num * block_size * 100 // total_size
                    else:
                        file_pct = 0
                    overall = (_idx * 100 + file_pct) // total
                    self.progress.emit(min(overall, 99), "")

                logger.info(f"Downloading {url} → {target_path}")
                urllib.request.urlretrieve(url, str(target_path), reporthook)  # noqa: S310

            self.progress.emit(100, "")
            self.finished.emit(True, "")

        except InterruptedError:
            self.finished.emit(False, self._CANCEL_MSG)
        except Exception as e:
            logger.exception(f"Download failed: {e}")
            self.finished.emit(False, str(e))


class ModelDownloadDialog(QDialog):
    """Dialog to select a SAM2 model variant and download weights."""

    def __init__(
        self,
        parent: QWidget | None = None,
        active_variant: str = DEFAULT_MODEL,
        request_id: str | None = None,
    ) -> None:
        super().__init__(parent)
        self.setWindowTitle("SAM2 Model")
        self.setMinimumWidth(350)
        self._worker: _DownloadWorker | None = None
        self._request_id = request_id

        layout = QVBoxLayout()

        layout.addWidget(QLabel("Select a model variant:"))

        # Model variant selector — mark downloaded variants
        self._variant_combo = QComboBox()
        active_idx = 0
        for idx, (key, info) in enumerate(MODELS.items()):
            suffix = " (downloaded)" if is_model_downloaded(key) else ""
            self._variant_combo.addItem(info["label"] + suffix, key)
            if key == active_variant:
                active_idx = idx
        self._variant_combo.setCurrentIndex(active_idx)
        layout.addWidget(self._variant_combo)

        # Size info
        self._size_label = QLabel()
        self._update_size_label()
        self._variant_combo.currentIndexChanged.connect(self._update_size_label)
        layout.addWidget(self._size_label)

        # Buttons
        button_box = QDialogButtonBox()
        self._action_btn = button_box.addButton(
            "Select",
            QDialogButtonBox.ButtonRole.AcceptRole,
        )
        button_box.addButton(QDialogButtonBox.StandardButton.Cancel)
        self._action_btn.clicked.connect(self._on_action)
        self._variant_combo.currentIndexChanged.connect(self._update_action_btn)
        button_box.rejected.connect(self._on_reject)
        layout.addWidget(button_box)
        self._update_action_btn()

        self.setLayout(layout)

    def selected_variant(self) -> str:
        return self._variant_combo.currentData()

    def _update_size_label(self) -> None:
        variant = self.selected_variant()
        size = MODELS[variant]["size_mb"]
        if is_model_downloaded(variant):
            self._size_label.setText("Model already downloaded.")
        else:
            self._size_label.setText(f"Estimated download size: ~{size} MB")

    def _update_action_btn(self, _index: int = 0) -> None:
        variant = self.selected_variant()
        if is_model_downloaded(variant):
            self._action_btn.setText("Select")
        else:
            self._action_btn.setText("Download")

    def _on_reject(self) -> None:
        if self._worker and self._worker.isRunning():
            self._worker.requestInterruption()
            self._worker.wait()
        self.reject()

    def _on_action(self) -> None:
        """Select if downloaded, download otherwise."""
        variant = self.selected_variant()
        if is_model_downloaded(variant):
            self.accept()
            return
        self._on_download()

    def _on_download(self) -> None:
        variant = self.selected_variant()
        model_info = MODELS[variant]
        model_dir = get_model_dir(variant)
        model_dir.mkdir(parents=True, exist_ok=True)

        # Build download list: (url, local_name, target_path)
        files = [
            (model_info["urls"][remote], local, model_dir / local)
            for remote, local in model_info["files"].items()
        ]

        self._progress = QProgressDialog(
            "Downloading model weights...",
            "Cancel",
            0,
            100,
            self,
        )
        self._progress.setWindowTitle("Downloading SAM2 Model")
        self._progress.setMinimumDuration(0)
        self._progress.setValue(0)

        self._worker = _DownloadWorker(files, self, request_id=self._request_id)
        self._worker.progress.connect(self._on_progress)
        self._worker.finished.connect(self._on_worker_finished)
        self._progress.canceled.connect(self._worker.requestInterruption)

        self._worker.start()

    def _on_progress(self, pct: int, label: str) -> None:
        self._progress.setValue(pct)
        if label:
            self._progress.setLabelText(label)

    def _on_worker_finished(self, success: bool, error: str) -> None:
        self._progress.close()
        self._worker = None

        if success:
            variant = self.selected_variant()
            failures = verify_download(variant)
            if failures:
                QMessageBox.critical(
                    self,
                    "Checksum Verification Failed",
                    f"SHA-256 mismatch for: {', '.join(failures)}\n"
                    "The downloaded files may be corrupted. Please retry.",
                )
                return
            logger.info(f"Model {variant} downloaded and verified")
            self.accept()
        elif error and error != _DownloadWorker._CANCEL_MSG:
            QMessageBox.critical(
                self,
                "Download Failed",
                f"Failed to download model weights:\n{error}",
            )
