"""공용 지형 창고 — 도엽(1:5,000)마다 2 m 지형을 처음 쓸 때 한 번 셈 · 영구 저장
(PLAN 607-2 · 607-3).

자리 = `storage/_shared/terrain/<도엽>/dem_<epsg>_<격자>m.npz` + `.key.json`
(모든 작업 폴더가 보는 정션 · 시놀로지로 데스크탑에도).
원천 = 전국 등고선 파일(EPSG:5179) · 그 도엽에 등고선이 없으면(섬 등)
그때만 도엽 zip 을 받아 같은 창고 폴더에(B04 받기 함수 · `prune_sheets` 안 부름).
도엽 + 여유 200 m 등고선으로 셈 → 도엽 안쪽 칸(칸 중심 위경도가 그 도엽)만 저장 → 이어 붙이면
통째 셈과 같음. 격자 원점 = 좌표 격자 배수(칸 모서리 = k·cell_m).
열쇠 = 계산 방식 판 + 여유 + 전국 파일 지문(크기 · 수정 시각) + 쓴 도엽 zip 해시.
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import os
import time
from pathlib import Path
from typing import Any, Callable

import numpy as np
from pyproj import Transformer

from B04_PreProcess.B04_PreProcess_Engine_MapSheet import (
    SHEET5K_SIZE_DEG,
    latlon_to_sheet5k,
    neighbors_3x3,
    sheet5k_to_bounds,
)
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import ELEVATION_KEYS, GridSpec
from common_util.common_util_atomic import atomic_write_npz
from common_util.common_util_json import atomic_write_json
from common_util.common_util_terrain_surface import (
    TERRAIN_ALGO,
    build_contour_surface,
    segments_from_lines,
)
from config.config_system import PROJECT_ROOT, STORAGE_BASE_DIR

logger = logging.getLogger(__name__)

TERRAIN_ROOT = Path(STORAGE_BASE_DIR) / "_shared" / "terrain"
NATIONAL_GPKG = Path(PROJECT_ROOT) / "resources" / "data_global_contours" / "national_contours.gpkg"
TERRAIN_MARGIN_M = 200.0
# ponytail: 잠금 파일이 이보다 오래되면 죽은 셈으로 보고 깸 · 도엽 하나 셈 1~5 분
LOCK_STALE_S = 15 * 60

Progress = Callable[[int, str], None] | None
Box = tuple[float, float, float, float]
Segments = tuple[np.ndarray, np.ndarray, np.ndarray]
_NO_SEGMENTS: Segments = (np.zeros((0, 2)), np.zeros((0, 2)), np.zeros(0))


def _names(sheet_no: str, epsg: int, cell_m: float) -> tuple[Path, Path, Path]:
    folder = TERRAIN_ROOT / str(sheet_no)
    stem = f"dem_{int(epsg)}_{cell_m:g}m"
    return folder / f"{stem}.npz", folder / f"{stem}.key.json", folder / f".{stem}.lock"


def _gpkg_stamp() -> str | None:
    try:
        stat = NATIONAL_GPKG.stat()
    except OSError:
        return None
    return f"{stat.st_size}:{stat.st_mtime_ns}"


def _sha256(path: Path) -> str | None:
    try:
        return hashlib.sha256(path.read_bytes()).hexdigest()
    except OSError:
        return None


def _key(epsg: int, cell_m: float, zips: dict[str, str | None]) -> dict[str, Any]:
    return {
        "algo": TERRAIN_ALGO,
        "margin_m": TERRAIN_MARGIN_M,
        "epsg": int(epsg),
        "cell_m": float(cell_m),
        "gpkg": _gpkg_stamp(),
        "zips": zips,
    }


def _stored_ok(sheet_no: str, epsg: int, cell_m: float) -> dict[str, Any] | None:
    """저장 열쇠가 지금 원천과 같으면 그 열쇠 — 아니면 None(다시 셈)."""
    npz, key_path, _ = _names(sheet_no, epsg, cell_m)
    try:
        stored = json.loads(key_path.read_text(encoding="utf-8"))
    except (OSError, ValueError):
        return None
    if not stored.get("empty") and not npz.is_file():
        return None
    zips = {s: _sha256(TERRAIN_ROOT / s / f"{s}.zip") for s in stored.get("zips") or {}}
    now = _key(epsg, cell_m, zips)
    if now["gpkg"] is None:  # 전국 파일 없는 PC — 있는 저장본을 믿음(시놀로지 덮어쓰기 핑퐁 막음)
        now["gpkg"] = stored.get("gpkg")
    stored_key = {k: stored.get(k) for k in now}
    return stored if stored_key == now else None


class _Lock:
    """도엽 · 좌표계 · 격자 하나를 한 셈만 — 잠금 파일(서버 여럿 · 스레드 여럿 공통)."""

    def __init__(self, path: Path) -> None:
        self.path = path

    def __enter__(self) -> "_Lock":
        self.path.parent.mkdir(parents=True, exist_ok=True)
        while True:
            try:
                fd = os.open(self.path, os.O_CREAT | os.O_EXCL | os.O_WRONLY)
                os.write(fd, f"{os.getpid()} {time.time():.0f}".encode())
                os.close(fd)
                return self
            except FileExistsError:
                try:
                    if time.time() - self.path.stat().st_mtime > LOCK_STALE_S:
                        self.path.unlink(missing_ok=True)
                        continue
                except OSError:
                    continue
                time.sleep(0.3)

    def __exit__(self, *exc: Any) -> None:
        self.path.unlink(missing_ok=True)


def sheets_for_box(box: Box, epsg: int) -> list[str]:
    """TM 상자(x0, y0, x1, y1)에 닿는 도엽 번호(북서부터)."""
    x0, y0, x1, y1 = box
    t = np.linspace(0.0, 1.0, 21)
    xs = np.concatenate([x0 + (x1 - x0) * t, np.full(21, x1), x1 - (x1 - x0) * t, np.full(21, x0)])
    ys = np.concatenate([np.full(21, y0), y0 + (y1 - y0) * t, np.full(21, y1), y1 - (y1 - y0) * t])
    lon, lat = Transformer.from_crs(f"EPSG:{int(epsg)}", "EPSG:4326", always_xy=True).transform(
        xs, ys
    )
    step = SHEET5K_SIZE_DEG
    sheets: list[str] = []
    for i in range(math.floor(lat.max() / step), math.floor(lat.min() / step) - 1, -1):
        for j in range(math.floor(lon.min() / step), math.floor(lon.max() / step) + 1):
            no = latlon_to_sheet5k((i + 0.5) * step, (j + 0.5) * step)
            if no not in sheets:
                sheets.append(no)
    return sheets


def _owned(sheet_no: str, epsg: int, cell_m: float) -> tuple[int, int, np.ndarray]:
    """도엽 안쪽 칸 — (첫 칸 열 번호 ix0 · 첫 칸 행 번호 iy0 · 소유 표 (ny, nx) y 오름차순)."""
    lon0, lat0, lon1, lat1 = sheet5k_to_bounds(sheet_no)
    to_tm = Transformer.from_crs("EPSG:4326", f"EPSG:{int(epsg)}", always_xy=True)
    t = np.linspace(0.0, 1.0, 21)
    lon = np.concatenate([lon0 + (lon1 - lon0) * t, np.full(21, lon1), lon0 + (lon1 - lon0) * t])
    lat = np.concatenate([np.full(21, lat0), lat0 + (lat1 - lat0) * t, np.full(21, lat1)])
    lon = np.concatenate([lon, np.full(21, lon0)])
    lat = np.concatenate([lat, lat0 + (lat1 - lat0) * t])
    xs, ys = to_tm.transform(lon, lat)
    ix0, ix1 = math.floor(xs.min() / cell_m) - 1, math.ceil(xs.max() / cell_m) + 1
    iy0, iy1 = math.floor(ys.min() / cell_m) - 1, math.ceil(ys.max() / cell_m) + 1
    cx = (np.arange(ix0, ix1) + 0.5) * cell_m
    cy = (np.arange(iy0, iy1) + 0.5) * cell_m
    gx, gy = np.meshgrid(cx, cy)
    clon, clat = Transformer.from_crs(f"EPSG:{int(epsg)}", "EPSG:4326", always_xy=True).transform(
        gx, gy
    )
    own = (clon >= lon0) & (clon < lon1) & (clat >= lat0) & (clat < lat1)
    return ix0, iy0, own


def _gpkg_segments(epsg: int, box: Box) -> Segments:
    """전국 등고선에서 TM 상자에 닿는 선분 — 끝 a · b (M,2) 사업 좌표 · 높이 z (M,)."""
    if not NATIONAL_GPKG.is_file():
        return _NO_SEGMENTS
    import pyogrio
    import shapely

    x0, y0, x1, y1 = box
    to_5179 = Transformer.from_crs(f"EPSG:{int(epsg)}", "EPSG:5179", always_xy=True)
    t = np.linspace(0.0, 1.0, 21)
    bx, by = to_5179.transform(
        np.concatenate([x0 + (x1 - x0) * t, x0 + (x1 - x0) * t, np.full(21, x0), np.full(21, x1)]),
        np.concatenate([np.full(21, y0), np.full(21, y1), y0 + (y1 - y0) * t, y0 + (y1 - y0) * t]),
    )
    pad = 20.0
    lo_x, lo_y, hi_x, hi_y = bx.min() - pad, by.min() - pad, bx.max() + pad, by.max() + pad
    frame = pyogrio.read_dataframe(
        NATIONAL_GPKG, bbox=(lo_x, lo_y, hi_x, hi_y), columns=["CTRLN_HG"]
    )
    if frame.empty:
        return _NO_SEGMENTS
    parts, owner = shapely.get_parts(frame.geometry.values, return_index=True)
    levels = frame["CTRLN_HG"].to_numpy(dtype=np.float64)[owner]
    coords, part = shapely.get_coordinates(parts, return_index=True)
    # 상자 밖 선분은 버림(격자 밖 선분은 결과에 안 닿음 · 꼭짓점 수백만 변환을 줄임)
    same = part[:-1] == part[1:]
    a, b = coords[:-1], coords[1:]
    hit = (
        same
        & (np.maximum(a[:, 0], b[:, 0]) >= lo_x)
        & (np.minimum(a[:, 0], b[:, 0]) <= hi_x)
        & (np.maximum(a[:, 1], b[:, 1]) >= lo_y)
        & (np.minimum(a[:, 1], b[:, 1]) <= hi_y)
    )
    seg = np.flatnonzero(hit)
    if not len(seg):
        return _NO_SEGMENTS
    to_tm = Transformer.from_crs("EPSG:5179", f"EPSG:{int(epsg)}", always_xy=True)
    ax, ay = to_tm.transform(a[seg, 0], a[seg, 1])
    bx2, by2 = to_tm.transform(b[seg, 0], b[seg, 1])
    return np.column_stack([ax, ay]), np.column_stack([bx2, by2]), levels[part[seg]]


def _zip_lines(sheet_nos: list[str], epsg: int) -> list[tuple[np.ndarray, float]]:
    """창고 폴더의 도엽 zip 등고선 → (꼭짓점 · 높이) 목록."""
    from B04_PreProcess.B04_PreProcess_Engine_SheetParser import SHEET_LAYERS, _read_layer_from_zip
    from B04_PreProcess.B04_PreProcess_Engine_SheetStore import _load_index
    from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import iter_linestrings

    prefix, code = SHEET_LAYERS["등고선"]
    lines: list[tuple[np.ndarray, float]] = []
    for no in sheet_nos:
        folder = TERRAIN_ROOT / no
        entry = _load_index(folder)["sheets"].get(no)
        if not entry:
            continue
        frame = _read_layer_from_zip(folder / entry["file"], prefix, code, entry["crs"])
        if frame is None:
            continue
        frame = frame.to_crs(f"EPSG:{int(epsg)}")
        column = next((k for k in ELEVATION_KEYS if k in frame.columns), None)
        if column is None:
            continue
        for geometry, level in zip(frame.geometry, frame[column]):
            if geometry is None or level is None:
                continue
            for line in iter_linestrings(geometry):
                lines.append((np.asarray(line.coords)[:, :2], float(level)))
    return lines


def _fetch_zips(sheet_no: str, report: Callable[[int, str], None]) -> tuple[dict, str | None]:
    """빠진 곳 — 도엽 + 이웃 8 장 zip 을 창고 폴더마다 받음. ({도엽: 해시}, 실패 까닭)."""
    from B04_PreProcess.B04_PreProcess_Engine_SheetStore import ensure_sheets

    zips: dict[str, str | None] = {}
    neighbors = neighbors_3x3(sheet_no)
    for i, no in enumerate([sheet_no] + [n for n in neighbors if n != sheet_no]):
        report(10 + 40 * i // 9, f"도엽 받기 {i + 1}/9 · {no}")
        result = ensure_sheets(TERRAIN_ROOT / no, [no])
        if no in result["available"]:
            zips[no] = _sha256(TERRAIN_ROOT / no / f"{no}.zip")
        elif no == sheet_no:
            reason = result["failed"].get(no, "")
            # 검색 결과 없음 = 바다 등 진짜 빈 도엽 · 그 밖(로그인 · 망)은 다음에 다시
            return zips, None if reason.startswith("LookupError") else reason
    return zips, None


def ensure_sheet_terrain(
    sheet_no: str, epsg: int, cell_m: float = 2.0, on_progress: Progress = None
) -> Path | None:
    """도엽 하나의 창고 지형 — 있으면 그대로 · 없거나 원천이 바뀌었으면 셈. 빈 도엽이면 None.

    받기 실패(로그인 · 망)는 RuntimeError — 저장 안 함(다음에 다시).
    """
    sheet_no = str(sheet_no)

    def report(pct: int, stage: str) -> None:
        if on_progress is not None:
            on_progress(pct, stage)

    npz, key_path, lock_path = _names(sheet_no, epsg, cell_m)
    stored = _stored_ok(sheet_no, epsg, cell_m)
    if stored is not None:
        return None if stored.get("empty") else npz
    with _Lock(lock_path):
        stored = _stored_ok(sheet_no, epsg, cell_m)  # 기다리는 동안 남이 셌을 수 있음
        if stored is not None:
            return None if stored.get("empty") else npz
        started = time.monotonic()
        report(2, f"전국 등고선 읽기 · {sheet_no}")
        ix0, iy0, own = _owned(sheet_no, epsg, cell_m)
        pad = math.ceil(TERRAIN_MARGIN_M / cell_m)
        rows, cols = own.shape
        spec = GridSpec(
            x_min=(ix0 - pad) * cell_m,
            y_max=(iy0 + rows + pad) * cell_m,
            cell_m=cell_m,
            n_rows=rows + 2 * pad,
            n_cols=cols + 2 * pad,
        )
        region = (
            spec.x_min,
            spec.y_max - spec.n_rows * cell_m,
            spec.x_min + spec.n_cols * cell_m,
            spec.y_max,
        )
        a, b, z = _gpkg_segments(epsg, region)
        x_lo, y_lo = ix0 * cell_m, iy0 * cell_m
        x_hi, y_hi = (ix0 + cols) * cell_m, (iy0 + rows) * cell_m
        has_own = bool(
            np.any(
                (np.maximum(a[:, 0], b[:, 0]) >= x_lo)
                & (np.minimum(a[:, 0], b[:, 0]) <= x_hi)
                & (np.maximum(a[:, 1], b[:, 1]) >= y_lo)
                & (np.minimum(a[:, 1], b[:, 1]) <= y_hi)
            )
        )
        zips: dict[str, str | None] = {}
        if not has_own:
            zips, failure = _fetch_zips(sheet_no, report)
            if failure:
                raise RuntimeError(f"도엽 {sheet_no} 받기 실패: {failure}")
            za, zb, zz = segments_from_lines(_zip_lines(list(zips), epsg))
            a, b, z = np.vstack([a, za]), np.vstack([b, zb]), np.concatenate([z, zz])
        report(55, f"지형 계산 · {sheet_no}")
        surface = build_contour_surface(spec, a, b, z) if len(z) else None
        key = _key(epsg, cell_m, zips)
        key_path.unlink(missing_ok=True)  # 쓰는 도중 끊기면 옛 열쇠가 새 파일을 덮지 않게
        if surface is None:
            npz.unlink(missing_ok=True)
            atomic_write_json(key_path, {**key, "empty": True})
            report(100, f"빈 도엽 · {sheet_no}")
            return None
        core = surface[pad : pad + rows, pad : pad + cols][::-1, :].astype(np.float32)
        core[~own] = np.nan
        report(90, f"저장 · {sheet_no}")
        atomic_write_npz(
            npz, z=core, ix0=np.int64(ix0), iy0=np.int64(iy0), cell_m=np.float64(cell_m)
        )
        atomic_write_json(
            key_path,
            {
                **key,
                "source": "zip" if zips else "gpkg",
                "seconds": round(time.monotonic() - started, 1),
            },
        )
        report(100, f"완료 · {sheet_no}")
        logger.info(
            "공용 지형 %s (EPSG:%s · %gm) %.1fs", sheet_no, epsg, cell_m, time.monotonic() - started
        )
        return npz


def load_terrain(
    epsg: int,
    *,
    sheets: list[str] | None = None,
    box: Box | None = None,
    cell_m: float = 2.0,
    on_progress: Progress = None,
) -> dict[str, Any]:
    """도엽 목록 또는 TM 상자 → 이어 붙인 격자(x · y 오름차순 · z (ny, nx) · valid_mask) + 메타."""
    if sheets is None:
        if box is None:
            raise ValueError("sheets 또는 box 가 필요함")
        sheets = sheets_for_box(box, epsg)
    sheets = [str(s) for s in sheets]
    pieces: list[tuple[int, int, np.ndarray]] = []
    computed: list[str] = []
    failed: dict[str, str] = {}
    total = max(len(sheets), 1)
    for i, no in enumerate(sheets):
        fresh = _stored_ok(no, epsg, cell_m) is None

        def inner(pct: int, stage: str, i: int = i) -> None:
            if on_progress is not None:
                on_progress(int((i * 100 + pct) / total), f"{stage} ({i + 1}/{total})")

        try:
            path = ensure_sheet_terrain(no, epsg, cell_m, inner if on_progress else None)
        except Exception as exc:  # noqa: BLE001 — 한 장이 죽어도 나머지는 이어 붙임
            failed[no] = f"{type(exc).__name__}: {exc}"
            continue
        if fresh:
            computed.append(no)
        if path is None:
            continue
        with np.load(path) as data:
            pieces.append((int(data["ix0"]), int(data["iy0"]), data["z"]))
    if box is not None:
        ix0, iy0 = math.ceil(box[0] / cell_m - 0.5), math.ceil(box[1] / cell_m - 0.5)
        ix1, iy1 = math.floor(box[2] / cell_m - 0.5) + 1, math.floor(box[3] / cell_m - 0.5) + 1
    elif pieces:
        ix0 = min(p[0] for p in pieces)
        iy0 = min(p[1] for p in pieces)
        ix1 = max(p[0] + p[2].shape[1] for p in pieces)
        iy1 = max(p[1] + p[2].shape[0] for p in pieces)
    else:
        ix0 = iy0 = ix1 = iy1 = 0
    z = np.full((max(iy1 - iy0, 0), max(ix1 - ix0, 0)), np.nan, dtype=np.float32)
    for px, py, pz in pieces:
        r0, c0 = max(py, iy0), max(px, ix0)
        r1, c1 = min(py + pz.shape[0], iy1), min(px + pz.shape[1], ix1)
        if r0 >= r1 or c0 >= c1:
            continue
        src = pz[r0 - py : r1 - py, c0 - px : c1 - px]
        dst = z[r0 - iy0 : r1 - iy0, c0 - ix0 : c1 - ix0]
        np.copyto(dst, src, where=np.isfinite(src))
    if on_progress is not None:
        on_progress(100, "지형 준비 끝")
    return {
        "x": (np.arange(ix0, ix1) + 0.5) * cell_m,
        "y": (np.arange(iy0, iy1) + 0.5) * cell_m,
        "z": z,
        "valid_mask": np.isfinite(z),
        "cell_m": float(cell_m),
        "epsg": int(epsg),
        "sheets": sheets,
        "computed": computed,
        "failed": failed,
    }