"""DEM (digital elevation model) store: upload once, preview, sample, build from points.

Rasters are kept under SPIPE_DEM_DIR (default: <repo>/data/dem) as GeoTIFF,
keyed by a random id; only the newest SPIPE_DEM_KEEP files are retained.
Elevations are returned in the raster's own units (normally metres).
"""
from __future__ import annotations

import json
import os
import shutil
import uuid
import warnings
import zipfile
from pathlib import Path

import numpy as np
import rasterio
import shapely
import shapely.geometry
from rasterio.enums import Resampling
from rasterio.io import MemoryFile
from rasterio.transform import from_origin
from pyproj import CRS, Transformer

from .batch import BatchError

DEM_DIR = Path(os.environ.get("SPIPE_DEM_DIR", Path(__file__).resolve().parents[1] / "data" / "dem"))
DEM_KEEP = int(os.environ.get("SPIPE_DEM_KEEP", "8"))
RASTER_EXTENSIONS = (".tif", ".tiff", ".asc", ".img", ".dem", ".bil", ".flt", ".hgt", ".vrt")
MAX_CELLS_FROM_POINTS = 25_000_000
_ID_OK = set("0123456789abcdef")


def _path(dem_id: str) -> Path:
    if len(dem_id) != 32 or not set(dem_id) <= _ID_OK:
        raise BatchError("DEM id가 올바르지 않습니다.")
    p = DEM_DIR / f"{dem_id}.tif"
    if not p.exists():
        raise BatchError("DEM을 찾을 수 없습니다. 다시 불러오세요 (서버가 오래된 DEM을 정리했을 수 있습니다).")
    return p


def _trim() -> None:
    files = sorted(DEM_DIR.glob("*.tif"), key=lambda p: p.stat().st_mtime, reverse=True)
    for old in files[DEM_KEEP:]:
        old.unlink(missing_ok=True)
        old.with_suffix(".json").unlink(missing_ok=True)


def _crs_of(user: str | None) -> CRS:
    if not user:
        raise BatchError("좌표계가 필요합니다 (로컬 좌표계에서는 DEM을 쓸 수 없습니다).")
    try:
        return CRS.from_user_input(user)
    except Exception as e:  # pyproj raises several types
        raise BatchError(f"좌표계를 읽을 수 없습니다: {user[:60]}") from e


def _info(dem_id: str, src, name: str) -> dict:
    if src.crs is None:
        raise BatchError("DEM에 좌표계 정보가 없습니다 (.prj 또는 GeoTIFF 좌표계를 넣어 주세요).")
    band = src.read(1, masked=True, out_shape=(min(src.height, 1024), min(src.width, 1024)), resampling=Resampling.nearest)
    to_wgs = Transformer.from_crs(src.crs, "EPSG:4326", always_xy=True)
    l, b, r, t = src.bounds
    corners = [list(to_wgs.transform(x, y)) for x, y in ((l, t), (r, t), (r, b), (l, b))]
    return {
        "id": dem_id,
        "name": name,
        "crs": src.crs.to_string(),
        "width": src.width,
        "height": src.height,
        "resolution": [abs(src.res[0]), abs(src.res[1])],
        "bounds": [l, b, r, t],
        "corners": corners,
        "min": float(band.min()) if band.count() else None,
        "max": float(band.max()) if band.count() else None,
        "nodata": src.nodata,
    }


