"""
Modified from https://courses.spatialthoughts.com/pyqgis-masterclass.html#
Original credits to "Customizing QGIS with Python Course" by Ujaval Gandhi
www.spatialthoughts.com
"""

from __future__ import annotations

import os
import tempfile
from pathlib import Path
from typing import TYPE_CHECKING

import structlog
from qgis.PyQt.QtGui import QIcon
from qgis.PyQt.QtWidgets import QDialog, QLabel, QMessageBox, QProgressDialog

if TYPE_CHECKING:
    from qgis.gui import QgisInterface
    from samgis_core.prediction_api.sam2_adapter import Sam2OnnxPredictor

    from .tasks import Sam2InferenceTask
from qgis.PyQt.QtCore import Qt
from qgis.PyQt.QtGui import QAction
from qgis.core import (
    QgsApplication,
    QgsRasterLayer,
    QgsVectorLayer,
    QgsProject,
    QgsWkbTypes,
)
from samgis_core.utilities.request_context import bind_request_id

from .vector_dialog import VectorPromptDialog
from .adapters import qgis_raster, qgis_vector

logger = structlog.get_logger(__name__)


plugin_dir = Path(__file__).parent
public_dir = plugin_dir / "public"
deps_dir = str(plugin_dir / "deps")


class SamgisPlugin:
    def __init__(self, iface: QgisInterface) -> None:
        self.samgis_toolbar = None
        self.label = None
        self.vector_prompt_action = None
        self.iface = iface

        # Lazy-init SAM model (loaded on first use to avoid slow plugin startup)
        self._sam_model: Sam2OnnxPredictor | None = None
        self._request_id = None
        from .model_registry import load_active_variant

        self._active_variant = load_active_variant()

    def initGui(self) -> None:
        # create a toolbar
        self.samgis_toolbar = self.iface.addToolBar("SamGIS selector Toolbar")

        # Icons
        logo_icon = QIcon(str(public_dir / "icon.svg"))
        settings_icon = QIcon(str(public_dir / "settings.svg"))

        # create a label
        self.label = QLabel("SamGIS:", self.samgis_toolbar)
        self.samgis_toolbar.addWidget(self.label)

        # Create local SAM inference action
        self.vector_prompt_action = QAction(logo_icon, "Local SAM", self.samgis_toolbar)
        self.samgis_toolbar.addAction(self.vector_prompt_action)
        self.vector_prompt_action.triggered.connect(self.show_vector_prompt_dialog)

        # Change model action
        self.change_model_action = QAction(
            settings_icon, "Change Model", self.samgis_toolbar
        )
        self._update_model_tooltip()
        self.samgis_toolbar.addAction(self.change_model_action)
        self.change_model_action.triggered.connect(self.show_change_model_dialog)

    def unload(self) -> None:
        del self.samgis_toolbar

    def _update_model_tooltip(self) -> None:
        from .model_registry import MODELS

        label = MODELS.get(self._active_variant, {}).get("label", self._active_variant)
        self.change_model_action.setToolTip(f"Active model: {label}")

    def _get_model_dir(self) -> Path | None:
        """Return model directory, triggering download dialog if needed."""
        from .model_registry import (
            is_model_downloaded,
            get_model_dir,
            save_active_variant,
        )

        variant = self._active_variant
        if is_model_downloaded(variant):
            return get_model_dir(variant)

        from .model_dialog import ModelDownloadDialog

        dialog = ModelDownloadDialog(
            self.iface.mainWindow(),
            active_variant=variant,
            request_id=self._request_id,
        )
        if dialog.exec() == QDialog.DialogCode.Accepted:
            variant = dialog.selected_variant()
            self._active_variant = variant
            save_active_variant(variant)
            self._update_model_tooltip()
            return get_model_dir(variant)

        self.iface.messageBar().pushWarning(
            "SAM",
            "Model download cancelled. Cannot run inference.",
        )
        return None

    def _ensure_onnxruntime(self) -> bool:
        """Check onnxruntime availability; offer non-blocking pip install if missing."""
        from .utility import try_import_module

        if try_import_module("onnxruntime") is not None:
            logger.debug("onnxruntime already available")
            return True

        answer = QMessageBox.question(
            self.iface.mainWindow(),
            "Missing dependency",
            "onnxruntime is required for inference but not installed.\n"
            "Install it now? This may take some minutes...",
            QMessageBox.StandardButton.Yes | QMessageBox.StandardButton.No,  # pyright: ignore[reportArgumentType] — PyQt stubs return int for scoped enum |
        )
        if answer != QMessageBox.StandardButton.Yes:
            self.iface.messageBar().pushInfo(
                "SAM",
                "onnxruntime is needed for inference. Click Local SAM to retry.",
            )
            return False

        from .utility import PipInstallWorker

        progress = QProgressDialog(
            "Installing onnxruntime...",
            None,
            0,
            0,
            self.iface.mainWindow(),
        )
        progress.setWindowModality(Qt.WindowModality.WindowModal)
        progress.setCancelButton(None)
        progress.show()

        worker = PipInstallWorker(
            "onnxruntime", target_dir=deps_dir, request_id=self._request_id
        )
        result = {}

        def on_finished(success: bool, msg: str) -> None:
            result["success"] = success
            result["msg"] = msg
            progress.close()

        worker.finished.connect(on_finished)
        worker.start()
        # Process events while worker runs (modal but responsive)
        while worker.isRunning():
            QgsApplication.processEvents()
        worker.wait()

        if result.get("success"):
            logger.info(f"onnxruntime installed successfully to {deps_dir}")
            self.iface.messageBar().pushSuccess(
                "SAM",
                "onnxruntime installed successfully.",
            )
            return True

        QMessageBox.critical(
            self.iface.mainWindow(),
            "Installation failed",
            result.get("msg", "Unknown error installing onnxruntime."),
        )
        return False

    def _get_sam_model(self) -> Sam2OnnxPredictor | None:
        """Lazy-load SAM2 model on first use. Returns None if model unavailable."""
        if self._sam_model is None:
            if not self._ensure_onnxruntime():
                return None
            model_dir = self._get_model_dir()
            if model_dir is None:
                return None
            from samgis_core.prediction_api.sam2_adapter import Sam2OnnxPredictor

            self._sam_model = Sam2OnnxPredictor(model_dir)
            logger.info(f"SAM2 model loaded from {model_dir}")
        return self._sam_model

    def show_change_model_dialog(self) -> None:
        """Show model selection dialog and switch to a different variant."""
        self._request_id = bind_request_id()
        from .model_dialog import ModelDownloadDialog
        from .model_registry import is_model_downloaded, save_active_variant

        dialog = ModelDownloadDialog(
            self.iface.mainWindow(),
            active_variant=self._active_variant,
            request_id=self._request_id,
        )
        if dialog.exec() != QDialog.DialogCode.Accepted:
            return

        variant = dialog.selected_variant()
        if not is_model_downloaded(variant):
            return

        # Reset model so next inference loads the new variant
        self._sam_model = None
        self._active_variant = variant
        save_active_variant(variant)
        self._update_model_tooltip()
        self.iface.messageBar().pushSuccess(
            "SAM",
            f"Model changed to {variant}. Will load on next inference.",
        )

    def _load_geojson_layer(self, geojson_str: str, layer_name: str) -> bool:
        """Write GeoJSON to temp file and load as QgsVectorLayer.

        Returns True if layer was added successfully.
        """
        tmp = tempfile.NamedTemporaryFile(  # noqa: SIM115 — delete=False: file must persist after close
            suffix=".geojson",
            delete=False,
            mode="w",
            prefix="samgis_",
        )
        tmp.write(geojson_str)
        tmp.close()
        logger.debug(f"Wrote GeoJSON to temp file: {tmp.name}")

        vlayer = QgsVectorLayer(tmp.name, layer_name, "ogr")
        if vlayer.isValid():
            QgsProject.instance().addMapLayer(vlayer)
            self.iface.mapCanvas().refresh()
            return True

        logger.warning(f"Failed to load GeoJSON from {tmp.name}")
        os.unlink(tmp.name)
        return False

    # TODO: add interactive map tools — click to place FG/BG points,
    #  drag to draw rectangles directly on the map canvas (QgsMapTool subclass)
    #  instead of requiring pre-existing vector layers.

    def show_vector_prompt_dialog(self) -> None:
        """Show dialog to select vector layers as SAM prompts (local inference)."""
        self._request_id = bind_request_id()
        from .prompt_builder import geo_to_sam_prompts

        layers = QgsProject.instance().mapLayers().values()
        raster_layers = [lyr for lyr in layers if isinstance(lyr, QgsRasterLayer)]

        if not raster_layers:
            self.iface.messageBar().pushWarning(
                "Warning", "No raster layers found in project."
            )
            return

        dialog = VectorPromptDialog(self.iface.mainWindow())
        if dialog.exec() != QDialog.DialogCode.Accepted:
            return

        raster = dialog.selected_raster()
        vector_layers = dialog.selected_vector_layers()

        if raster is None:
            self.iface.messageBar().pushWarning("Warning", "No raster layer selected.")
            return
        if not vector_layers:
            self.iface.messageBar().pushWarning("Warning", "No vector layers selected.")
            return

        # Extract raster data (image + transform from same source)
        image, transform, raster_w, raster_h, raster_crs = qgis_raster.get_raster_data(
            raster
        )
        logger.info(
            f"Raster '{raster.name()}': image={raster_w}x{raster_h}, crs={raster_crs}",
        )

        # Check for point layers missing 'label' field (warned after inference)
        layers_missing_label = [
            vlayer.name()
            for vlayer in vector_layers
            if (
                vlayer.geometryType() == QgsWkbTypes.PointGeometry
                and vlayer.fields().indexOf("label") < 0
            )
        ]

        # Collect geometries from all selected vector layers
        all_geometries = []
        for vlayer in vector_layers:
            geoms = qgis_vector.extract_geometries(vlayer, raster_crs)
            all_geometries.extend(geoms)

        logger.info(
            "Collected %d geometries from %d vector layers",
            len(all_geometries),
            len(vector_layers),
        )

        # Convert to SAM prompts
        prompts = geo_to_sam_prompts(all_geometries, transform, raster_w, raster_h)

        if not prompts:
            self.iface.messageBar().pushWarning(
                "Warning", "No valid prompts generated (all out of bounds?)."
            )
            return

        # SAM2 supports at most 1 rectangle per predict() call
        n_rects = sum(1 for p in prompts if p["type"] == "rectangle")
        if n_rects > 1:
            self.iface.messageBar().pushWarning(
                "SAM",
                f"SAM2 supports at most 1 rectangle per prediction, "
                f"but {n_rects} were found. Please use a single polygon feature.",
            )
            return

        sam = self._get_sam_model()
        if sam is None:
            return

        # Submit inference as background task (UI stays responsive)
        from .tasks import Sam2InferenceTask

        vector_names = ", ".join(v.name() for v in vector_layers)
        layer_name = f"SAM local - {raster.name()} - {vector_names}"

        task = Sam2InferenceTask(
            sam_model=sam,
            image=image,
            prompts=prompts,
            transform=transform,
            raster_crs=raster_crs,
            raster_w=raster_w,
            raster_h=raster_h,
            layer_name=layer_name,
            layers_missing_label=layers_missing_label,
            request_id=self._request_id,
        )
        task.taskCompleted.connect(lambda: self._on_inference_complete(task))
        task.taskTerminated.connect(lambda: self._on_inference_failed(task))

        # Keep a reference so the task isn't garbage-collected
        self._current_task = task
        QgsApplication.taskManager().addTask(task)
        self.iface.messageBar().pushInfo("SAM", "Inference started...")

    def _on_inference_complete(self, task: Sam2InferenceTask) -> None:
        """Callback on main thread after successful inference."""
        self._current_task = None

        if task.geojson_str is None:
            self.iface.messageBar().pushCritical("SAM", "Inference produced no output.")
            return

        if self._load_geojson_layer(task.geojson_str, task.layer_name):
            self.iface.messageBar().pushSuccess(
                "SAM",
                f"Added {task.n_shapes} polygons",
            )
        else:
            self.iface.messageBar().pushCritical(
                "SAM",
                "Failed to create vector layer from results.",
            )

        if task.layers_missing_label:
            names = ", ".join(task.layers_missing_label)
            self.iface.messageBar().pushWarning(
                "SAM",
                f"Layer(s) '{names}' have no 'label' field — "
                f"all points default to foreground (1).",
            )

    def _on_inference_failed(self, task: Sam2InferenceTask) -> None:
        """Callback on main thread after failed/cancelled inference."""
        self._current_task = None
        if task.error:
            self.iface.messageBar().pushCritical(
                "SAM", f"Inference failed: {task.error}"
            )
        else:
            self.iface.messageBar().pushWarning("SAM", "Inference was cancelled.")
