from __future__ import annotations

import csv
import glob
import json
import math
import os
from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np
import pdal
from osgeo import gdal, ogr, osr
from scipy import ndimage

from .config import Config, LakeConfig, RiverErosionProfileConfig
from .gdal_utils import ensure_dir, ensure_parent, open_dataset

gdal.UseExceptions()


def _river_debug_enabled() -> bool:
    flag = os.getenv("RIVER_EROSION_DEBUG", "")
    return flag.strip().lower() in {"1", "true", "yes", "y"}


def _river_debug(message: str) -> None:
    if _river_debug_enabled():
        print(f"[river_erosion:debug] {message}")


def _parse_tile_key(tile_key: str) -> Optional[Tuple[int, int]]:
    if not tile_key:
        return None
    parts = tile_key.split("_")
    if len(parts) != 2:
        return None
    try:
        return int(parts[0]), int(parts[1])
    except ValueError:
        return None


def _tile_geotransform(bounds: Tuple[float, float, float, float], width: int, height: int) -> Tuple[float, float, float, float, float, float]:
    xmin, ymin, xmax, ymax = bounds
    xdenom = float(width - 1) if width > 1 else float(width)
    ydenom = float(height - 1) if height > 1 else float(height)
    xres = (xmax - xmin) / xdenom
    yres = (ymax - ymin) / ydenom
    return xmin, xres, 0.0, ymax, 0.0, -yres


def _pixel_size_m(gt: Tuple[float, ...]) -> float:
    return (abs(gt[1]) + abs(gt[5])) / 2.0


def _bounds_to_layer(
    bounds: Tuple[float, float, float, float],
    buffer_m: float,
    to_layer: osr.CoordinateTransformation,
) -> Tuple[float, float, float, float]:
    xmin, ymin, xmax, ymax = bounds
    xmin -= buffer_m
    ymin -= buffer_m
    xmax += buffer_m
    ymax += buffer_m
    corners = [(xmin, ymin), (xmin, ymax), (xmax, ymin), (xmax, ymax)]
    xs: list[float] = []
    ys: list[float] = []
    for x, y in corners:
        lon, lat, _ = to_layer.TransformPoint(x, y)
        xs.append(lon)
        ys.append(lat)
    return min(xs), min(ys), max(xs), max(ys)


def _normalize_order(value, profile: RiverErosionProfileConfig) -> Optional[int]:
    if value is None:
        return None
    try:
        order = int(value)
    except (TypeError, ValueError):
        return None
    if profile.invert_order:
        max_order = max(1, int(profile.invert_max_order))
        if order <= 0:
            order = max_order
        order = max(1, min(max_order, order))
        order = max_order + 1 - order
    return order


def _order_depth(order: int, profile: RiverErosionProfileConfig) -> Optional[float]:
    if order < profile.min_order:
        return None
    delta = order - profile.min_order
    depth = profile.depth_base_m + profile.depth_per_order_m * delta
    depth = min(max(depth, profile.depth_min_m), profile.depth_max_m)
    return depth


def _refine_mask(
    mask: np.ndarray,
    fill_radius: int = 3,
    smooth_sigma: float = 2.0,
) -> np.ndarray:
    if fill_radius > 0:
        struct = ndimage.generate_binary_structure(2, 1)
        mask = ndimage.binary_closing(mask, structure=struct, iterations=fill_radius).astype(np.float32)
    if smooth_sigma > 0:
        mask = ndimage.gaussian_filter(mask, sigma=smooth_sigma)
    return mask


def _build_water_mask_mosaic(
    laz_paths: Iterable[str],
    bounds: Tuple[float, float, float, float],
    width: int,
    height: int,
    water_class: int = 9,
) -> Optional[np.ndarray]:
    laz_paths = [path for path in laz_paths if path and os.path.exists(path)]
    if not laz_paths:
        return None

    xmin, ymin, xmax, ymax = bounds
    pipeline_json: list[dict] = []
    for path in laz_paths:
        pipeline_json.append({
            "type": "readers.las",
            "filename": path,
        })
    if len(laz_paths) > 1:
        pipeline_json.append({"type": "filters.merge"})
    pipeline_json.append({
        "type": "filters.crop",
        "bounds": f"([{xmin},{xmax}],[{ymin},{ymax}])",
    })
    pipeline_json.append({
        "type": "filters.range",
        "limits": f"Classification[{water_class}:{water_class}]",
    })

    try:
        pipeline = pdal.Pipeline(json.dumps(pipeline_json))
        count = pipeline.execute()
    except RuntimeError as exc:
        print(f"[river_erosion] PDAL error building water mask: {exc}")
        return None

    if count == 0:
        return None

    arrays = pipeline.arrays
    if not arrays:
        return None

    if len(arrays) == 1:
        xs = arrays[0]["X"]
        ys = arrays[0]["Y"]
    else:
        xs = np.concatenate([arr["X"] for arr in arrays])
        ys = np.concatenate([arr["Y"] for arr in arrays])

    xdenom = float(width - 1) if width > 1 else float(width)
    ydenom = float(height - 1) if height > 1 else float(height)
    xres = (xmax - xmin) / xdenom
    yres = (ymax - ymin) / ydenom

    col = ((xs - xmin) / xres).astype(int)
    row = ((ymax - ys) / yres).astype(int)

    valid = (col >= 0) & (col < width) & (row >= 0) & (row < height)
    if not np.any(valid):
        return None

    mask = np.zeros((height, width), dtype=np.float32)
    mask[row[valid], col[valid]] = 1.0
    return mask


