"""QgsTask subclass for non-blocking SAM2 inference."""

from __future__ import annotations

from typing import TYPE_CHECKING

import structlog

import numpy as np
from qgis.core import QgsTask
from samgis_core.utilities.request_context import (
    bind_request_id,
    clear_request_context,
)

logger = structlog.get_logger(__name__)


class Sam2InferenceTask(QgsTask):
    """Run SAM2 inference + vectorization in a background thread.

    Results are stored as instance attributes for the main-thread callback.
    """

    if TYPE_CHECKING:
        from samgis_core.prediction_api.sam2_adapter import Sam2OnnxPredictor

    def __init__(
        self,
        sam_model: Sam2OnnxPredictor,
        image: np.ndarray,
        prompts: list[dict],
        transform: tuple[float, ...],
        raster_crs: str,
        raster_w: int,
        raster_h: int,
        layer_name: str,
        layers_missing_label: list[str],
        request_id: str | None = None,
    ) -> None:
        super().__init__("SAM2 inference", QgsTask.CanCancel)
        self._sam = sam_model
        self._image = image
        self._prompts = prompts
        self._transform = transform
        self._raster_crs = raster_crs
        self._raster_w = raster_w
        self._raster_h = raster_h

        # Stored for the main-thread callback
        self.layer_name = layer_name
        self.layers_missing_label = layers_missing_label

        self._request_id = request_id

        # Results (populated by run())
        self.geojson_str: str | None = None
        self.n_shapes: int = 0
        self.error: str | None = None

    def run(self) -> bool:
        rid = bind_request_id(self._request_id)
        logger.info(f"inference task started, request_id={rid}")
        try:
            return self._run_pipeline()
        except Exception as e:
            self.error = str(e)
            logger.exception("Sam2InferenceTask failed: %s", e)
            return False
        finally:
            clear_request_context()

    def _run_pipeline(self) -> bool:
        # --- 1. Encode image (set_image) ---
        self.setProgress(5)
        logger.info("SAM2 encoding image shape=%s", self._image.shape)
        self._sam.set_image(self._image)

        if self.isCanceled():
            return False
        self.setProgress(30)

        # --- 2. Predict masks ---
        logger.info("SAM2 predicting masks for %d prompts", len(self._prompts))
        masks, ious = self._sam.predict(self._prompts)
        logger.info("SAM2 masks shape=%s, ious=%s", masks.shape, ious)

        if self.isCanceled():
            return False
        self.setProgress(60)

        # --- 3. Select best mask ---
        best_idx = np.argmax(ious)
        mask = masks[best_idx]
        logger.info(
            "Selected mask %d (IoU=%.4f): shape=%s, nonzero=%d/%d",
            best_idx,
            ious[best_idx],
            mask.shape,
            np.count_nonzero(mask),
            mask.size,
        )

        # Resize if mask dimensions don't match raster
        if mask.shape != (self._raster_h, self._raster_w):
            logger.info(
                "Mask size mismatch: mask=%s, raster=%dx%d — resizing",
                mask.shape,
                self._raster_w,
                self._raster_h,
            )
            from .mask_utils import _resize_nearest

            mask = _resize_nearest(mask, self._raster_h, self._raster_w)

        if self.isCanceled():
            return False
        self.setProgress(70)

        # --- 4. Vectorize mask to GeoJSON ---
        from .adapters.gdal_vectorizer import GdalVectorizer

        vectorizer = GdalVectorizer()
        logger.info(
            "Vectorizing: mask=%s, transform=%s, crs=%s",
            mask.shape,
            self._transform,
            self._raster_crs,
        )
        self.geojson_str = vectorizer.vectorize(
            mask,
            self._transform,
            source_crs=self._raster_crs,
        )
        self.n_shapes = vectorizer.n_shapes

        self.setProgress(100)
        return True
