from __future__ import annotations

import os
import time

import numpy as np
from qgis.core import Qgis, QgsMessageLog

from .i18n import tr

# Raster formats that pip-installed rasterio may not support reliably.
# These go through GDAL windowed read instead.
_GDAL_ONLY_FORMATS = {
    ".ecw", ".sid", ".jp2", ".j2k", ".j2c",
    ".nitf", ".ntf", ".img", ".hdf", ".hdf5", ".he5", ".nc",
    ".gpkg",
    ".pdf",
}

# Online/remote raster providers that need rendering before encoding
ONLINE_PROVIDERS = frozenset(["wms", "wmts", "xyz", "arcgismapserver", "wcs"])


def _normalize_to_uint8(bands, nodata_value=None):
    """Normalize a multi-band array to (H, W, 3) uint8 using per-band percentile stretch.

    Args:
        bands: numpy array of shape (C, H, W), any numeric dtype
        nodata_value: optional nodata value to mask before computing percentiles

    Returns:
        numpy array of shape (H, W, 3) uint8
    """
    num_bands = bands.shape[0]
    if num_bands == 1:
        bands = np.repeat(bands, 3, axis=0)
    elif num_bands == 2:
        bands = np.stack([bands[0], bands[1], bands[0]], axis=0)
    elif num_bands > 3:
        bands = bands[:3, :, :]

    # Build nodata mask: True where pixel is valid
    if nodata_value is not None:
        valid_mask = bands[0] != nodata_value
    else:
        valid_mask = np.ones(bands.shape[1:], dtype=bool)

    # Also mask NaN for float types
    if np.issubdtype(bands.dtype, np.floating):
        for b in range(bands.shape[0]):
            valid_mask = valid_mask & ~np.isnan(bands[b])

    is_uint8 = (bands.dtype == np.uint8)
    result = np.zeros((3, bands.shape[1], bands.shape[2]), dtype=np.uint8)

    for b in range(3):
        band = bands[b].astype(np.float64)
        valid_pixels = band[valid_mask]

        if valid_pixels.size == 0:
            continue

        p2, p98 = np.percentile(valid_pixels, [2, 98])

        if is_uint8:
            # Only stretch if histogram is compressed
            if p98 - p2 < 220:
                if p98 > p2:
                    stretched = np.clip(band, p2, p98)
                    stretched = (stretched - p2) / (p98 - p2) * 255
                    result[b] = stretched.astype(np.uint8)
                else:
                    result[b] = band.astype(np.uint8)
            else:
                result[b] = bands[b]
        else:
            if p98 > p2:
                stretched = np.clip(band, p2, p98)
                stretched = (stretched - p2) / (p98 - p2) * 255
                result[b] = stretched.astype(np.uint8)
            # else: stays zeros

    # Zero out nodata pixels
    if nodata_value is not None or np.issubdtype(bands.dtype, np.floating):
        nodata_mask = ~valid_mask
        if np.any(nodata_mask):
            for b in range(3):
                result[b][nodata_mask] = 0

    # CHW -> HWC, returned C-contiguous. tobytes() in the worker handshake is
    # order-correct even on a transpose view, but a contiguous buffer keeps the
    # SAM predictor (which may take .data / buffer views internally) on the
    # well-tested path and avoids a latent footgun for future callers.
    return np.ascontiguousarray(np.transpose(result, (1, 2, 0)))


def _apply_colormap(indices, colormap):
    """Map a (H, W) palette-index array to (H, W, 3) uint8 RGB.

    A paletted raster (land cover, thematic classes) stores class INDICES, not
    colours. Reading the raw index band and stretching it (the default path)
    produces a meaningless grey ramp instead of the coloured classes the user
    sees. Expanding through the colour table restores the real classes.

    Args:
        indices: (H, W) integer array of palette indices.
        colormap: {index: (r, g, b[, a])} mapping (GDAL colour table or
                  rasterio ``src.colormap``).

    Returns:
        (H, W, 3) uint8 RGB, or None when the input cannot be mapped. Indices
        outside the table render black.
    """
    idx = np.asarray(indices)
    if idx.ndim != 2 or idx.size == 0 or not colormap:
        return None
    max_i = int(idx.max(initial=0))
    lut = np.zeros((max_i + 1, 3), dtype=np.uint8)
    for i, color in colormap.items():
        if 0 <= int(i) <= max_i and color is not None and len(color) >= 3:
            lut[int(i), 0] = int(color[0]) & 0xFF
            lut[int(i), 1] = int(color[1]) & 0xFF
            lut[int(i), 2] = int(color[2]) & 0xFF
    flat = np.clip(idx.astype(np.int64), 0, max_i)
    return lut[flat]


