"""Shim module providing osgeo.gdal and osgeo.osr interfaces via rasterio.

This allows LiCSBAS to run on Windows without compiling GDAL Python bindings.
Only the subset of GDAL/OSR functions actually used by LiCSBAS is implemented.

Used functions:
  gdal.Open(filename) -> Dataset
  gdal.GetDriverByName('GTiff') -> Driver
  gdal.Warp("", src, format='MEM', ...) -> Dataset
  gdal.Info(filename, ...) -> str/dict
  gdal.GDT_Float32, gdal.GDT_Byte
  gdal.DEMProcessing (not implemented — only used by color_geotiff)
  osr.SpatialReference()
"""

import numpy as np
import rasterio
from rasterio.transform import from_bounds
from rasterio.warp import reproject, Resampling, calculate_default_transform
from rasterio.crs import CRS
import json

# GDAL data type constants
GDT_Float32 = "float32"
GDT_Byte = "uint8"
GDT_Float64 = "float64"
GDT_Int16 = "int16"
GDT_Int32 = "int32"
GDT_UInt16 = "uint16"


class SpatialReference:
    """Minimal osr.SpatialReference shim."""

    def __init__(self):
        self._epsg = None

    def ImportFromEPSG(self, code):
        self._epsg = code

    def ExportToWkt(self):
        if self._epsg:
            return CRS.from_epsg(self._epsg).to_wkt()
        return ""


class Band:
    """Shim for gdal.Band."""

    def __init__(self, data, nodata=None, parent_dataset=None):
        self._data = data
        self._nodata = nodata
        self._parent = parent_dataset

    def ReadAsArray(self):
        return self._data

    def WriteArray(self, array):
        self._data[:] = array  # in-place copy to preserve reference in Dataset._bands_data

    def SetNoDataValue(self, val):
        self._nodata = val
        if self._parent:
            self._parent._nodata = val

    def GetNoDataValue(self):
        return self._nodata

    def FlushCache(self):
        if self._parent:
            self._parent.FlushCache()


class Dataset:
    """Shim for gdal.Dataset — wraps rasterio for read/write."""

    def __init__(self, filename=None, array=None, transform=None, crs=None, nodata=None):
        self._filename = filename
        self._transform = transform
        self._crs = crs
        self._nodata = nodata
        self._bands_data = []

        if filename and array is None:
            with rasterio.open(filename) as src:
                self._width = src.width
                self._height = src.height
                self._count = src.count
                self._transform = src.transform
                self._crs = src.crs
                self._nodata = src.nodata
                for i in range(1, src.count + 1):
                    self._bands_data.append(src.read(i))
        elif array is not None:
            if array.ndim == 2:
                self._bands_data = [array]
                self._height, self._width = array.shape
                self._count = 1
            else:
                self._bands_data = [array[i] for i in range(array.shape[0])]
                self._count = array.shape[0]
                self._height, self._width = array.shape[1], array.shape[2]

    def ReadAsArray(self):
        if self._count == 1:
            return self._bands_data[0].copy()
        return np.array(self._bands_data)

    def GetRasterBand(self, idx):
        return Band(self._bands_data[idx - 1], self._nodata, parent_dataset=self)

    def GetGeoTransform(self):
        t = self._transform
        return (t.c, t.a, t.b, t.f, t.d, t.e)

    def GetProjection(self):
        if self._crs:
            return self._crs.to_wkt()
        return ""

    def SetGeoTransform(self, gt):
        from rasterio.transform import Affine
        self._transform = Affine(gt[1], gt[2], gt[0], gt[4], gt[5], gt[3])

    def SetProjection(self, wkt):
        self._crs = CRS.from_wkt(wkt)

    @property
    def RasterXSize(self):
        return self._width

    @property
    def RasterYSize(self):
        return self._height

    @property
    def RasterCount(self):
        return self._count

    def SetMetadataItem(self, key, value, domain=""):
        """Store metadata (written to file on FlushCache if writable)."""
        if not hasattr(self, "_metadata"):
            self._metadata = {}
        self._metadata[key] = value

    def FlushCache(self):
        """Write dataset to disk if it was created via Driver.Create()."""
        if getattr(self, "_write_on_flush", False) and self._filename:
            _dtype_map = {
                "float32": rasterio.float32,
                "float64": rasterio.float64,
                "uint8": rasterio.uint8,
                "int16": rasterio.int16,
                "int32": rasterio.int32,
                "uint16": rasterio.uint16,
            }
            np_dtype = str(self._bands_data[0].dtype)
            rio_dtype = _dtype_map.get(np_dtype, rasterio.float32)

            profile = {
                "driver": "GTiff",
                "width": self._width,
                "height": self._height,
                "count": self._count,
                "dtype": rio_dtype,
                "crs": self._crs,
                "transform": self._transform,
            }
            if self._nodata is not None:
                profile["nodata"] = self._nodata

            # Apply compression if available
            if hasattr(self, "_create_options"):
                for opt in self._create_options:
                    if "=" in opt:
                        k, v = opt.split("=", 1)
                        profile[k.lower()] = v

            with rasterio.open(self._filename, "w", **profile) as dst:
                for i, band_data in enumerate(self._bands_data):
                    dst.write(band_data, i + 1)