def _layer_has_field(layer: ogr.Layer, field_name: Optional[str]) -> bool:
    if not field_name:
        return False
    layer_defn = layer.GetLayerDefn()
    if layer_defn is None:
        return False
    return layer_defn.GetFieldIndex(field_name) != -1


def _build_order_lines(
    layer_sources: List[Tuple[ogr.Layer, str]],
    bounds: Tuple[float, float, float, float],
    width: int,
    height: int,
    profile: RiverErosionProfileConfig,
    to_tile: osr.CoordinateTransformation,
    to_layer: osr.CoordinateTransformation,
    tile_srs: osr.SpatialReference,
) -> np.ndarray:
    gt = _tile_geotransform(bounds, width, height)
    driver = gdal.GetDriverByName("MEM")
    order_ds = driver.Create("", width, height, 1, gdal.GDT_Float32)
    order_ds.SetGeoTransform(gt)
    order_ds.SetProjection(tile_srs.ExportToWkt())
    band = order_ds.GetRasterBand(1)
    band.Fill(0.0)

    if not any(_layer_has_field(layer, profile.order_field) for layer, _ in layer_sources):
        _river_debug(f"Order field '{profile.order_field}' not found in HydroRIVERS layer(s).")
        return band.ReadAsArray().astype(np.float32)

    buffer_m = max(profile.smooth_sigma_m * 3.0, 250.0)
    minx, miny, maxx, maxy = _bounds_to_layer(bounds, buffer_m, to_layer)

    buffered_geoms: list[ogr.Geometry] = []
    burn_values: list[float] = []

    for layer, _ in layer_sources:
        layer.SetSpatialFilterRect(minx, miny, maxx, maxy)
        layer.ResetReading()
        for feat in layer:
            order = _normalize_order(feat.GetField(profile.order_field), profile)
            if order is None:
                continue
            geom = feat.GetGeometryRef()
            if geom is None:
                continue
            geom = geom.Clone()
            geom.Transform(to_tile)
            if geom.IsEmpty():
                continue
            buffered_geoms.append(geom)
            burn_values.append(float(order))

    if not buffered_geoms:
        return band.ReadAsArray().astype(np.float32)

    ordered = sorted(zip(buffered_geoms, burn_values), key=lambda item: item[1])
    ordered_geoms = [geom for geom, _ in ordered]
    ordered_values = [float(value) for _, value in ordered]

    if hasattr(gdal, "RasterizeGeometries"):
        gdal.RasterizeGeometries(
            order_ds,
            [1],
            ordered_geoms,
            burn_values=ordered_values,
            options=["MERGE_ALG=REPLACE", "ALL_TOUCHED=TRUE"],
        )
    else:
        mem_driver = ogr.GetDriverByName("MEM") or ogr.GetDriverByName("Memory")
        if mem_driver is None:
            raise SystemExit("[river_erosion] OGR MEM driver missing; update GDAL.")
        mem_ds = mem_driver.CreateDataSource("")
        mem_layer = mem_ds.CreateLayer("burn", srs=tile_srs, geom_type=ogr.wkbLineString)
        mem_layer.CreateField(ogr.FieldDefn("burn", ogr.OFTReal))
        layer_defn = mem_layer.GetLayerDefn()
        for geom, value in zip(ordered_geoms, ordered_values):
            feature = ogr.Feature(layer_defn)
            feature.SetField("burn", value)
            feature.SetGeometry(geom)
            mem_layer.CreateFeature(feature)
            feature = None
        gdal.RasterizeLayer(
            order_ds,
            [1],
            mem_layer,
            options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"],
        )
        mem_ds = None

    return band.ReadAsArray().astype(np.float32)