def _read_palette_rgb_rasterio(src, window, out_h, out_w):
    """Return (H, W, 3) uint8 for a paletted rasterio source, or None when the
    source is not paletted / has no usable colour table. Reads band 1 with
    NEAREST resampling so class indices are never interpolated (bilinear on
    indices would invent nonexistent classes at edges)."""
    try:
        from rasterio.enums import ColorInterp, Resampling

        interp = src.colorinterp
        if not interp or interp[0] != ColorInterp.palette:
            return None
        cmap = src.colormap(1)
        if not cmap:
            return None
        idx = src.read(
            1, window=window, out_shape=(out_h, out_w),
            resampling=Resampling.nearest,
        )
        return _apply_colormap(idx, cmap)
    except Exception:  # noqa: BLE001 - fall back to raw normalization on any issue
        return None


def _read_palette_rgb_gdal(ds, col_off, row_off, actual_w, actual_h, out_w, out_h):
    """Return (H, W, 3) uint8 for a paletted GDAL dataset, or None when the
    band is not paletted. GDAL windowed reads with buf_x/ysize use nearest
    resampling by default, so class indices are preserved."""
    try:
        from osgeo import gdal

        band1 = ds.GetRasterBand(1)
        ctable = band1.GetColorTable()
        if ctable is None or band1.GetColorInterpretation() != gdal.GCI_PaletteIndex:
            return None
        idx = band1.ReadAsArray(
            col_off, row_off, actual_w, actual_h,
            buf_xsize=out_w, buf_ysize=out_h,
        )
        cmap = {ci: ctable.GetColorEntry(ci) for ci in range(ctable.GetCount())}
        return _apply_colormap(idx, cmap)
    except Exception:  # noqa: BLE001 - fall back to raw normalization on any issue
        return None


def _fetch_online_bands(provider, extent, width, height):
    """Fetch raw band data from an online raster provider.

    Handles ARGB32 formats directly and fetches individual bands for
    all other data types (Byte, UInt16, Int16, UInt32, Int32, Float32, Float64).

    Args:
        provider: QgsRasterDataProvider
        extent: QgsRectangle for the area to fetch
        width: pixel width to request
        height: pixel height to request

    Returns:
        (bands_array, is_argb, error) where:
        - bands_array: numpy array (C, H, W) or (H, W, 3) uint8 if ARGB
        - is_argb: True if result is already RGB uint8 (H, W, 3)
        - error: error string or None
    """
    block = provider.block(1, extent, width, height)
    if block is None or not block.isValid():
        return None, False, "Provider block fetch failed"

    block_w = block.width()
    block_h = block.height()
    if block_w == 0 or block_h == 0:
        return None, False, "Provider returned empty block"

    raw_bytes = block.data()
    if raw_bytes is None or len(raw_bytes) == 0:
        return None, False, "Provider returned empty data"

    raw_data = bytes(raw_bytes)
    dt = block.dataType()

    # ARGB32 formats: direct extraction
    is_argb32 = dt == Qgis.DataType.ARGB32
    is_argb32_pre = dt == Qgis.DataType.ARGB32_Premultiplied
    is_argb = is_argb32 or is_argb32_pre
    if is_argb:
        arr = np.frombuffer(raw_data, dtype=np.uint8).reshape(
            block_h, block_w, 4).copy()
        # Qt BGRA byte order -> RGB
        image_np = np.stack(
            [arr[:, :, 2], arr[:, :, 1], arr[:, :, 0]], axis=-1)
        return image_np, True, None

    # Map Qgis data types to numpy dtypes
    dtype_map = {
        Qgis.DataType.Byte: np.uint8,
        Qgis.DataType.UInt16: np.uint16,
        Qgis.DataType.Int16: np.int16,
        Qgis.DataType.UInt32: np.uint32,
        Qgis.DataType.Int32: np.int32,
        Qgis.DataType.Float32: np.float32,
        Qgis.DataType.Float64: np.float64,
    }

    np_dtype = dtype_map.get(dt)
    if np_dtype is None:
        # Unknown type: try ARGB32 interpretation as fallback
        if len(raw_data) == block_w * block_h * 4:
            arr = np.frombuffer(raw_data, dtype=np.uint8).reshape(
                block_h, block_w, 4).copy()
            image_np = np.stack(
                [arr[:, :, 2], arr[:, :, 1], arr[:, :, 0]], axis=-1)
            return image_np, True, None
        return None, False, f"Unsupported data type: {dt}"

    band_count = min(provider.bandCount(), 3)
    bands = []

    # First band already fetched
    band1 = np.frombuffer(raw_data, dtype=np_dtype).reshape(
        block_h, block_w).copy()
    bands.append(band1)

    # Fetch remaining bands
    for band_idx in range(2, band_count + 1):
        b = provider.block(band_idx, extent, block_w, block_h)
        if b is not None and b.isValid():
            b_data = bytes(b.data())
            if len(b_data) > 0:
                band_arr = np.frombuffer(
                    b_data, dtype=np_dtype
                ).reshape(block_h, block_w).copy()
                bands.append(band_arr)

    bands_array = np.stack(bands, axis=0)
    return bands_array, False, None


