"""Adapter functions to extract data from QgsRasterLayer for SAM consumption."""

from __future__ import annotations

from typing import TYPE_CHECKING

import structlog

import numpy as np

if TYPE_CHECKING:
    from qgis.core import Qgis, QgsRasterLayer

logger = structlog.get_logger(__name__)

# TODO: Band ordering assumes 1=R, 2=G, 3=B. May need to check layer.renderer()
#  for band mapping on non-standard rasters. Revisit if needed.


def get_raster_data(layer: QgsRasterLayer) -> tuple[np.ndarray, tuple, int, int, str]:
    """Read raster image + derive transform from actual image dimensions.

    Returns (image, transform, width, height, crs) all consistent with each other.
    The transform is computed from the image shape, NOT from layer.width()/height(),
    to avoid mismatches when QGIS reports different dimensions than the pixel data.
    """
    image = get_image(layer)
    h, w = image.shape[:2]
    extent = layer.extent()
    pixel_w = extent.width() / w
    pixel_h = extent.height() / h
    transform = (
        extent.xMinimum(),
        pixel_w,
        0.0,
        extent.yMaximum(),
        0.0,
        -pixel_h,
    )
    layer_crs = layer.crs()
    crs = layer_crs.authid()
    logger.debug(
        "get_raster_data CRS: authid=%r, isValid=%s, description=%r, toWkt[:80]=%r",
        crs,
        layer_crs.isValid(),
        layer_crs.description(),
        layer_crs.toWkt()[:80],
    )
    if not crs:
        crs = layer_crs.toWkt()
        logger.warning(f"authid empty, falling back to toWkt ({len(crs)} chars)")
    if not crs:
        from qgis.core import QgsProject

        project_crs = QgsProject.instance().crs()
        if project_crs.isValid():
            crs = project_crs.authid() or project_crs.toWkt()
            logger.warning(f"CRS unset, falling back to project CRS: {crs!r}")
        else:
            crs = "EPSG:4326"
            logger.warning(f"CRS unset and project CRS invalid, defaulting to {crs}")

    # Diagnostic: check if layer dimensions match actual image dimensions
    layer_w, layer_h = layer.width(), layer.height()
    if layer_w != w or layer_h != h:
        logger.warning(
            "DIMENSION MISMATCH: layer reports %dx%d but image array is %dx%d",
            layer_w,
            layer_h,
            w,
            h,
        )
    logger.debug(
        "get_raster_data: image=%dx%d (shape=%s dtype=%s), crs=%s, "
        "extent=(%.6f,%.6f,%.6f,%.6f), pixel_size=(%.8f,%.8f), transform=%s",
        w,
        h,
        image.shape,
        image.dtype,
        crs,
        extent.xMinimum(),
        extent.yMinimum(),
        extent.xMaximum(),
        extent.yMaximum(),
        pixel_w,
        pixel_h,
        transform,
    )
    return image, transform, w, h, crs


def get_image(layer: QgsRasterLayer) -> np.ndarray:
    """Read raster bands as (H, W, 3) uint8 ndarray suitable for SAM input.

    Auto-normalizes non-uint8 rasters to 0-255 range.
    """

    provider = layer.dataProvider()
    extent = layer.extent()
    w = layer.width()
    h = layer.height()
    band_count = layer.bandCount()

    logger.debug(
        "get_image: reading %d bands, layer=%dx%d, source=%s",
        band_count,
        w,
        h,
        layer.source(),
    )

    bands = []
    for band_idx in range(1, min(band_count, 3) + 1):
        block = provider.block(band_idx, extent, w, h)
        dtype = provider.dataType(band_idx)
        np_dtype = _qgis_dtype_to_numpy(dtype)
        buf = bytes(block.data())
        arr = np.frombuffer(buf, dtype=np_dtype).reshape(h, w)
        bands.append(arr)

    # Handle grayscale: repeat single band to 3 channels
    while len(bands) < 3:
        bands.append(bands[0])

    img = np.stack(bands, axis=-1)
    logger.debug(
        "get_image: result shape=%s, dtype=%s, min=%s, max=%s",
        img.shape,
        img.dtype,
        img.min(),
        img.max(),
    )

    # Normalize to uint8 if needed
    if img.dtype != np.uint8:
        img = _normalize_to_uint8(img)

    return img


def get_transform(
    layer: QgsRasterLayer,
) -> tuple[float, float, float, float, float, float]:
    """Extract 6-coefficient GDAL geotransform from layer extent and dimensions."""
    extent = layer.extent()
    w = layer.width()
    h = layer.height()
    pixel_width = extent.width() / w
    pixel_height = extent.height() / h
    transform = (
        extent.xMinimum(),
        pixel_width,
        0.0,
        extent.yMaximum(),
        0.0,
        -pixel_height,
    )
    logger.debug("get_transform: %s (layer reports %dx%d)", transform, w, h)
    return transform


def get_crs(layer: QgsRasterLayer) -> str:
    """Return CRS authority identifier, e.g. 'EPSG:4326'."""
    return layer.crs().authid()


def get_extent(layer: QgsRasterLayer) -> tuple[float, float, float, float]:
    """Return (xmin, ymin, xmax, ymax)."""
    extent = layer.extent()
    return (
        extent.xMinimum(),
        extent.yMinimum(),
        extent.xMaximum(),
        extent.yMaximum(),
    )


def get_extent_as_4326(layer: QgsRasterLayer) -> tuple[float, float, float, float]:
    """Return raster extent reprojected to EPSG:4326 as (xmin, ymin, xmax, ymax)."""
    from qgis.core import (
        QgsCoordinateReferenceSystem,
        QgsCoordinateTransform,
        QgsProject,
    )

    source_crs = layer.crs()
    dest_crs = QgsCoordinateReferenceSystem("EPSG:4326")
    transform = QgsCoordinateTransform(source_crs, dest_crs, QgsProject.instance())
    extent_4326 = transform.transformBoundingBox(layer.extent())
    logger.debug(
        "get_extent_as_4326: source_crs=%s, extent_4326=(%.6f,%.6f,%.6f,%.6f)",
        source_crs.authid(),
        extent_4326.xMinimum(),
        extent_4326.yMinimum(),
        extent_4326.xMaximum(),
        extent_4326.yMaximum(),
    )
    return (
        extent_4326.xMinimum(),
        extent_4326.yMinimum(),
        extent_4326.xMaximum(),
        extent_4326.yMaximum(),
    )


def _qgis_dtype_to_numpy(qgis_dtype: Qgis.DataType) -> np.dtype:
    """Map QGIS raster data type to numpy dtype."""
    from qgis.core import Qgis

    mapping = {
        Qgis.DataType.Byte: np.uint8,
        Qgis.DataType.Int16: np.int16,
        Qgis.DataType.UInt16: np.uint16,
        Qgis.DataType.Int32: np.int32,
        Qgis.DataType.UInt32: np.uint32,
        Qgis.DataType.Float32: np.float32,
        Qgis.DataType.Float64: np.float64,
    }
    return np.dtype(mapping.get(qgis_dtype, np.uint8))


def _normalize_to_uint8(arr: np.ndarray) -> np.ndarray:
    """Min-max normalize array to 0-255 uint8."""
    arr_min = arr.min()
    arr_max = arr.max()
    if arr_max == arr_min:
        return np.zeros_like(arr, dtype=np.uint8)
    normalized = (arr - arr_min) / (arr_max - arr_min) * 255
    return normalized.astype(np.uint8)