def _build_flow_distance(
    layer_sources: List[Tuple[ogr.Layer, str]],
    bounds: Tuple[float, float, float, float],
    width: int,
    height: int,
    dist_field: str,
    to_tile: osr.CoordinateTransformation,
    to_layer: osr.CoordinateTransformation,
    tile_srs: osr.SpatialReference,
) -> np.ndarray:
    gt = _tile_geotransform(bounds, width, height)
    driver = gdal.GetDriverByName("MEM")
    dist_ds = driver.Create("", width, height, 1, gdal.GDT_Float32)
    dist_ds.SetGeoTransform(gt)
    dist_ds.SetProjection(tile_srs.ExportToWkt())
    band = dist_ds.GetRasterBand(1)
    band.Fill(0.0)

    if not any(_layer_has_field(layer, dist_field) for layer, _ in layer_sources):
        _river_debug(f"Distance field '{dist_field}' not found in HydroRIVERS layer(s).")
        return band.ReadAsArray().astype(np.float32)

    buffer_m = 50.0
    minx, miny, maxx, maxy = _bounds_to_layer(bounds, buffer_m, to_layer)

    buffered_geoms: list[ogr.Geometry] = []
    burn_values: list[float] = []

    for layer, _ in layer_sources:
        layer.SetSpatialFilterRect(minx, miny, maxx, maxy)
        layer.ResetReading()
        for feat in layer:
            dist_val = feat.GetField(dist_field)
            if dist_val is None:
                continue
            try:
                dist_km = float(dist_val)
            except (TypeError, ValueError):
                continue
            geom = feat.GetGeometryRef()
            if geom is None:
                continue
            geom = geom.Clone()
            geom.Transform(to_tile)
            if geom.IsEmpty():
                continue
            buffered_geoms.append(geom)
            burn_values.append(dist_km)

    if not buffered_geoms:
        return band.ReadAsArray().astype(np.float32)

    ordered = sorted(zip(buffered_geoms, burn_values), key=lambda item: item[1])
    ordered_geoms = [geom for geom, _ in ordered]
    ordered_values = [float(value) for _, value in ordered]

    if hasattr(gdal, "RasterizeGeometries"):
        gdal.RasterizeGeometries(
            dist_ds,
            [1],
            ordered_geoms,
            burn_values=ordered_values,
            options=["MERGE_ALG=REPLACE", "ALL_TOUCHED=TRUE"],
        )
    else:
        mem_driver = ogr.GetDriverByName("MEM") or ogr.GetDriverByName("Memory")
        if mem_driver is None:
            raise SystemExit("[river_erosion] OGR MEM driver missing; update GDAL.")
        mem_ds = mem_driver.CreateDataSource("")
        mem_layer = mem_ds.CreateLayer("burn", srs=tile_srs, geom_type=ogr.wkbLineString)
        mem_layer.CreateField(ogr.FieldDefn("burn", ogr.OFTReal))
        layer_defn = mem_layer.GetLayerDefn()
        for geom, value in zip(ordered_geoms, ordered_values):
            feature = ogr.Feature(layer_defn)
            feature.SetField("burn", value)
            feature.SetGeometry(geom)
            mem_layer.CreateFeature(feature)
            feature = None
        gdal.RasterizeLayer(
            dist_ds,
            [1],
            mem_layer,
            options=["ATTRIBUTE=burn", "ALL_TOUCHED=TRUE"],
        )
        mem_ds = None

    return band.ReadAsArray().astype(np.float32)


def _fill_nan_with_nearest(data: np.ndarray) -> np.ndarray:
    if not np.isnan(data).any():
        return data
    valid = np.isfinite(data)
    if not np.any(valid):
        return np.zeros_like(data, dtype=np.float32)
    _, indices = ndimage.distance_transform_edt(~valid, return_indices=True)
    filled = data[tuple(indices)]
    return filled.astype(np.float32)


def _resize_mask(mask: np.ndarray, height: int, width: int) -> np.ndarray:
    if mask.shape == (height, width):
        return mask
    zoom_y = height / mask.shape[0]
    zoom_x = width / mask.shape[1]
    resized = ndimage.zoom(mask, (zoom_y, zoom_x), order=1)
    if resized.shape[0] != height or resized.shape[1] != width:
        resized = resized[:height, :width]
        pad_y = max(0, height - resized.shape[0])
        pad_x = max(0, width - resized.shape[1])
        if pad_y or pad_x:
            resized = np.pad(resized, ((0, pad_y), (0, pad_x)), mode="edge")
    return resized.astype(np.float32)


def _apply_water_mask_to_ortho(tile_id: str, mask: np.ndarray, cfg: Config) -> None:
    if not cfg.ortho.apply_water_mask:
        return
    ortho_path = os.path.join(cfg.export.ortho_dir, f"{tile_id}.jpg")
    if not os.path.exists(ortho_path):
        return
    river_dir = getattr(cfg.ortho, "river_dir", cfg.export.ortho_dir)
    ensure_dir(river_dir)
    out_path = os.path.join(river_dir, f"{tile_id}.jpg")

    ds = gdal.Open(ortho_path, gdal.GA_ReadOnly)
    if ds is None:
        print(f"[river_erosion] Could not open ortho {ortho_path}")
        return

    width = ds.RasterXSize
    height = ds.RasterYSize
    if width <= 0 or height <= 0:
        ds = None
        return

    data = ds.ReadAsArray()
    if data is None:
        ds = None
        return

    if data.ndim == 2:
        data = np.stack([data, data, data], axis=0)

    if data.shape[0] < 3:
        ds = None
        return

    rgb = data[:3].astype(np.float32)
    mask_res = _resize_mask(mask, height, width)
    mask_res = np.clip(mask_res, 0.0, 1.0)
    threshold = float(cfg.ortho.water_mask_threshold)
    water_mask = mask_res >= threshold
    trans_opts = gdal.TranslateOptions(
        outputType=gdal.GDT_Byte,
        format="JPEG",
        creationOptions=[f"QUALITY={cfg.ortho.jpeg_quality}", "WORLDFILE=YES"],
    )
    if not np.any(water_mask):
        try:
            gdal.Translate(out_path, ds, options=trans_opts)
        except RuntimeError as exc:
            print(f"[river_erosion] Ortho write failed for {tile_id}: {exc}")
        ds = None
        print(f"[river_erosion] Wrote ortho {out_path} (copied; empty mask)")
        return

    mode = str(getattr(cfg.ortho, "water_color_mode", "median") or "median").lower()
    if mode == "fixed":
        water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32)
    else:
        if water_mask.any():
            sample = rgb[:, water_mask]
            if sample.size >= 3:
                water_color = np.median(sample, axis=1)
            else:
                water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32)
        else:
            water_color = np.array(cfg.ortho.water_fallback_rgb, dtype=np.float32)

    blend = float(cfg.ortho.water_blend)
    weights = mask_res * blend
    for c in range(3):
        rgb[c] = rgb[c] * (1.0 - weights) + water_color[c] * weights
    rgb = np.clip(rgb, 0.0, 255.0).astype(np.uint8)

    mem_driver = gdal.GetDriverByName("MEM")
    out_ds = mem_driver.Create("", width, height, 3, gdal.GDT_Byte)
    out_ds.SetGeoTransform(ds.GetGeoTransform())
    out_ds.SetProjection(ds.GetProjection())
    for c in range(3):
        out_ds.GetRasterBand(c + 1).WriteArray(rgb[c])

    try:
        gdal.Translate(out_path, out_ds, options=trans_opts)
    except RuntimeError as exc:
        print(f"[river_erosion] Ortho write failed for {tile_id}: {exc}")
    out_ds = None
    ds = None
    print(f"[river_erosion] Wrote ortho {out_path}")