def _render_layer_to_image(layer, extent, width, height):
    """Render a layer to an RGB image using QGIS map renderer (fallback).

    Works with any layer type (WMS, WMTS, XYZ, WCS, vector tiles, etc.)

    Args:
        layer: QgsMapLayer to render
        extent: QgsRectangle for the area
        width: pixel width
        height: pixel height

    Returns:
        (image_np, error) where image_np is (H, W, 3) uint8 or None
    """
    try:
        from qgis.core import QgsMapRendererCustomPainterJob, QgsMapSettings
        from qgis.PyQt.QtCore import QSize
        from qgis.PyQt.QtGui import QImage, QPainter

        img = QImage(QSize(width, height), QImage.Format.Format_RGB32)
        img.fill(0)

        settings = QgsMapSettings()
        settings.setOutputSize(QSize(width, height))
        settings.setExtent(extent)
        settings.setLayers([layer])
        settings.setDestinationCrs(layer.crs())
        settings.setBackgroundColor(img.pixelColor(0, 0))

        painter = QPainter(img)
        job = QgsMapRendererCustomPainterJob(settings, painter)
        job.start()
        job.waitForFinished()
        painter.end()

        # QImage -> numpy
        img = img.convertToFormat(QImage.Format.Format_RGB32)
        ptr = img.bits()
        ptr.setsize(height * width * 4)
        arr = np.frombuffer(ptr, dtype=np.uint8).reshape(
            height, width, 4).copy()
        # BGRA -> RGB
        image_np = np.stack(
            [arr[:, :, 2], arr[:, :, 1], arr[:, :, 0]], axis=-1)
        return image_np, None

    except Exception as e:
        return None, f"Renderer fallback failed: {str(e)}"


def _needs_gdal_conversion(raster_path):
    """Check if raster format requires GDAL conversion for rasterio."""
    ext = os.path.splitext(raster_path)[1].lower()
    return ext in _GDAL_ONLY_FORMATS


