"""
Processing pipeline: raster generation and styling.
"""

import os

import processing
import numpy as np
from osgeo import gdal
from qgis.core import (
    QgsRasterLayer,
    QgsRasterShader, QgsColorRampShader,
    QgsSingleBandPseudoColorRenderer,
)
from qgis.PyQt.QtGui import QColor

# Rainbow ramp stops shared by apply_elevation_style and render_elevation_composite
_ELEV_RAMP = [
    (0.0,  (0x8B, 0x00, 0xFF)),
    (0.2,  (0x00, 0x00, 0xFF)),
    (0.4,  (0x00, 0xFF, 0xFF)),
    (0.6,  (0x00, 0xFF, 0x00)),
    (0.8,  (0xFF, 0xFF, 0x00)),
    (1.0,  (0xFF, 0x00, 0x00)),
]

_RDYLGN_RAMP = [
    '#a50026',
    '#d73027',
    '#f46d43',
    '#fdae61',
    '#fee08b',
    '#ffffbf',
    '#d9ef8b',
    '#a6d96a',
    '#66bd63',
    '#1a9850',
    '#006837',
]


def _single_band_percentile_range(path: str, pmin=2.0, pmax=98.0):
    ds = gdal.Open(path)
    if ds is None or ds.RasterCount != 1:
        return None

    band = ds.GetRasterBand(1)
    max_pixels = 2_000_000
    scale = max(1.0, (ds.RasterXSize * ds.RasterYSize / max_pixels) ** 0.5)
    buf_xsize = max(1, int(ds.RasterXSize / scale))
    buf_ysize = max(1, int(ds.RasterYSize / scale))
    arr = band.ReadAsArray(buf_xsize=buf_xsize, buf_ysize=buf_ysize).astype(np.float32)

    nodata = band.GetNoDataValue()
    valid = np.isfinite(arr)
    if nodata is not None:
        valid &= arr != nodata
    values = arr[valid]
    ds = None

    if values.size == 0:
        return None

    lo, hi = np.percentile(values, [pmin, pmax])
    lo, hi = float(lo), float(hi)
    if not np.isfinite(lo) or not np.isfinite(hi) or hi <= lo:
        return None
    return lo, hi


def generate_chm(dsm_path: str, dtm_path: str, output_path: str) -> str:
    """CHM = DSM - DTM。DSM を DTM グリッドに合わせてから演算する。"""
    dtm_ds = gdal.Open(dtm_path)
    gt = dtm_ds.GetGeoTransform()
    xmin = gt[0]
    ymax = gt[3]
    xmax = xmin + gt[1] * dtm_ds.RasterXSize
    ymin = ymax + gt[5] * dtm_ds.RasterYSize
    xres, yres = gt[1], abs(gt[5])
    srs = dtm_ds.GetProjection()
    dtm_ds = None

    tmp_dsm = output_path + '.tmp.tif'
    gdal.Warp(tmp_dsm, dsm_path,
              outputBounds=(xmin, ymin, xmax, ymax),
              xRes=xres, yRes=yres,
              dstSRS=srs, format='GTiff')

    processing.run('gdal:rastercalculator', {
        'INPUT_A': tmp_dsm, 'BAND_A': 1,
        'INPUT_B': dtm_path, 'BAND_B': 1,
        'FORMULA': 'A-B',
        'OUTPUT': output_path,
    })

    if os.path.isfile(tmp_dsm):
        os.remove(tmp_dsm)
    return output_path


def generate_vegetation_index(ortho_path: str, output_path: str) -> str:
    """VARI = (G - R) / (G + R - B) from RGB orthophoto (bands 1=R, 2=G, 3=B).
    nodata 領域は -9999 をセットし QGIS で透明描画させる。"""
    from osgeo import gdal
    _NODATA = -9999.0
    ds = gdal.Open(ortho_path)
    r = ds.GetRasterBand(1).ReadAsArray().astype(np.float32)
    g = ds.GetRasterBand(2).ReadAsArray().astype(np.float32)
    b = ds.GetRasterBand(3).ReadAsArray().astype(np.float32)

    # nodata マスク：アルファバンド（Band 4）優先、なければ nodata 値、最後に全ゼロ画素
    if ds.RasterCount >= 4 and ds.GetRasterBand(4).GetColorInterpretation() == gdal.GCI_AlphaBand:
        alpha = ds.GetRasterBand(4).ReadAsArray()
        nodata_mask = (alpha == 0)
    else:
        nodata_val = ds.GetRasterBand(1).GetNoDataValue()
        if nodata_val is not None:
            nodata_mask = (r == nodata_val)
        else:
            nodata_mask = (r == 0) & (g == 0) & (b == 0)

    with np.errstate(divide='ignore', invalid='ignore'):
        vari = np.nan_to_num((g - r) / (g + r - b)).astype(np.float32)
    vari[(vari > 1e30) | (vari < -1e30)] = _NODATA
    vari[nodata_mask] = _NODATA

    h, w = r.shape
    driver = gdal.GetDriverByName('GTiff')
    ds_out = driver.Create(output_path, w, h, 1, gdal.GDT_Float32)
    ds_out.SetGeoTransform(ds.GetGeoTransform())
    ds_out.SetProjection(ds.GetProjection())
    band = ds_out.GetRasterBand(1)
    band.WriteArray(vari)
    band.SetNoDataValue(_NODATA)
    ds_out.FlushCache()
    ds_out = ds = None
    return output_path