def _write_manifest(
    manifest_path: str,
    rows: Iterable[Dict[str, str]],
    global_min: float,
    global_max: float,
) -> None:
    ensure_parent(manifest_path)
    header = [
        "tile_id",
        "xmin",
        "ymin",
        "xmax",
        "ymax",
        "global_min",
        "global_max",
        "out_res",
        "tile_key",
        "tile_min",
        "tile_max",
    ]
    with open(manifest_path, "w", encoding="utf-8", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=header)
        writer.writeheader()
        for row in rows:
            out = dict(row)
            out["global_min"] = f"{global_min:.6f}"
            out["global_max"] = f"{global_max:.6f}"
            writer.writerow(out)


def _write_mask_png(
    out_path: str,
    mask: np.ndarray,
    bounds: Tuple[float, float, float, float],
    tile_srs: osr.SpatialReference,
    *,
    bands: int = 1,
) -> None:
    ensure_parent(out_path)
    height, width = mask.shape[-2], mask.shape[-1]
    gt = _tile_geotransform(bounds, width, height)
    driver = gdal.GetDriverByName("MEM")
    ds = driver.Create("", width, height, bands, gdal.GDT_Byte)
    ds.SetGeoTransform(gt)
    ds.SetProjection(tile_srs.ExportToWkt())
    if bands == 1:
        ds.GetRasterBand(1).WriteArray(mask.astype(np.uint8))
    else:
        for idx in range(bands):
            ds.GetRasterBand(idx + 1).WriteArray(mask[idx].astype(np.uint8))
    options = gdal.TranslateOptions(format="PNG", creationOptions=["WORLDFILE=YES"])
    try:
        gdal.Translate(out_path, ds, options=options)
    except RuntimeError as exc:
        print(f"[river_erosion] Mask write failed for {out_path}: {exc}")
    ds = None


def _resolve_hydrorivers_path(source_path: str) -> Optional[str]:
    if source_path and os.path.exists(source_path):
        return source_path
    fallback = os.path.join(
        "raw",
        "hydrorivers",
        "hydrorivers_eu_shp",
        "HydroRIVERS_v10_eu_shp",
        "HydroRIVERS_v10_eu.shp",
    )
    if os.path.exists(fallback):
        return fallback
    if source_path and os.path.isdir(source_path):
        candidates = glob.glob(os.path.join(source_path, "**", "*.shp"), recursive=True)
        if candidates:
            candidates.sort()
            for cand in candidates:
                if "hydrorivers" in cand.lower():
                    return cand
            return candidates[0]
    return None


def _laz_path_for_tile(tile_id: str, source_dir: str) -> Optional[str]:
    parts = tile_id.split("_")
    if len(parts) < 6:
        return None
    x_idx = parts[2]
    y_idx = parts[3]
    laz_name = f"bdom20rgbi_32_{x_idx}_{y_idx}_2_rp.laz"
    laz_path = os.path.join(source_dir, laz_name)
    return laz_path if os.path.exists(laz_path) else None