def _read_crop_with_gdal(raster_path, center_x, center_y, crop_size,
                         scale_factor, layer_extent):
    """Read a windowed crop directly from a GDAL-supported raster (JP2, ECW, etc.).

    Instead of converting the entire file to GeoTIFF, reads only the needed
    pixels using GDAL windowed access.

    Returns (image_np, crop_info, error_msg, error_code). On success, error_msg
    and error_code are both None. On failure error_code is a stable English
    snake_case identifier safe for analytics aggregation across locales.
    """
    ext = os.path.splitext(raster_path)[1].upper()

    try:
        from osgeo import gdal
    except ImportError:
        return None, None, tr(
            "{ext} format is not directly supported. "
            "GDAL is not available.\n"
            "Please convert your raster to GeoTIFF (.tif) before using "
            "AI Segmentation."
        ).format(ext=ext), "crop_error_gdal_unavailable"

    ds = None
    # rasterio bundles its own proj.db in site-packages, and we force
    # PROJ_DATA/GDAL_DATA to point there so rasterio works in the venv.
    # GDAL (from osgeo) has its own bundled data and will hit a
    # "DATABASE.LAYOUT.VERSION.MINOR = N whereas a number >= M is expected"
    # error when forced to read rasterio's older proj.db. Shadow the vars
    # via GDAL's own config-option table (scoped to this GDAL instance,
    # like _scope_rasterio_data_paths does for rasterio) instead of
    # os.environ, so no other thread or plugin in the process ever sees
    # them missing while this call runs.
    _proj_data_backup = gdal.GetConfigOption("PROJ_DATA")
    _proj_lib_backup = gdal.GetConfigOption("PROJ_LIB")
    _gdal_data_backup = gdal.GetConfigOption("GDAL_DATA")
    gdal.SetConfigOption("PROJ_DATA", "")
    gdal.SetConfigOption("PROJ_LIB", "")
    gdal.SetConfigOption("GDAL_DATA", "")
    try:
        ds = gdal.Open(raster_path)
        if ds is None:
            return None, None, tr(
                "Cannot open {ext} file. The format may not be supported "
                "by your QGIS installation.\n"
                "Please convert your raster to GeoTIFF (.tif) before using "
                "AI Segmentation."
            ).format(ext=ext), "crop_error_unsupported_format"

        raster_width = ds.RasterXSize
        raster_height = ds.RasterYSize
        gt = ds.GetGeoTransform()

        use_layer_extent = False
        if layer_extent:
            xmin_le, ymin_le, xmax_le, ymax_le = layer_extent
            if gt is None or gt == (0, 1, 0, 0, 0, 1):
                use_layer_extent = True
            else:
                left_near = abs(gt[0]) < 10
                top_near = abs(gt[3]) < 10
                right_near = abs(gt[0] + gt[1] * raster_width) < 10
                bottom_near = abs(gt[3] + gt[5] * raster_height) < 10
                if left_near and top_near and right_near and bottom_near:
                    use_layer_extent = True

        if use_layer_extent and layer_extent:
            xmin_le, ymin_le, xmax_le, ymax_le = layer_extent
            pixel_size_x = (xmax_le - xmin_le) / raster_width
            pixel_size_y = (ymax_le - ymin_le) / raster_height
            bounds_left = xmin_le
            bounds_top = ymax_le
        else:
            pixel_size_x = abs(gt[1])
            pixel_size_y = abs(gt[5])
            bounds_left = gt[0]
            bounds_top = gt[3]

        col_center = (center_x - bounds_left) / pixel_size_x
        row_center = (bounds_top - center_y) / pixel_size_y

        read_size = int(crop_size * scale_factor)
        half = read_size // 2
        col_off = max(0, int(round(col_center - half)))
        row_off = max(0, int(round(row_center - half)))

        actual_width = min(read_size, raster_width - col_off)
        actual_height = min(read_size, raster_height - row_off)

        if actual_width <= 0 or actual_height <= 0:
            return None, None, "Click is outside the raster bounds", "crop_error_outside_bounds"

        num_bands = min(ds.RasterCount, 3)
        if num_bands == 0:
            return None, None, "Raster has no bands", "crop_error_no_bands"

        if scale_factor > 1.0:
            out_h = max(1, min(crop_size, int(actual_height / scale_factor)))
            out_w = max(1, min(crop_size, int(actual_width / scale_factor)))
        elif scale_factor < 1.0:
            out_h = min(crop_size, max(1, int(actual_height / scale_factor)))
            out_w = min(crop_size, max(1, int(actual_width / scale_factor)))
        else:
            out_h = actual_height
            out_w = actual_width

        # Paletted (colour-table) raster: expand class indices to their true
        # colours instead of stretching raw indices to grey.
        palette_rgb = _read_palette_rgb_gdal(
            ds, col_off, row_off, actual_width, actual_height, out_w, out_h)
        if palette_rgb is not None:
            image_np = palette_rgb
            ds = None
        else:
            bands = []
            for b_idx in range(1, num_bands + 1):
                band = ds.GetRasterBand(b_idx)
                data = band.ReadAsArray(
                    col_off, row_off, actual_width, actual_height,
                    buf_xsize=out_w, buf_ysize=out_h
                )
                # GDAL returns None (not an exception) when the block cannot be
                # decoded: corrupt tile, missing codec (JPEG-in-TIFF), broken
                # overview. np.stack on a None would raise a cryptic error that
                # surfaces as crop_error_read_failed. Catch it here with guidance.
                if data is None:
                    return None, None, tr(
                        "Could not read pixels from this {ext} file. The file may "
                        "be corrupt, truncated, or use a compression your GDAL "
                        "build cannot decode.\n"
                        "Try opening it in QGIS to confirm it displays, or convert "
                        "it to GeoTIFF (.tif) before using AI Segmentation."
                    ).format(ext=ext), "crop_error_read_failed"
                bands.append(data)

            nodata = ds.GetRasterBand(1).GetNoDataValue()
            ds = None

            tile_data = np.stack(bands, axis=0)
            image_np = _normalize_to_uint8(tile_data, nodata_value=nodata)

        if out_h < crop_size or out_w < crop_size:
            pad_bottom = crop_size - out_h
            pad_right = crop_size - out_w
            image_np = np.pad(
                image_np,
                ((0, pad_bottom), (0, pad_right), (0, 0)),
                mode="reflect"
            )

        crop_minx = bounds_left + col_off * pixel_size_x
        crop_maxx = bounds_left + (col_off + actual_width) * pixel_size_x
        crop_maxy = bounds_top - row_off * pixel_size_y
        crop_miny = bounds_top - (row_off + actual_height) * pixel_size_y

        crop_info = {
            "bounds": (crop_minx, crop_miny, crop_maxx, crop_maxy),
            "img_shape": (out_h, out_w),
            "col_off": col_off,
            "row_off": row_off,
        }

        QgsMessageLog.logMessage(
            f"Read {ext} crop directly via GDAL: {out_w}x{out_h} at ({col_off}, {row_off})",
            "AI Segmentation", level=Qgis.MessageLevel.Info
        )
        return image_np, crop_info, None, None

    except Exception as e:
        return None, None, tr(
            "Failed to read {ext} file: {error}\n"
            "Please convert your raster to GeoTIFF (.tif) manually."
        ).format(ext=ext, error=str(e)), "crop_error_read_failed"

    finally:
        ds = None
        gdal.SetConfigOption("PROJ_DATA", _proj_data_backup)
        gdal.SetConfigOption("PROJ_LIB", _proj_lib_backup)
        gdal.SetConfigOption("GDAL_DATA", _gdal_data_backup)