def generate_hillshade(dtm_path: str, output_path: str) -> str:
    """Hillshade from DTM."""
    processing.run('gdal:hillshade', {
        'INPUT': dtm_path,
        'BAND': 1,
        'Z_FACTOR': 1.0,
        'SCALE': 1.0,
        'AZIMUTH': 315.0,
        'ALTITUDE': 45.0,
        'OUTPUT': output_path,
    })
    return output_path


def render_elevation_composite(elev_path: str, hs_path: str, output_path: str) -> str:
    """Bake elevation color ramp + hillshade (Multiply 70%) into a single RGB GeoTIFF."""
    ds_elev = gdal.Open(elev_path)
    band = ds_elev.GetRasterBand(1)
    elev = band.ReadAsArray().astype(np.float32)
    nodata = band.GetNoDataValue()

    valid = elev if nodata is None else elev[elev != nodata]
    valid = valid[np.isfinite(valid)]
    if valid.size == 0:
        lo, hi = 0.0, 1.0
    else:
        lo, hi = float(np.nanmin(valid)), float(np.nanmax(valid))
    span = max(hi - lo, 1e-10)

    # nodata マスク（透明化対象）
    nodata_mask = ~np.isfinite(elev)
    if nodata is not None:
        nodata_mask |= (elev == nodata)

    # Normalize elevation to 0-1 and apply rainbow ramp
    t = np.clip((elev - lo) / span, 0.0, 1.0)
    h, w = elev.shape
    rgb = np.zeros((h, w, 3), dtype=np.float32)
    for i in range(len(_ELEV_RAMP) - 1):
        t0, c0 = _ELEV_RAMP[i]
        t1, c1 = _ELEV_RAMP[i + 1]
        in_seg = (t >= t0) & (t <= t1)
        alpha = np.where(in_seg, (t - t0) / max(t1 - t0, 1e-10), 0.0)
        for ch in range(3):
            rgb[:, :, ch] += np.where(in_seg, (1 - alpha) * c0[ch] + alpha * c1[ch], 0.0)

    # Hillshade Multiply at 70% opacity: result = base * (hs * 0.7 + 0.3)
    ds_hs = gdal.Open(hs_path)
    hs = ds_hs.GetRasterBand(1).ReadAsArray().astype(np.float32) / 255.0
    for ch in range(3):
        rgb[:, :, ch] = np.clip(rgb[:, :, ch] * (hs * 0.7 + 0.3), 0, 255)

    # Write RGBA GeoTIFF（nodata 領域はアルファ 0 で透明）
    driver = gdal.GetDriverByName('GTiff')
    ds_out = driver.Create(output_path, w, h, 4, gdal.GDT_Byte)
    ds_out.SetGeoTransform(ds_elev.GetGeoTransform())
    ds_out.SetProjection(ds_elev.GetProjection())
    for ch in range(3):
        ds_out.GetRasterBand(ch + 1).WriteArray(rgb[:, :, ch].astype(np.uint8))
    alpha_band = np.where(nodata_mask, 0, 255).astype(np.uint8)
    band4 = ds_out.GetRasterBand(4)
    band4.WriteArray(alpha_band)
    band4.SetColorInterpretation(gdal.GCI_AlphaBand)
    ds_out.FlushCache()
    ds_out = ds_hs = ds_elev = None
    return output_path


def apply_vegetation_style(layer: QgsRasterLayer) -> None:
    """Apply WebODM-like RdYlGn styling to a raw, single-band vegetation index."""
    source = layer.source().split('|', 1)[0]
    value_range = _single_band_percentile_range(source)
    if value_range is None:
        return

    lo, hi = value_range
    ramp = QgsColorRampShader(lo, hi)
    ramp.setColorRampType(QgsColorRampShader.Type.Interpolated)
    step = (hi - lo) / (len(_RDYLGN_RAMP) - 1)
    ramp.setColorRampItemList([
        QgsColorRampShader.ColorRampItem(lo + step * i, QColor(color))
        for i, color in enumerate(_RDYLGN_RAMP)
    ])
    shader = QgsRasterShader(lo, hi)
    shader.setRasterShaderFunction(ramp)
    renderer = QgsSingleBandPseudoColorRenderer(layer.dataProvider(), 1, shader)
    renderer.setClassificationMin(lo)
    renderer.setClassificationMax(hi)
    layer.setRenderer(renderer)
    layer.triggerRepaint()