def _find_lake_depths(
    manual_geoms: List[ogr.Geometry],
    manual_srs: osr.SpatialReference,
    lake_cfg: LakeConfig,
    hydrolakes_path: str,
) -> List[Tuple[float, Dict[str, Any]]]:
    results = []
    default_depth = lake_cfg.default_depth_m
    tolerance = lake_cfg.match_tolerance_m

    if not manual_geoms:
        return []

    if not os.path.exists(hydrolakes_path):
        _river_debug(f"HydroLAKES path not found: {hydrolakes_path}")
        return [(default_depth, {"error": "HydroLAKES missing"}) for _ in manual_geoms]

    ds = ogr.Open(hydrolakes_path)
    if ds is None:
        return [(default_depth, {"error": "Could not open HydroLAKES"}) for _ in manual_geoms]
    
    layer = ds.GetLayer(0)
    layer_srs = layer.GetSpatialRef()
    
    # Transformation from Manual (Tile) to HydroLAKES
    transform = None
    if not manual_srs.IsSame(layer_srs):
        transform = osr.CoordinateTransformation(manual_srs, layer_srs)

    # Reverse Transform (HydroLAKES -> Manual) for distance check
    rev_transform = None
    if not manual_srs.IsSame(layer_srs):
        rev_transform = osr.CoordinateTransformation(layer_srs, manual_srs)

    for geom in manual_geoms:
        # Working geometry in HydroLAKES CRS for spatial filter
        search_geom = geom.Clone()
        if transform:
            search_geom.Transform(transform)
        
        centroid = search_geom.Centroid()
        pt = centroid.GetPoint() # (lon, lat)
        
        # Simple spatial filter around centroid (approx 0.1 deg ~ 10km)
        search_buffer = 0.1 
        layer.SetSpatialFilterRect(pt[0] - search_buffer, pt[1] - search_buffer, pt[0] + search_buffer, pt[1] + search_buffer)
        
        best_dist = float("inf")
        best_depth = default_depth
        best_id = -1
        matched = False
        
        layer.ResetReading()
        for feat in layer:
            feat_geom = feat.GetGeometryRef()
            if feat_geom is None:
                continue
                
            # Clone and transform feature to Manual CRS (Meters) for accurate distance
            feat_geom_m = feat_geom.Clone()
            if rev_transform:
                feat_geom_m.Transform(rev_transform)
                
            dist = feat_geom_m.Distance(geom)
            
            if dist < best_dist:
                best_dist = dist
                best_id = feat.GetFID()
                # Get depth
                d = feat.GetField("Depth_avg")
                if d is not None and float(d) > 0:
                    best_depth = float(d)
                else:
                    best_depth = default_depth
        
        if best_dist <= tolerance:
            matched = True
            results.append((best_depth, {"id": best_id, "distance": best_dist, "matched": True}))
        else:
            # If closest is too far, use default
            results.append((default_depth, {"id": best_id, "distance": best_dist, "matched": False, "note": "Outside tolerance"}))
            
    return results


def _load_manual_mask(
    tile_id: str, 
    search_dir: str = "raw/water_masks"
) -> Tuple[Optional[np.ndarray], Optional[np.ndarray], Optional[np.ndarray]]:
    if not os.path.exists(search_dir):
        return None, None, None
        
    # Try exact match first
    # Expected pattern: {tile_id}_mask_viz.png or similar
    # User file: dgm01_32_324_5507_rp_mask_viz.png
    # Tile ID:   dgm01_32_324_5507_1_rp
    
    # Heuristic: try to construct the filename from tile_id parts
    # If tile_id has "_1_rp", try removing "_1"
    
    candidates = [
        f"{tile_id}.png",
        f"{tile_id}_mask.png",
        f"{tile_id}_mask_viz.png",
        f"{tile_id}_viz.png"
    ]
    
    # Try stripping "_1"
    base_id = tile_id.replace("_1_rp", "_rp")
    candidates.extend([
        f"{base_id}.png",
        f"{base_id}_mask.png",
        f"{base_id}_mask_viz.png",
        f"{base_id}_viz.png"
    ])
    
    match_path = None
    for cand in candidates:
        path = os.path.join(search_dir, cand)
        if os.path.exists(path):
            match_path = path
            break
            
    if not match_path:
        return None, None, None
        
    ds = gdal.Open(match_path)
    if ds is None:
        return None, None, None
        
    # Expecting RGB
    if ds.RasterCount < 3:
        return None, None, None
        
    red = ds.GetRasterBand(1).ReadAsArray()
    green = ds.GetRasterBand(2).ReadAsArray()
    blue = ds.GetRasterBand(3).ReadAsArray()
    
    # Thresholds (assuming UInt8 0-255 or UInt16 0-65535)
    # If UInt16, max is 65535.
    # Check max value to guess? Or just > 0.
    
    # Semantic:
    # Blue = Lake (Water)
    # Green = Island (No Erosion)
    # Red = Bridge (No Erosion + Fix)
    
    # Masks are boolean (1 where color is present)
    mask_r = (red > 0)
    mask_g = (green > 0)
    mask_b = (blue > 0)
    
    return mask_b, mask_g, mask_r