def extract_crop_from_raster(raster_path, center_x, center_y, crop_size=1024,
                             layer_crs_wkt=None, layer_extent=None,
                             scale_factor=1.0):
    """Extract a crop_size x crop_size RGB crop centered on (center_x, center_y).

    Args:
        raster_path: Path to the raster file
        center_x, center_y: Center of crop in geo/pixel coordinates
        crop_size: Size of the crop in pixels (default 1024)
        layer_crs_wkt: Optional CRS WKT for non-georeferenced rasters
        layer_extent: Optional (xmin, ymin, xmax, ymax) for non-georeferenced rasters
        scale_factor: Controls native read size relative to crop_size.
            When > 1.0, reads more native pixels and downsamples to crop_size.
            When < 1.0, reads fewer native pixels and upsamples to crop_size.
            Clamped to [0.25, 8.0] by the caller.

    Returns:
        (image_np, crop_info, error) where:
        - image_np: (H, W, 3) uint8 numpy array
        - crop_info: dict with 'bounds' (minx, miny, maxx, maxy) and 'img_shape' (H, W)
        - error: error string or None on success
    """
    # Handle GDAL-only formats (JP2, ECW, etc.) with direct windowed read.
    # Checked BEFORE the rasterio import: this path reads through QGIS's own
    # GDAL, so it must keep working when the venv rasterio is broken/missing.
    if _needs_gdal_conversion(raster_path):
        return _read_crop_with_gdal(
            raster_path, center_x, center_y, crop_size,
            scale_factor, layer_extent
        )

    try:
        import rasterio
        from rasterio.enums import Resampling
        from rasterio.windows import Window
    except ImportError:
        return None, None, "rasterio is not available", "crop_error_rasterio_unavailable"

    try:
        with rasterio.open(raster_path) as src:
            raster_width = src.width
            raster_height = src.height
            raster_transform = src.transform

            # Detect non-georeferenced mode
            use_layer_extent = False
            if layer_extent:
                xmin_le, ymin_le, xmax_le, ymax_le = layer_extent
                if src.crs is None:
                    use_layer_extent = True
                else:
                    rb = src.bounds
                    left_near = abs(rb.left) < 10
                    bottom_near = abs(rb.bottom) < 10
                    right_near = abs(rb.right - raster_width) < 10
                    top_near = abs(rb.top - raster_height) < 10
                    if left_near and bottom_near and right_near and top_near:
                        use_layer_extent = True

            if use_layer_extent and layer_extent:
                xmin_le, ymin_le, xmax_le, ymax_le = layer_extent
                pixel_size_x = (xmax_le - xmin_le) / raster_width
                pixel_size_y = (ymax_le - ymin_le) / raster_height
                bounds_left = xmin_le
                bounds_top = ymax_le
            else:
                pixel_size_x = abs(raster_transform.a)
                pixel_size_y = abs(raster_transform.e)
                bounds_left = src.bounds.left
                bounds_top = src.bounds.top

            # Convert geo coords to pixel coords
            col_center = (center_x - bounds_left) / pixel_size_x
            row_center = (bounds_top - center_y) / pixel_size_y

            # When scale_factor > 1, read a larger native window
            read_size = int(crop_size * scale_factor)
            half = read_size // 2
            col_off = max(0, int(round(col_center - half)))
            row_off = max(0, int(round(row_center - half)))

            actual_width = min(read_size, raster_width - col_off)
            actual_height = min(read_size, raster_height - row_off)

            if actual_width <= 0 or actual_height <= 0:
                return None, None, "Click is outside the raster bounds", "crop_error_outside_bounds"

            window = Window(col_off, row_off, actual_width, actual_height)

            if scale_factor > 1.0:
                out_h = min(crop_size, int(actual_height / scale_factor))
                out_w = min(crop_size, int(actual_width / scale_factor))
                out_h = max(1, out_h)
                out_w = max(1, out_w)
                tile_data = src.read(
                    window=window,
                    out_shape=(src.count, out_h, out_w),
                    resampling=Resampling.bilinear
                )
            elif scale_factor < 1.0:
                out_h = min(crop_size, max(1, int(actual_height / scale_factor)))
                out_w = min(crop_size, max(1, int(actual_width / scale_factor)))
                tile_data = src.read(
                    window=window,
                    out_shape=(src.count, out_h, out_w),
                    resampling=Resampling.bilinear
                )
            else:
                tile_data = src.read(window=window)
                out_h = actual_height
                out_w = actual_width

            nodata = src.nodata
            # Paletted (colour-table) raster: expand class indices to their true
            # colours (read band 1 with nearest so indices are never blended).
            palette_rgb = _read_palette_rgb_rasterio(src, window, out_h, out_w)
            if palette_rgb is not None:
                image_np = palette_rgb
            else:
                image_np = _normalize_to_uint8(tile_data, nodata_value=nodata)

            # Pad to full crop_size if crop was clipped at raster edge.
            # Uses reflect padding instead of black borders for better
            # SAM context at image boundaries.
            if out_h < crop_size or out_w < crop_size:
                pad_bottom = crop_size - out_h
                pad_right = crop_size - out_w
                image_np = np.pad(
                    image_np,
                    ((0, pad_bottom), (0, pad_right), (0, 0)),
                    mode="reflect"
                )

            # Compute geo bounds for this crop (covers the full read area)
            crop_minx = bounds_left + col_off * pixel_size_x
            crop_maxx = bounds_left + (col_off + actual_width) * pixel_size_x
            crop_maxy = bounds_top - row_off * pixel_size_y
            crop_miny = bounds_top - (row_off + actual_height) * pixel_size_y

            crop_info = {
                "bounds": (crop_minx, crop_miny, crop_maxx, crop_maxy),
                "img_shape": (out_h, out_w),
                "col_off": col_off,
                "row_off": row_off,
            }

            return image_np, crop_info, None, None

    except Exception as e:
        # Fallback to GDAL if rasterio fails (unsupported driver, etc.)
        QgsMessageLog.logMessage(
            f"rasterio failed ({str(e)}), trying GDAL fallback...",
            "AI Segmentation", level=Qgis.MessageLevel.Warning
        )
        return _read_crop_with_gdal(
            raster_path, center_x, center_y, crop_size,
            scale_factor, layer_extent
        )