def _store(src_path: str, name: str) -> dict:
    DEM_DIR.mkdir(parents=True, exist_ok=True)
    dem_id = uuid.uuid4().hex
    out = DEM_DIR / f"{dem_id}.tif"
    try:
        with rasterio.open(src_path) as src:
            if src.count < 1:
                raise BatchError("래스터 밴드가 없습니다.")
            profile = src.profile.copy()
            profile.update(driver="GTiff", count=1, compress="deflate", tiled=True, blockxsize=256, blockysize=256, BIGTIFF="IF_SAFER")
            if src.width < 256 or src.height < 256:
                profile.update(tiled=False)
                profile.pop("blockxsize", None)
                profile.pop("blockysize", None)
            with rasterio.open(out, "w", **profile) as dst:
                for _, window in src.block_windows(1):
                    dst.write(src.read(1, window=window), 1, window=window)
        with rasterio.open(out) as src:
            info = _info(dem_id, src, name)
    except rasterio.errors.RasterioIOError as e:
        out.unlink(missing_ok=True)
        raise BatchError(f"래스터를 읽을 수 없습니다: {e}") from e
    except BatchError:
        out.unlink(missing_ok=True)
        raise
    out.with_suffix(".json").write_text(json.dumps(info, ensure_ascii=False), encoding="utf-8")
    _trim()
    return info


def import_dem(upload_path: str, filename: str) -> dict:
    """Store an uploaded raster (or a zip holding one, e.g. ASC + PRJ) as a DEM."""
    ext = os.path.splitext(filename.lower())[1]
    if ext == ".zip":
        try:
            with zipfile.ZipFile(upload_path) as z:
                names = [n for n in z.namelist() if os.path.splitext(n.lower())[1] in RASTER_EXTENSIONS and not n.startswith("__MACOSX")]
        except zipfile.BadZipFile as e:
            raise BatchError("zip 파일이 손상되었습니다.") from e
        if not names:
            raise BatchError(f"zip 안에 래스터({', '.join(RASTER_EXTENSIONS)})가 없습니다.")
        return _store("/vsizip/" + upload_path.replace("\\", "/") + "/" + names[0], filename)
    if ext not in RASTER_EXTENSIONS:
        raise BatchError(f"지원하지 않는 형식입니다: {ext} (GeoTIFF, ASC, IMG 또는 이를 담은 zip)")
    # GDAL picks the driver from the extension; keep it on the temp copy.
    named = upload_path + ext
    shutil.copyfile(upload_path, named)
    try:
        return _store(named, filename)
    finally:
        os.remove(named)


def dem_info(dem_id: str) -> dict:
    p = _path(dem_id)
    return json.loads(p.with_suffix(".json").read_text(encoding="utf-8"))