def erode_rivers(cfg: Config) -> int:
    re_cfg = cfg.river_erosion
    if not os.path.exists(cfg.export.manifest_path):
        print(f"[river_erosion] Missing manifest: {cfg.export.manifest_path}. Run heightmap export first.")
        return 1

    source_path = _resolve_hydrorivers_path(re_cfg.source_path)
    if not source_path:
        print(f"[river_erosion] Missing HydroRIVERS source: {re_cfg.source_path}")
        return 1

    if not re_cfg.enabled:
        print("[river_erosion] Warning: river_erosion.enabled=false; proceeding because --erode-rivers was requested.")

    try:
        ds = ogr.Open(source_path)
    except RuntimeError as exc:
        print(f"[river_erosion] Could not open {source_path} ({exc})")
        return 1
    if ds is None:
        print(f"[river_erosion] Could not open {source_path}")
        return 1

    layer = ds.GetLayerByName(re_cfg.source_layer) if re_cfg.source_layer else None
    if layer is None:
        layer = ds.GetLayer(0)
    if layer is None:
        print(f"[river_erosion] Could not find layer in {source_path}")
        return 1

    layer_sources = [(layer, os.path.basename(source_path))]

    layer_srs = layer.GetSpatialRef()
    if layer_srs is None:
        print("[river_erosion] Warning: HydroRIVERS layer CRS missing; assuming EPSG:4326.")
        layer_srs = osr.SpatialReference()
        layer_srs.ImportFromEPSG(4326)
    tile_srs = osr.SpatialReference()
    tile_srs.ImportFromEPSG(25832)
    layer_srs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
    tile_srs.SetAxisMappingStrategy(osr.OAMS_TRADITIONAL_GIS_ORDER)
    to_tile = osr.CoordinateTransformation(layer_srs, tile_srs)
    to_layer = osr.CoordinateTransformation(tile_srs, layer_srs)

    with open(cfg.export.manifest_path, newline="", encoding="utf-8") as f:
        reader = csv.DictReader(f)
        if not reader.fieldnames:
            print(f"[river_erosion] Empty manifest: {cfg.export.manifest_path}")
            return 1
        required = {"tile_id", "xmin", "ymin", "xmax", "ymax", "global_min", "global_max", "out_res"}
        missing = required - set(reader.fieldnames)
        if missing:
            print(f"[river_erosion] Manifest missing columns: {', '.join(sorted(missing))}")
            return 1
        rows = list(reader)

    if not rows:
        print(f"[river_erosion] No tiles in manifest: {cfg.export.manifest_path}")
        return 1

    key_lookup: Dict[str, Dict[str, str]] = {}
    for row in rows:
        tile_id = row["tile_id"].strip()
        tile_key = row.get("tile_key", "").strip()
        if tile_key:
            key_lookup[tile_key] = row

    lidar_cfg = getattr(re_cfg, "lidar", None)
    source_dir = getattr(lidar_cfg, "source_dir", "raw/bdom20rgbi")
    water_class = getattr(lidar_cfg, "classification_water", 9)
    depth_profile = re_cfg.vr
    fill_radius = getattr(lidar_cfg, "fill_holes_radius", 3)
    smooth_sigma = getattr(lidar_cfg, "bank_slope_sigma", depth_profile.smooth_sigma_m)

    output_dir = re_cfg.output_dir
    manifest_out = re_cfg.manifest_vr
    ensure_dir(output_dir)

    rows_out: List[Dict[str, str]] = []
    global_min = math.inf
    global_max = -math.inf

    for row in rows:
        tile_id = row["tile_id"].strip()
        tile_key = row.get("tile_key", "").strip()
        center_key = _parse_tile_key(tile_key)
        if center_key is None:
            _river_debug(f"Tile {tile_id} missing/invalid tile_key; processing without neighbors.")
            neighbors = [(0, 0, row)]
        else:
            neighbors = []
            for dy in (-1, 0, 1):
                for dx in (-1, 0, 1):
                    neighbor_key = f"{center_key[0] + dx}_{center_key[1] + dy}"
                    neighbor_row = key_lookup.get(neighbor_key)
                    if neighbor_row is None:
                        continue
                    neighbors.append((dx, dy, neighbor_row))

        center_bounds = (
            float(row["xmin"]),
            float(row["ymin"]),
            float(row["xmax"]),
            float(row["ymax"]),
        )
        out_res = int(row["out_res"])
        step = out_res - 1
        tile_w = center_bounds[2] - center_bounds[0]
        tile_h = center_bounds[3] - center_bounds[1]

        if not neighbors:
            neighbors = [(0, 0, row)]

        min_dx = min(item[0] for item in neighbors)
        max_dx = max(item[0] for item in neighbors)
        min_dy = min(item[1] for item in neighbors)
        max_dy = max(item[1] for item in neighbors)
        tiles_x = max_dx - min_dx + 1
        tiles_y = max_dy - min_dy + 1

        mosaic_w = tiles_x * step + 1
        mosaic_h = tiles_y * step + 1

        if center_key and cfg.tile_key.enabled:
            grid_x = center_key[0] * cfg.tile_key.tile_size_x
            grid_y = center_key[1] * cfg.tile_key.tile_size_y
            
            mosaic_xmin = grid_x + min_dx * cfg.tile_key.tile_size_x
            mosaic_ymin = grid_y + min_dy * cfg.tile_key.tile_size_y
            mosaic_xmax = mosaic_xmin + tiles_x * cfg.tile_key.tile_size_x
            mosaic_ymax = mosaic_ymin + tiles_y * cfg.tile_key.tile_size_y
        else:
            mosaic_xmin = center_bounds[0] + min_dx * tile_w
            mosaic_ymin = center_bounds[1] + min_dy * tile_h
            mosaic_xmax = mosaic_xmin + tile_w * tiles_x
            mosaic_ymax = mosaic_ymin + tile_h * tiles_y
            
        mosaic_bounds = (mosaic_xmin, mosaic_ymin, mosaic_xmax, mosaic_ymax)

        height_mosaic = np.full((mosaic_h, mosaic_w), np.nan, dtype=np.float32)

        center_idx = None
        for idx, (dx, dy, _) in enumerate(neighbors):
            if dx == 0 and dy == 0:
                center_idx = idx
                break

        ordered_neighbors = neighbors[:]
        if center_idx is not None:
            center_item = ordered_neighbors.pop(center_idx)
            ordered_neighbors.append(center_item)

        for dx, dy, nrow in ordered_neighbors:
            tile_id_n = nrow["tile_id"].strip()
            png_path = os.path.join(cfg.export.heightmap_dir, f"{tile_id_n}.png")
            if not os.path.exists(png_path):
                print(f"[river_erosion] Missing PNG for {tile_id_n}: {png_path}")
                continue
            ds_tile = open_dataset(png_path, f"[river_erosion] Could not open {png_path}.")
            band = ds_tile.GetRasterBand(1)
            raw = band.ReadAsArray().astype(np.float32)
            ds_tile = None
            if raw is None or raw.size == 0:
                continue

            if raw.shape[0] != out_res or raw.shape[1] != out_res:
                out_res = raw.shape[0]
                step = out_res - 1

            scale_min = float(nrow.get("tile_min") or nrow.get("global_min"))
            scale_max = float(nrow.get("tile_max") or nrow.get("global_max"))
            if scale_max <= scale_min:
                scale_max = scale_min + 1e-3

            height = scale_min + (raw / 65535.0) * (scale_max - scale_min)

            x0 = (dx - min_dx) * step
            y0 = (max_dy - dy) * step
            height_mosaic[y0 : y0 + out_res, x0 : x0 + out_res] = height

        height_mosaic = _fill_nan_with_nearest(height_mosaic)

        laz_paths = []
        for _, _, nrow in neighbors:
            laz_path = _laz_path_for_tile(nrow["tile_id"].strip(), source_dir)
            if laz_path:
                laz_paths.append(laz_path)
        laz_paths = sorted(set(laz_paths))

        mask = _build_water_mask_mosaic(laz_paths, mosaic_bounds, mosaic_w, mosaic_h, water_class)
        if mask is None:
            _river_debug(f"Tile {tile_id}: no LiDAR water mask; copying source.")
            mask = np.zeros((mosaic_h, mosaic_w), dtype=np.float32)
        else:
            mask = _refine_mask(mask, fill_radius, smooth_sigma)

        order_line = _build_order_lines(
            layer_sources,
            mosaic_bounds,
            mosaic_w,
            mosaic_h,
            depth_profile,
            to_tile,
            to_layer,
            tile_srs,
        )

        line_mask = order_line > 0.0
        mask_bin = mask > 0.1
        depth_map = np.zeros_like(order_line, dtype=np.float32)
        labels = None
        num = 0
        indices = None
        order_nearest = None
        dist_nearest = None

        if np.any(line_mask):
            _, indices = ndimage.distance_transform_edt(~line_mask, return_indices=True)
            order_nearest = order_line[tuple(indices)]

        if np.any(mask_bin):
            labels, num = ndimage.label(mask_bin)
            fallback_depth = float(getattr(depth_profile, "fallback_depth_m", 2.0))
            if num > 0:
                for comp in range(1, num + 1):
                    comp_mask = labels == comp
                    if order_nearest is not None:
                        comp_orders = order_nearest[comp_mask]
                        comp_orders = comp_orders[comp_orders > 0.0]
                        if comp_orders.size > 0:
                            order_val = int(round(float(np.median(comp_orders))))
                            depth_val = _order_depth(order_val, depth_profile)
                        else:
                            depth_val = None
                    else:
                        depth_val = None

                    if depth_val is None or depth_val <= 0.0:
                        depth_val = fallback_depth

                    if depth_val and depth_val > 0.0:
                        depth_map[comp_mask] = depth_val * mask[comp_mask]

        # --- Manual Mask Integration ---
        # Load manual mask for the center tile (if available)
        man_b, man_g, man_r = _load_manual_mask(tile_id)
        
        if man_b is not None or man_g is not None or man_r is not None:
            # Determine slice for center tile in the mosaic
            x0_c = (0 - min_dx) * step
            y0_c = (max_dy - 0) * step
            sl_y = slice(y0_c, y0_c + out_res)
            sl_x = slice(x0_c, x0_c + out_res)
            
            # 1. Apply Erasers (Green=Island, Red=Bridge)
            if man_g is not None:
                depth_map[sl_y, sl_x][man_g] = 0.0
                mask[sl_y, sl_x][man_g] = 0.0
            
            if man_r is not None:
                depth_map[sl_y, sl_x][man_r] = 0.0
                mask[sl_y, sl_x][man_r] = 0.0
                
            # 2. Apply Lakes (Blue)
            if man_b is not None:
                # Force water mask
                mask[sl_y, sl_x][man_b] = 1.0
                
                # Calculate depths
                lbl_lakes, num_lakes = ndimage.label(man_b)
                if num_lakes > 0:
                    gt_c = _tile_geotransform(center_bounds, out_res, out_res)
                    lake_geoms = []
                    # Get centers of mass
                    # center_of_mass returns list of tuples if index is a sequence
                    centers = ndimage.center_of_mass(man_b, lbl_lakes, range(1, num_lakes + 1))
                        
                    for (r, c) in centers:
                        # Pixel to World
                        # x = xmin + c * xres
                        # y = ymax + r * yres (yres is negative)
                        wx = gt_c[0] + c * gt_c[1]
                        wy = gt_c[3] + r * gt_c[5]
                        pt = ogr.Geometry(ogr.wkbPoint)
                        pt.AddPoint(wx, wy)
                        lake_geoms.append(pt)
                    
                    # Lookup depths
                    lake_depths_meta = _find_lake_depths(
                        lake_geoms, 
                        tile_srs, 
                        re_cfg.lakes, 
                        re_cfg.lakes.source_path
                    )
                    
                    # Apply depths
                    for idx, (depth, meta) in enumerate(lake_depths_meta):
                        comp_idx = idx + 1
                        lake_mask = (lbl_lakes == comp_idx)
                        # Apply to mosaic depth map
                        depth_map[sl_y, sl_x][lake_mask] = depth
                        
                        match_status = "MATCH" if meta.get("matched") else "DEFAULT"
                        dist_info = f"{meta.get('distance', -1):.1f}m"
                        _river_debug(
                            f"Tile {tile_id} Manual Lake {idx+1}: "
                            f"Depth={depth:.2f}m ({match_status}, {dist_info})"
                        )

        height_mosaic = height_mosaic - depth_map

        flow_slope = float(getattr(depth_profile, "flow_slope_m_per_km", 0.0))
        dist_field = str(getattr(depth_profile, "flow_distance_field", "DIST_DN_KM") or "DIST_DN_KM")
        if flow_slope > 0.0 and np.any(line_mask) and indices is not None and labels is not None and num > 0:
            dist_line = _build_flow_distance(
                layer_sources,
                mosaic_bounds,
                mosaic_w,
                mosaic_h,
                dist_field,
                to_tile,
                to_layer,
                tile_srs,
            )
            dist_nearest = dist_line[tuple(indices)]
            for comp in range(1, num + 1):
                comp_mask = labels == comp
                if not np.any(comp_mask):
                    continue
                dist_vals = dist_nearest[comp_mask]
                heights = height_mosaic[comp_mask]
                if dist_vals.size == 0 or heights.size == 0:
                    continue
                order = np.argsort(-dist_vals)
                dist_sorted = dist_vals[order]
                heights_sorted = heights[order]
                adjusted = heights_sorted + flow_slope * dist_sorted
                enforced = np.minimum.accumulate(adjusted)
                new_heights_sorted = enforced - flow_slope * dist_sorted
                new_heights = heights.copy()
                new_heights[order] = new_heights_sorted
                height_mosaic[comp_mask] = np.minimum(heights, new_heights)


        x0_center = (0 - min_dx) * step
        y0_center = (max_dy - 0) * step
        cropped = height_mosaic[y0_center : y0_center + out_res, x0_center : x0_center + out_res]

        if cfg.ortho.apply_water_mask:
            mask_center = mask[y0_center : y0_center + out_res, x0_center : x0_center + out_res]
            _apply_water_mask_to_ortho(tile_id, mask_center, cfg)

        mask_center = mask[y0_center : y0_center + out_res, x0_center : x0_center + out_res]
        threshold = float(cfg.ortho.water_mask_threshold)
        water_bin = mask_center >= threshold

        mask_dir = os.path.join(cfg.work.work_dir, "river_masks")
        ensure_dir(mask_dir)
        mask_path = os.path.join(mask_dir, f"{tile_id}_mask.png")
        _write_mask_png(mask_path, (water_bin.astype(np.uint8) * 255), center_bounds, tile_srs)

        tile_min = float(np.min(cropped))
        tile_max = float(np.max(cropped))
        if tile_max <= tile_min:
            tile_max = tile_min + 1e-3

        scaled = (cropped - tile_min) / (tile_max - tile_min)
        out_u16 = np.clip(scaled * 65535.0, 0.0, 65535.0).astype(np.uint16)

        global_min = min(global_min, tile_min)
        global_max = max(global_max, tile_max)

        out_path = os.path.join(output_dir, f"{tile_id}.png")
        mem_driver = gdal.GetDriverByName("MEM")
        out_ds = mem_driver.Create("", out_u16.shape[1], out_u16.shape[0], 1, gdal.GDT_UInt16)
        out_ds.SetGeoTransform(_tile_geotransform(center_bounds, out_res, out_res))
        out_ds.SetProjection(tile_srs.ExportToWkt())
        out_ds.GetRasterBand(1).WriteArray(out_u16)
        trans_opts = gdal.TranslateOptions(
            outputType=gdal.GDT_UInt16,
            format="PNG",
            creationOptions=["WORLDFILE=YES"],
        )
        try:
            gdal.Translate(out_path, out_ds, options=trans_opts)
        except RuntimeError as exc:
            print(f"[river_erosion] Translate failed for {tile_id}: {exc}")
            continue
        out_ds = None

        rows_out.append(
            {
                "tile_id": tile_id,
                "xmin": f"{center_bounds[0]:.6f}",
                "ymin": f"{center_bounds[1]:.6f}",
                "xmax": f"{center_bounds[2]:.6f}",
                "ymax": f"{center_bounds[3]:.6f}",
                "out_res": str(out_res),
                "tile_key": tile_key,
                "tile_min": f"{tile_min:.6f}",
                "tile_max": f"{tile_max:.6f}",
            }
        )

        print(f"[river_erosion] Wrote {out_path}")

    if not rows_out:
        print("[river_erosion] No tiles processed; nothing written.")
        return 1

    _write_manifest(manifest_out, rows_out, global_min, global_max)
    print(f"[river_erosion] Manifest: {manifest_out}")
    print(f"[river_erosion] Summary: wrote {len(rows_out)} tiles.")

    return 0