class OnlineCropFetcher:
    """Step-driven online-tile crop fetch.

    One provider block read + stabilization check is one ``step()``. The
    caller drives the retry back-off between steps: the interactive path
    schedules them with QTimer so the waits never block the GUI thread, while
    extract_crop_from_online_layer drives the same steps in a plain blocking
    loop for the headless/recovery contract. The retry policy (attempt count,
    progressive delays, blank-tile handling) lives here once, so both paths
    behave identically.

    The provider is the user's LIVE layer provider. The resampling state the
    fetch forces (provider resampling on, bilinear zoomed-in/out) is
    snapshotted by begin() and reverted by restore(); the caller MUST call
    restore() on every exit so the user's layer is never left mutated.
    """

    _MAX_RETRIES = 8
    _RETRY_DELAY = 1.0

    def __init__(self, layer, center_x, center_y, canvas_mupp, crop_size=1024):
        from qgis.core import QgsRectangle

        self._layer = layer
        self._crop_size = crop_size
        self.error = None
        self.error_code = None
        self._mutated = False
        self._orig_in = None
        self._orig_out = None
        self._orig_enabled = None
        self._attempt = 0
        self._prev_data = None
        self._reload_pending = False

        provider = layer.dataProvider()
        self._provider = provider
        if provider is None:
            self.error = tr("Layer data provider is not available.")
            self.error_code = "crop_error_online_provider_unavailable"
            return

        half_size = crop_size * canvas_mupp / 2.0
        self._extent = QgsRectangle(
            center_x - half_size, center_y - half_size,
            center_x + half_size, center_y + half_size
        )
        QgsMessageLog.logMessage(
            f"Online crop request: center=({center_x:.6f}, {center_y:.6f}), "
            f"mupp={canvas_mupp:.6f}, extent=({self._extent.xMinimum():.2f}, "
            f"{self._extent.yMinimum():.2f}, {self._extent.xMaximum():.2f}, "
            f"{self._extent.yMaximum():.2f}), CRS={layer.crs().authid()}",
            "AI Segmentation", level=Qgis.MessageLevel.Info
        )

    def begin(self):
        """Snapshot the live provider's resampling state, then force bilinear
        for the fetch. Safe to call once; restore() reverts whatever was
        captured, even if a setter below raises partway through."""
        if self._provider is None:
            return
        provider = self._provider
        try:
            self._orig_in = provider.zoomedInResamplingMethod()
            self._orig_out = provider.zoomedOutResamplingMethod()
        except (AttributeError, RuntimeError):
            # Provider without the resampling API: fetch at native resampling,
            # nothing to revert.
            return
        # isProviderResamplingEnabled() is the getter paired with
        # enableProviderResampling(); on a build that lacks it the enabled flag
        # stays unknown (None) and restore() leaves it on.
        if hasattr(provider, "isProviderResamplingEnabled"):
            try:
                self._orig_enabled = provider.isProviderResamplingEnabled()
            except (AttributeError, RuntimeError):
                self._orig_enabled = None
        # Past this point a revert is owed regardless of what the setters do.
        self._mutated = True
        try:
            provider.enableProviderResampling(True)
            provider.setZoomedInResamplingMethod(
                provider.ResamplingMethod.Bilinear)
            provider.setZoomedOutResamplingMethod(
                provider.ResamplingMethod.Bilinear)
        except (AttributeError, RuntimeError):
            pass

    def restore(self):
        """Revert the provider resampling state captured by begin().
        Idempotent and best-effort (the provider may be mid-teardown when a
        superseded async fetch unwinds)."""
        if not self._mutated:
            return
        self._mutated = False
        provider = self._provider
        try:
            provider.setZoomedInResamplingMethod(self._orig_in)
            provider.setZoomedOutResamplingMethod(self._orig_out)
            if self._orig_enabled is not None:
                provider.enableProviderResampling(self._orig_enabled)
        except (AttributeError, RuntimeError):
            pass

    def step(self):
        """Perform one block read + stabilization check. Returns (action,
        delay_seconds):

        - ("stabilized", 0.0): two equal non-blank reads; go to finish().
        - ("refetch", 0.5): first valid non-blank read; reload issued, wait
          then step() again.
        - ("retry", delay): not stable yet; wait ``delay`` then step() again
          (the reload is issued at the START of the next step, matching the
          blocking loop's wait-then-reload order).
        - ("exhausted", 0.0): retry budget spent; go to finish() with whatever
          the provider returned.
        """
        provider = self._provider
        if self._reload_pending:
            # The prior retry has served its back-off; refresh the tiles now,
            # exactly as the blocking loop reloaded before each re-read.
            provider.reloadData()
            self._reload_pending = False
        attempt = self._attempt
        block = provider.block(1, self._extent, self._crop_size, self._crop_size)
        if block is not None and block.isValid():
            cur_data = bytes(block.data())
            # An all-zero block means the tiles have not downloaded yet (or the
            # area has no coverage). Two blank reads in a row must NOT count as
            # "stabilized": a slow network would trip the equality check and
            # abandon the retry budget after ~1s, shipping a blank crop. Only
            # real pixel content can end the loop early; blank reads always
            # fall through to the progressive retry branch.
            if cur_data.strip(b"\x00"):
                if self._prev_data is not None and cur_data == self._prev_data:
                    return ("stabilized", 0.0)
                self._prev_data = cur_data
                if attempt == 0:
                    # First valid fetch: reload once and re-check whether the
                    # tiles are still loading/updating.
                    provider.reloadData()
                    self._attempt = attempt + 1
                    return ("refetch", 0.5)
        if attempt < self._MAX_RETRIES - 1:
            delay = self._RETRY_DELAY * (1 + attempt * 0.5)  # 1.0, 1.5, 2.0, ...
            QgsMessageLog.logMessage(
                f"Online tile fetch attempt {attempt + 1} - "
                f"retrying in {delay:.1f}s...",
                "AI Segmentation", level=Qgis.MessageLevel.Warning
            )
            self._reload_pending = True
            self._attempt = attempt + 1
            return ("retry", delay)
        return ("exhausted", 0.0)

    def finish(self):
        """Read the band data for the stabilized extent and build the crop.
        Returns the extract_crop_from_raster-style 4-tuple (image_np,
        crop_info, error, error_code)."""
        provider = self._provider
        # Fetch bands using unified helper
        bands_result, is_argb, fetch_err = _fetch_online_bands(
            provider, self._extent, self._crop_size, self._crop_size)

        # If provider fetch failed, try canvas renderer fallback
        if fetch_err is not None:
            QgsMessageLog.logMessage(
                f"Provider fetch failed ({fetch_err}), trying renderer "
                "fallback...",
                "AI Segmentation", level=Qgis.MessageLevel.Warning
            )
            image_np, render_err = _render_layer_to_image(
                self._layer, self._extent, self._crop_size, self._crop_size)
            if render_err is not None:
                return None, None, tr(
                    "Failed to fetch tiles from the online layer. "
                    "Check your network connection."
                ), "crop_error_online_fetch_failed"
        elif is_argb:
            # Already RGB uint8 (H, W, 3) from ARGB32 path
            image_np = bands_result
        else:
            # Raw bands (C, H, W) - normalize with nodata
            nodata = None
            try:
                nodata = provider.sourceNoDataValue(1)
            except Exception:
                pass  # nosec B110
            image_np = _normalize_to_uint8(bands_result, nodata_value=nodata)

        height = image_np.shape[0]
        width = image_np.shape[1]

        # Check for blank tiles
        total_sum = int(image_np.sum())
        if total_sum == 0:
            # Tiles fetched but all-zero: stale provider cache, zoom outside
            # the service's range, or genuinely no coverage. The canvas
            # renderer reads what QGIS actually displays, so one render-path
            # retry rescues the stale-cache case before giving up.
            rendered, render_err = _render_layer_to_image(
                self._layer, self._extent, self._crop_size, self._crop_size)
            if render_err is None and rendered is not None and int(rendered.sum()) > 0:
                QgsMessageLog.logMessage(
                    "Blank tiles from provider, rescued by renderer fallback",
                    "AI Segmentation", level=Qgis.MessageLevel.Warning
                )
                image_np = rendered
                height = image_np.shape[0]
                width = image_np.shape[1]
            else:
                return None, None, tr(
                    "Online layer returned blank tiles for this area. The "
                    "current zoom level may be outside the service's range, "
                    "or this area has no coverage. Zoom to a level where the "
                    "layer is visible on the map, then try again."
                ), "crop_error_online_blank_tiles"

        crop_info = {
            "bounds": (self._extent.xMinimum(), self._extent.yMinimum(),
                       self._extent.xMaximum(), self._extent.yMaximum()),
            "img_shape": (height, width),
        }
        return image_np, crop_info, None, None