def preview_png(dem_id: str, max_size: int = 1600) -> bytes:
    """Hillshaded colour relief, transparent where there is no data."""
    with rasterio.open(_path(dem_id)) as src:
        k = max(1, int(np.ceil(max(src.width, src.height) / max_size)))
        h, w = max(1, src.height // k), max(1, src.width // k)
        z = src.read(1, masked=True, out_shape=(h, w), resampling=Resampling.bilinear).astype("float64")
        cell = abs(src.res[0]) * src.width / w
        geographic = src.crs and src.crs.is_geographic
    if geographic:
        cell *= 111_000
    mask = np.ma.getmaskarray(z)
    zf = z.filled(np.nan)
    lo, hi = np.nanpercentile(zf, [2, 98]) if np.isfinite(zf).any() else (0.0, 1.0)
    t = np.clip((zf - lo) / max(hi - lo, 1e-9), 0, 1)
    # Green -> yellow -> brown -> white elevation ramp.
    stops = np.array([[0.0, 64, 140, 80], [0.35, 190, 200, 110], [0.7, 160, 110, 60], [1.0, 245, 245, 240]])
    rgb = np.stack([np.interp(np.nan_to_num(t), stops[:, 0], stops[:, c]) for c in (1, 2, 3)])
    gy, gx = np.gradient(np.nan_to_num(zf, nan=float(np.nanmean(zf)) if np.isfinite(zf).any() else 0.0), cell)
    slope = np.arctan(np.hypot(gx, gy))
    aspect = np.arctan2(-gx, gy)
    az, alt = np.radians(315), np.radians(45)
    shade = np.clip(np.sin(alt) * np.cos(slope) + np.cos(alt) * np.sin(slope) * np.cos(az - aspect), 0, 1)
    rgb = rgb * (0.55 + 0.45 * shade)
    alpha = np.where(mask | ~np.isfinite(zf), 0, 255)
    img = np.concatenate([rgb, alpha[None]]).clip(0, 255).astype("uint8")
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", rasterio.errors.NotGeoreferencedWarning)  # a plain image
        with MemoryFile() as mem:
            with mem.open(driver="PNG", width=w, height=h, count=4, dtype="uint8") as dst:
                dst.write(img)
            return mem.read()


def sample(dem_id: str, points: list, points_crs: str | None) -> list:
    """Bilinear DEM values at [x, y] points given in `points_crs` (None outside / no data)."""
    if not points:
        return []
    pts = np.asarray(points, dtype="float64").reshape(-1, 2)
    with rasterio.open(_path(dem_id)) as src:
        tr = Transformer.from_crs(_crs_of(points_crs), src.crs, always_xy=True)
        xs, ys = tr.transform(pts[:, 0], pts[:, 1])
        inv = ~src.transform
        xs, ys = np.asarray(xs), np.asarray(ys)
        cols, rows = inv.a * xs + inv.b * ys + inv.c, inv.d * xs + inv.e * ys + inv.f
        # pixel centres sit at +0.5
        cols, rows = np.asarray(cols) - 0.5, np.asarray(rows) - 0.5
        ok = np.isfinite(cols) & np.isfinite(rows) & (cols >= -0.5) & (rows >= -0.5) & (cols <= src.width - 0.5) & (rows <= src.height - 0.5)
        out: list = [None] * len(pts)
        if not ok.any():
            return out
        c0 = int(max(0, np.floor(cols[ok].min()))); c1 = int(min(src.width - 1, np.floor(cols[ok].max()) + 1))
        r0 = int(max(0, np.floor(rows[ok].min()))); r1 = int(min(src.height - 1, np.floor(rows[ok].max()) + 1))
        win = rasterio.windows.Window(c0, r0, c1 - c0 + 1, r1 - r0 + 1)
        if win.width * win.height > 400_000_000:
            raise BatchError("지점들이 너무 넓게 퍼져 있습니다. 나눠서 실행하세요.")
        z = src.read(1, window=win, masked=True).astype("float64").filled(np.nan)
    for k in np.where(ok)[0]:
        c, r = cols[k] - c0, rows[k] - r0
        ci, ri = int(np.clip(np.floor(c), 0, z.shape[1] - 1)), int(np.clip(np.floor(r), 0, z.shape[0] - 1))
        cj, rj = min(ci + 1, z.shape[1] - 1), min(ri + 1, z.shape[0] - 1)
        fc, fr = float(np.clip(c - ci, 0, 1)), float(np.clip(r - ri, 0, 1))
        q = np.array([z[ri, ci], z[ri, cj], z[rj, ci], z[rj, cj]])
        wts = np.array([(1 - fc) * (1 - fr), fc * (1 - fr), (1 - fc) * fr, fc * fr])
        good = np.isfinite(q) & (wts > 0)
        if good.any():
            out[k] = round(float(np.sum(q[good] * wts[good]) / np.sum(wts[good])), 3)
        elif np.isfinite(q[0]):
            out[k] = round(float(q[0]), 3)
    return out


def from_points(points: list, points_crs: str | None, cell: float, name: str = "점 표고 DEM", target_crs: str | None = None) -> dict:
    """TIN-style DEM: linear interpolation over the Delaunay triangulation of [x, y, z] points."""
    from scipy.interpolate import LinearNDInterpolator

    pts = np.asarray(points, dtype="float64").reshape(-1, 3)
    pts = pts[np.all(np.isfinite(pts), axis=1)]
    if len(pts) < 3:
        raise BatchError("표고 점이 3개 이상 필요합니다.")
    crs = _crs_of(points_crs)
    if target_crs:
        # e.g. GIS points arrive in WGS84; grid them in the model's plane CRS.
        dst = _crs_of(target_crs)
        x, y = Transformer.from_crs(crs, dst, always_xy=True).transform(pts[:, 0], pts[:, 1])
        pts = np.column_stack([x, y, pts[:, 2]])
        crs = dst
    if crs.is_geographic:
        raise BatchError("경위도 좌표계 점으로는 만들 수 없습니다. 평면 좌표계를 쓰세요.")
    if not (cell > 0):
        raise BatchError("격자 크기는 0보다 커야 합니다.")
    x0, y0 = pts[:, 0].min(), pts[:, 1].min()
    x1, y1 = pts[:, 0].max(), pts[:, 1].max()
    w, h = int(np.ceil((x1 - x0) / cell)) + 1, int(np.ceil((y1 - y0) / cell)) + 1
    if w * h > MAX_CELLS_FROM_POINTS:
        raise BatchError(f"격자가 너무 큽니다 ({w}x{h}). 격자 크기를 키우세요.")
    try:
        interp = LinearNDInterpolator(pts[:, :2], pts[:, 2])
    except Exception as e:  # Qhull errors on collinear input
        raise BatchError("점들이 한 줄 위에 있어 삼각망을 만들 수 없습니다.") from e
    gx = x0 + (np.arange(w) + 0.5) * cell
    gy = y1 - (np.arange(h) + 0.5) * cell
    xx, yy = np.meshgrid(gx, gy)
    z = interp(xx, yy).astype("float32")
    nodata = -9999.0
    z[~np.isfinite(z)] = nodata
    DEM_DIR.mkdir(parents=True, exist_ok=True)
    tmp = DEM_DIR / f"tmp_{uuid.uuid4().hex}.tif"
    with rasterio.open(tmp, "w", driver="GTiff", width=w, height=h, count=1, dtype="float32", crs=crs.to_wkt(),
                       transform=from_origin(x0, y1, cell, cell), nodata=nodata, compress="deflate") as dst:
        dst.write(z, 1)
    try:
        return _store(str(tmp), name)
    finally:
        tmp.unlink(missing_ok=True)


def terrain_tile(dem_id: str, z: int, x: int, y: int) -> bytes:
    """Web-mercator XYZ tile of the DEM in Mapbox/MapLibre 'terrarium' encoding
    (elevation = R*256 + G + B/256 - 32768)."""
    from rasterio.warp import reproject

    if not (0 <= z <= 18 and 0 <= x < 2 ** z and 0 <= y < 2 ** z):
        raise BatchError("타일 번호가 올바르지 않습니다.")
    size = 256
    world = 20037508.342789244
    span = 2 * world / 2 ** z
    x0, y1 = -world + x * span, world - y * span
    dst = np.full((size, size), np.nan, dtype="float32")
    with rasterio.open(_path(dem_id)) as src:
        reproject(
            source=rasterio.band(src, 1), destination=dst,
            dst_transform=from_origin(x0, y1, span / size, span / size), dst_crs="EPSG:3857",
            dst_nodata=np.nan, resampling=Resampling.bilinear,
        )
    # Outside the DEM: its lowest value, so the terrain does not drop to sea level in a cliff.
    fill = float(np.nanmin(dst)) if np.isfinite(dst).any() else float(dem_info(dem_id).get("min") or 0.0)
    v = np.nan_to_num(dst, nan=fill).astype("float64") + 32768.0
    r = np.floor(v / 256)
    g = np.floor(v - r * 256)
    b = np.floor((v - r * 256 - g) * 256)
    img = np.stack([r, g, b]).clip(0, 255).astype("uint8")
    with warnings.catch_warnings():
        warnings.simplefilter("ignore", rasterio.errors.NotGeoreferencedWarning)
        with MemoryFile() as mem:
            with mem.open(driver="PNG", width=size, height=size, count=3, dtype="uint8") as d:
                d.write(img)
            return mem.read()


# ---------------------------------------------------------------- hydrology (유역·하천)

MAX_HYDRO_CELLS = 1_000_000
_D8 = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)]