class Driver:
    """Shim for gdal.Driver."""

    def __init__(self, name):
        self._name = name

    def Create(self, filename, width, height, bands=1, dtype=GDT_Float32, options=None):
        ds = Dataset()
        ds._filename = filename
        ds._width = width
        ds._height = height
        ds._count = bands

        np_dtype = np.float32
        if dtype == GDT_Byte or dtype == "uint8":
            np_dtype = np.uint8

        ds._bands_data = [np.zeros((height, width), dtype=np_dtype) for _ in range(bands)]
        ds._write_on_flush = True
        if options:
            ds._create_options = options
        return ds


def Open(filename, mode=0):
    """Shim for gdal.Open()."""
    try:
        return Dataset(filename)
    except Exception:
        return None


def GetDriverByName(name):
    """Shim for gdal.GetDriverByName()."""
    return Driver(name)


def Warp(dest, src, format="GTiff", outputBounds=None, width=None, height=None,
         resampleAlg=None, srcNodata=None, dstNodata=None, **kwargs):
    """Shim for gdal.Warp() — reproject/clip a raster."""
    resamp_map = {
        None: Resampling.nearest,
        "near": Resampling.nearest,
        "bilinear": Resampling.bilinear,
        "cubic": Resampling.cubic,
        "lanczos": Resampling.lanczos,
    }
    resampling = resamp_map.get(resampleAlg, Resampling.nearest)

    if isinstance(src, str):
        src_ds = Dataset(src)
    elif isinstance(src, Dataset):
        src_ds = src
    else:
        src_ds = Dataset(src)

    src_data = src_ds.ReadAsArray()
    src_transform = src_ds._transform
    src_crs = src_ds._crs or CRS.from_epsg(4326)

    if outputBounds and width and height:
        lon_w, lat_s, lon_e, lat_n = outputBounds
        dst_transform = from_bounds(lon_w, lat_s, lon_e, lat_n, width, height)
        dst_crs = src_crs

        if src_data.ndim == 2:
            dst_data = np.full((height, width), np.nan, dtype=np.float32)
            reproject(
                source=src_data.astype(np.float32),
                destination=dst_data,
                src_transform=src_transform,
                src_crs=src_crs,
                dst_transform=dst_transform,
                dst_crs=dst_crs,
                resampling=resampling,
                src_nodata=srcNodata,
                dst_nodata=dstNodata,
            )
        else:
            dst_data = np.full((src_data.shape[0], height, width), np.nan, dtype=np.float32)
            for i in range(src_data.shape[0]):
                reproject(
                    source=src_data[i].astype(np.float32),
                    destination=dst_data[i],
                    src_transform=src_transform,
                    src_crs=src_crs,
                    dst_transform=dst_transform,
                    dst_crs=dst_crs,
                    resampling=resampling,
                    src_nodata=srcNodata,
                    dst_nodata=dstNodata,
                )

        result = Dataset(array=dst_data, transform=dst_transform, crs=dst_crs)
        return result
    else:
        return src_ds


def Info(filename, computeMinMax=False, format=None, **kwargs):
    """Shim for gdal.Info()."""
    try:
        with rasterio.open(filename) as src:
            if format == "json":
                info = {
                    "size": [src.width, src.height],
                    "bands": [],
                }
                for i in range(1, src.count + 1):
                    band_info = {"band": i}
                    if computeMinMax:
                        data = src.read(i)
                        valid = data[~np.isnan(data)] if np.issubdtype(data.dtype, np.floating) else data
                        if len(valid) > 0:
                            band_info["computedMin"] = float(np.min(valid))
                            band_info["computedMax"] = float(np.max(valid))
                    info["bands"].append(band_info)
                return info
            return f"Driver: GTiff\nSize: {src.width}x{src.height}\nBands: {src.count}"
    except Exception:
        return None


def IdentifyDriver(filename):
    """Shim for gdal.IdentifyDriver()."""
    try:
        with rasterio.open(filename) as src:
            return True
    except Exception:
        return None


def VersionInfo():
    """Return GDAL version string."""
    return rasterio.gdal_version().replace(".", "0")