def _blocking_wait(seconds):
    """Block the calling thread for ``seconds`` while pumping the event loop.
    Used only by the synchronous (headless/recovery) online fetch driver; the
    interactive path schedules its waits off the event loop instead."""
    from qgis.core import QgsApplication

    deadline = time.monotonic() + seconds
    while time.monotonic() < deadline:
        QgsApplication.processEvents()
        time.sleep(0.05)


def extract_crop_from_online_layer(layer, center_x, center_y, canvas_mupp,
                                   crop_size=1024):
    """Extract a crop_size x crop_size RGB crop from an online layer.

    SYNCHRONOUS and blocking (up to ~18s of tile-fetch retries): used by the
    headless/MCP and recovery callers, whose contract is a blocking call. The
    interactive path drives the same OnlineCropFetcher steps asynchronously so
    the GUI never freezes.

    Args:
        layer: QgsRasterLayer (WMS, WMTS, XYZ, WCS, ArcGIS)
        center_x, center_y: Center of crop in layer CRS coordinates
        canvas_mupp: Map units per pixel
        crop_size: Size of the crop in pixels (default 1024)

    Returns:
        (image_np, crop_info, error, error_code) - same format as
        extract_crop_from_raster.
    """
    fetcher = OnlineCropFetcher(layer, center_x, center_y, canvas_mupp, crop_size)
    if fetcher.error is not None:
        return None, None, fetcher.error, fetcher.error_code
    try:
        fetcher.begin()
        while True:
            action, delay = fetcher.step()
            if action in ("stabilized", "exhausted"):
                break
            _blocking_wait(delay)
        return fetcher.finish()
    except Exception as e:  # noqa: BLE001
        return None, None, str(e), "crop_error_online_exception"
    finally:
        fetcher.restore()