def _hydro_grid(dem_id: str):
    """DEM (downsampled to <= MAX_HYDRO_CELLS) -> D8 receivers and flow accumulation
    on a priority-flood filled surface (pits and flats drain)."""
    import heapq

    with rasterio.open(_path(dem_id)) as src:
        k = max(1, int(np.ceil(np.sqrt(src.width * src.height / MAX_HYDRO_CELLS))))
        h, w = max(2, src.height // k), max(2, src.width // k)
        z = src.read(1, masked=True, out_shape=(h, w), resampling=Resampling.average).astype("float64")
        transform = src.transform @ src.transform.scale(src.width / w, src.height / h)
        crs = src.crs
    if crs.is_geographic:
        raise BatchError("유역 분석에는 평면 좌표계 DEM이 필요합니다.")
    valid = ~np.ma.getmaskarray(z)
    filled = np.where(valid, z.filled(np.nan), np.inf)
    done = ~valid.copy()
    heap: list = []
    edge = np.zeros_like(valid)
    edge[0, :] = edge[-1, :] = edge[:, 0] = edge[:, -1] = True
    pad = np.pad(~valid, 1, constant_values=True)
    near_nodata = np.zeros_like(valid)
    for dr, dc in _D8:
        near_nodata |= pad[1 + dr:1 + dr + h, 1 + dc:1 + dc + w]
    for r, c in zip(*np.where(valid & (edge | near_nodata))):
        heapq.heappush(heap, (filled[r, c], int(r), int(c)))
        done[r, c] = True
    eps = 1e-4
    order = []
    while heap:
        e, r, c = heapq.heappop(heap)
        order.append(r * w + c)
        for dr, dc in _D8:
            rr, cc = r + dr, c + dc
            if 0 <= rr < h and 0 <= cc < w and not done[rr, cc]:
                done[rr, cc] = True
                filled[rr, cc] = max(filled[rr, cc], e + eps)
                heapq.heappush(heap, (filled[rr, cc], rr, cc))
    n = h * w
    flat = filled.ravel()
    recv = np.full(n, -1, dtype=np.int64)
    best = np.zeros(n)
    cell = abs(transform.a)
    rows, cols = np.divmod(np.arange(n), w)
    for dr, dc in _D8:
        rr, cc = rows + dr, cols + dc
        ok = (rr >= 0) & (rr < h) & (cc >= 0) & (cc < w)
        nb = np.where(ok, rr * w + cc, 0)
        with np.errstate(invalid="ignore"):
            drop = np.where(ok, (flat - flat[nb]) / (cell * np.hypot(dr, dc)), -np.inf)
        drop[~np.isfinite(drop)] = -np.inf
        take = drop > best
        best[take] = drop[take]
        recv[take] = nb[take]
    acc = valid.ravel().astype("float64")
    for i in reversed(order):  # from the highest cell down
        j = recv[i]
        if j >= 0:
            acc[j] += acc[i]
    return {"h": h, "w": w, "recv": recv, "acc": acc, "transform": transform, "crs": crs, "cell_area": abs(transform.a * transform.e)}


def streams(dem_id: str, min_area_km2: float) -> dict:
    """Stream lines (WGS84 GeoJSON): cells draining at least `min_area_km2`."""
    g = _hydro_grid(dem_id)
    thr = max(2.0, min_area_km2 * 1e6 / g["cell_area"])
    w, tr = g["w"], g["transform"]
    idx = np.where((g["acc"] >= thr) & (g["recv"] >= 0))[0]
    if len(idx) > 300_000:
        raise BatchError("하천 셀이 너무 많습니다. 최소 유역 면적을 키우세요.")
    if not len(idx):
        return {"type": "FeatureCollection", "features": []}

    def xy(i):
        r, c = divmod(int(i), w)
        return tr @ (c + 0.5, r + 0.5)

    segs = np.array([shapely.LineString([xy(i), xy(g["recv"][i])]) for i in idx], dtype=object)
    merged = shapely.get_parts(shapely.line_merge(shapely.union_all(segs)))
    to_wgs = Transformer.from_crs(g["crs"], "EPSG:4326", always_xy=True)
    pts = shapely.points([xy(i) for i in idx])
    tree = shapely.STRtree(pts)
    acc_at = g["acc"][idx]
    feats = []
    for line in merged:
        near = tree.query(line, predicate="dwithin", distance=abs(tr.a) * 0.6)
        area = float(acc_at[near].max() * g["cell_area"] / 1e6) if len(near) else 0.0
        coords = np.asarray(line.coords)
        x, y = to_wgs.transform(coords[:, 0], coords[:, 1])
        feats.append({"type": "Feature", "properties": {"area_km2": round(area, 3), "length_m": round(float(line.length), 1)},
                      "geometry": {"type": "LineString", "coordinates": [[float(a), float(b)] for a, b in zip(x, y)]}})
    return {"type": "FeatureCollection", "features": feats}


def watershed(dem_id: str, points: list, points_crs: str | None, snap_m: float = 50.0) -> dict:
    """Catchment polygon upstream of each pour point (snapped to the largest flow within snap_m)."""
    from rasterio import features as rfeatures

    g = _hydro_grid(dem_id)
    h, w, tr = g["h"], g["w"], g["transform"]
    to_dem = Transformer.from_crs(_crs_of(points_crs), g["crs"], always_xy=True)
    to_wgs = Transformer.from_crs(g["crs"], "EPSG:4326", always_xy=True)
    order = np.argsort(g["recv"], kind="stable")
    recv_sorted = g["recv"][order]
    inv = ~tr
    rad = max(1, int(round(snap_m / abs(tr.a))))
    acc2 = g["acc"].reshape(h, w)
    feats = []
    for k, p in enumerate(points):
        x, y = to_dem.transform(float(p[0]), float(p[1]))
        c = inv.a * x + inv.b * y + inv.c
        r = inv.d * x + inv.e * y + inv.f
        r, c = int(r), int(c)
        if not (0 <= r < h and 0 <= c < w):
            continue
        r0, r1, c0, c1 = max(0, r - rad), min(h, r + rad + 1), max(0, c - rad), min(w, c + rad + 1)
        win = acc2[r0:r1, c0:c1]
        rr, cc = np.unravel_index(int(np.argmax(win)), win.shape)
        start = (r0 + rr) * w + (c0 + cc)
        mask = np.zeros(h * w, dtype=bool)
        mask[start] = True
        stack = [start]
        while stack:
            j = stack.pop()
            lo = np.searchsorted(recv_sorted, j, "left")
            hi = np.searchsorted(recv_sorted, j, "right")
            for i in order[lo:hi]:
                if not mask[i]:
                    mask[i] = True
                    stack.append(int(i))
        m2 = mask.reshape(h, w)
        polys = [shapely.geometry.shape(geom) for geom, val in rfeatures.shapes(m2.astype("uint8"), mask=m2, transform=tr) if val == 1]
        if not polys:
            continue
        poly = shapely.union_all(np.array(polys, dtype=object))
        poly = shapely.transform(poly, lambda a: np.column_stack(to_wgs.transform(a[:, 0], a[:, 1])))
        feats.append({"type": "Feature", "properties": {"point": k, "area_km2": round(float(mask.sum() * g["cell_area"] / 1e6), 4)},
                      "geometry": shapely.geometry.mapping(poly)})
    return {"type": "FeatureCollection", "features": feats}