"""도엽별 공용 하강장(609-1) — 배수 하강 방향장을 도엽마다 한 번 짓고 공용 창고에 영구 저장.

자리 = `storage/_shared/terrain/<도엽>/descent_<좌표계>_<셀>m.npz` + `.key.json`(607 지형 창고 옆).
원천 = 전국 등고선 · 그 도엽에 없으면 607 길로 받은 도엽 zip. 도엽 + 여유 100 m 로 짓고(짓는 법은
`build_contour_descent` 그대로 — 1 m · 같은 굽기 · 같은 EDT) 도엽 안쪽 칸(칸 중심 위경도)만 남김 —
칸마다 방향 코드(0~7 = `_NEIGHBOURS` · 8 = 제자리 · -1 = 무효) + 밴드 표고. 프로젝트는 격자를 덮는
도엽을 이어 붙여 읽기만 — 코드라 도엽 경계를 넘는 물길도 그대로 이어짐. 짓는 규칙이 바뀌면 판을
올림.
"""

from __future__ import annotations

import hashlib
import json
import logging
import math
import time
from pathlib import Path
from typing import Any

import numpy as np
from pyproj import CRS, Transformer
from pyproj.exceptions import CRSError
from shapely.geometry import LineString, mapping, shape

from B04_PreProcess.B04_PreProcess_Engine_MapSheet import sheet5k_to_bounds
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Descent import (
    _NEIGHBOURS,
    DESCENT_TILE_CELLS,
    DESCENT_WINDOW_MARGIN_CELLS,
    ContourDescent,
    build_contour_descent,
)
from B04_PreProcess.B04_PreProcess_Engine_Watershed_Grid import (
    AZIMUTH_INVALID,
    AZIMUTH_SINK,
    AZIMUTH_STEPS,
    GridSpec,
)
from common_util.common_util_atomic import atomic_write_npz
from common_util.common_util_crs import tm_epsg_for_lon
from common_util.common_util_json import atomic_write_json
from common_util.common_util_terrain_store import (
    NATIONAL_GPKG,
    TERRAIN_ROOT,
    _crs,
    _gpkg_stamp,
    _Lock,
    _names,
    _owned,
    _sha256,
    _zip_lines,
    ensure_sheet_terrain,
    sheets_for_box,
)
from config.config_system import DRAINAGE_CONTOUR_MIN_LENGTH_M

logger = logging.getLogger(__name__)

SHARED_DESCENT_VERSION = 1
SHARED_DESCENT_MARGIN_M = 100.0
_CODE_SINK = 8
# 3×3 자리 (dr+1)*3 + (dc+1) → 코드 · 코드 → (dr, dc) · 방위
_CODE_OF_SHIFT = np.array([0, 1, 2, 3, _CODE_SINK, 4, 5, 6, 7], dtype=np.int8)
_SHIFT_ROW = np.array([r for r, _ in _NEIGHBOURS] + [0])
_SHIFT_COL = np.array([c for _, c in _NEIGHBOURS] + [0])
_CODE_AZIMUTH = np.append(
    np.rint(np.arctan2(_SHIFT_ROW[:8], _SHIFT_COL[:8]) / (2.0 * math.pi / AZIMUTH_STEPS)).astype(
        np.int16
    )
    % AZIMUTH_STEPS,
    AZIMUTH_SINK,
).astype(np.int16)


Crs = int | str  # EPSG 번호 또는 번호 없는 좌표계 글(PRJ WKT — 용화 등 Korean 1985 + TOWGS84)


def descent_crs(crs: str) -> Crs | None:
    """프로젝트 좌표계 → EPSG 번호(있으면) 또는 그 글 · 못 읽으면 None(프로젝트 저장본 길)."""
    try:
        return CRS.from_user_input(crs).to_epsg() or crs
    except CRSError:
        return None


def _crs_tag(crs: Crs) -> str:
    """파일 이름 · 열쇠의 좌표계 — 번호 또는 WKT 지문(같은 PRJ 를 쓰는 프로젝트끼리 같음)."""
    if isinstance(crs, int) or str(crs).isdigit():
        return str(int(crs))
    wkt = CRS.from_user_input(crs).to_wkt()
    return "wkt" + hashlib.sha1(wkt.encode(), usedforsecurity=False).hexdigest()[:10]


def _shared_names(sheet_no: str, crs: Crs, cell_m: float) -> tuple[Path, Path, Path]:
    stem = f"descent_{_crs_tag(crs)}_{cell_m:g}m"
    folder = TERRAIN_ROOT / str(sheet_no)
    return folder / f"{stem}.npz", folder / f"{stem}.key.json", folder / f".{stem}.lock"


def _shared_key(crs: Crs, cell_m: float, zips: dict[str, str | None]) -> dict[str, Any]:
    return {
        "version": SHARED_DESCENT_VERSION,
        "margin_m": SHARED_DESCENT_MARGIN_M,
        "rule": [DRAINAGE_CONTOUR_MIN_LENGTH_M, DESCENT_WINDOW_MARGIN_CELLS, DESCENT_TILE_CELLS],
        "crs": _crs_tag(crs),
        "cell_m": float(cell_m),
        "gpkg": _gpkg_stamp(),
        "zips": zips,
    }


def _shared_ok(sheet_no: str, crs: Crs, cell_m: float) -> dict[str, Any] | None:
    """저장 열쇠가 지금 원천 · 규칙과 같으면 그 열쇠 — 아니면 None(다시 지음)."""
    npz, key_path, _ = _shared_names(sheet_no, crs, 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 = _shared_key(crs, cell_m, zips)
    if now["gpkg"] is None:  # 전국 파일 없는 PC — 있는 저장본을 믿음(607 과 같음)
        now["gpkg"] = stored.get("gpkg")
    return stored if {k: stored.get(k) for k in now} == now else None


def _feature(line: Any, level: float) -> dict[str, Any]:
    return {"geometry": mapping(line), "properties": {"elevation": float(level)}}


def _gpkg_features(crs: Crs, region: tuple[float, float, float, float]) -> list[dict[str, Any]]:
    """전국 등고선에서 TM 상자에 닿는 선 — 통째 선(길이 거름이 프로젝트 길과 같게) · 사업 좌표."""
    if not NATIONAL_GPKG.is_file():
        return []
    import pyogrio
    import shapely

    x0, y0, x1, y1 = region
    t = np.linspace(0.0, 1.0, 21)
    bx, by = Transformer.from_crs(_crs(crs), "EPSG:5179", always_xy=True).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]),
    )
    frame = pyogrio.read_dataframe(
        NATIONAL_GPKG,
        bbox=(bx.min() - 20.0, by.min() - 20.0, bx.max() + 20.0, by.max() + 20.0),
        columns=["CTRLN_HG"],
    )
    if frame.empty:
        return []
    to_tm = Transformer.from_crs("EPSG:5179", _crs(crs), always_xy=True)
    lines = shapely.transform(
        frame.geometry.values, lambda xy: np.column_stack(to_tm.transform(xy[:, 0], xy[:, 1]))
    )
    levels = frame["CTRLN_HG"].to_numpy(dtype=np.float64)
    return [
        _feature(line, level)
        for line, level in zip(lines, levels)
        if line is not None and np.isfinite(level)
    ]


def _touches(features: list[dict[str, Any]], box: tuple[float, float, float, float]) -> bool:
    x0, y0, x1, y1 = box
    for feature in features:
        bx0, by0, bx1, by1 = shape(feature["geometry"]).bounds
        if bx1 >= x0 and bx0 <= x1 and by1 >= y0 and by0 <= y1:
            return True
    return False


def _codes(descent: ContourDescent) -> np.ndarray:
    """방향장 → 칸 코드 (R, C) int8."""
    spec = descent.spec
    index = np.arange(spec.size, dtype=np.int64)
    receiver = descent.receiver.astype(np.int64)
    shift = (receiver // spec.n_cols - index // spec.n_cols + 1) * 3 + (
        receiver % spec.n_cols - index % spec.n_cols + 1
    )
    code = _CODE_OF_SHIFT[shift]
    code[~descent.valid.reshape(-1)] = -1
    return code.reshape(spec.n_rows, spec.n_cols)


def ensure_sheet_descent(sheet_no: str, crs: Crs, cell_m: float = 1.0) -> Path | None:
    """도엽 하나의 공용 하강장 — 있으면 그대로 · 없거나 원천이 바뀌었으면 지음. 빈 도엽이면 None.

    빠진 곳 받기 실패(로그인 · 망)는 RuntimeError — 저장 안 함(다음에 다시).
    """
    sheet_no = str(sheet_no)
    npz, key_path, lock_path = _shared_names(sheet_no, crs, cell_m)
    if (stored := _shared_ok(sheet_no, crs, cell_m)) is not None:
        return None if stored.get("empty") else npz
    with _Lock(lock_path):
        if (stored := _shared_ok(sheet_no, crs, cell_m)) is not None:
            return None if stored.get("empty") else npz
        started = time.monotonic()
        ix0, iy0, own = _owned(sheet_no, crs, cell_m)
        pad = math.ceil(SHARED_DESCENT_MARGIN_M / cell_m)
        rows, cols = own.shape
        spec = GridSpec(
            (ix0 - pad) * cell_m,
            (iy0 + rows + pad) * cell_m,
            cell_m,
            rows + 2 * pad,
            cols + 2 * pad,
        )
        x_max, y_min = spec.x_min + spec.n_cols * cell_m, spec.y_max - spec.n_rows * cell_m
        features = _gpkg_features(crs, (spec.x_min, y_min, x_max, spec.y_max))
        zips: dict[str, str | None] = {}
        own_box = (ix0 * cell_m, iy0 * cell_m, (ix0 + cols) * cell_m, (iy0 + rows) * cell_m)
        if not _touches(features, own_box):
            # 빠진 곳(섬 등) — 607 길로 도엽 zip 받기 · 그 열쇠의 zip 을 원천으로
            lon0, _, lon1, _ = sheet5k_to_bounds(sheet_no)
            tm = crs if isinstance(crs, int) else tm_epsg_for_lon((lon0 + lon1) / 2)
            ensure_sheet_terrain(sheet_no, tm)
            terrain_key = json.loads(_names(sheet_no, tm, 2.0)[1].read_text(encoding="utf-8"))
            zips = terrain_key.get("zips") or {}
            features += [_feature(LineString(xy), z) for xy, z in _zip_lines(list(zips), crs)]
        descent = build_contour_descent(spec, features) if features else None
        key = _shared_key(crs, cell_m, zips)
        key_path.unlink(missing_ok=True)  # 쓰는 도중 끊기면 옛 열쇠가 새 파일을 덮지 않게
        if descent is None or not descent.valid.any():
            npz.unlink(missing_ok=True)
            atomic_write_json(key_path, {**key, "empty": True})
            return None
        inner = (slice(pad, pad + rows), slice(pad, pad + cols))
        code = np.where(own[::-1, :], _codes(descent)[inner], -1).astype(np.int8)
        band = np.where(code >= 0, descent.band_elevation[inner], np.nan).astype(np.float32)
        atomic_write_npz(
            npz,
            code=code,
            band=band,
            origin=np.array([ix0 * cell_m, (iy0 + rows) * cell_m, cell_m]),
            levels=np.array(descent.levels, dtype=np.float64),
        )
        seconds = round(time.monotonic() - started, 1)
        atomic_write_json(
            key_path, {**key, "source": "zip" if zips else "gpkg", "seconds": seconds}
        )
        logger.info("공용 하강장 %s (%s · %gm) %.1fs", sheet_no, _crs_tag(crs), cell_m, seconds)
        return npz


def shared_descent(crs: Crs, spec: GridSpec) -> ContourDescent | None:
    """`spec` 격자를 덮는 도엽 하강장을 이어 붙임(없는 도엽은 지음).

    못 지은 도엽 · 칸 어긋남 = None(프로젝트 길로).
    """

    started = time.perf_counter()
    rows, cols, cell_m = spec.n_rows, spec.n_cols, spec.cell_m
    code = np.full((rows, cols), -1, dtype=np.int8)
    band = np.full((rows, cols), np.nan, dtype=np.float32)
    levels: set[float] = set()
    box = (spec.x_min, spec.y_max - rows * cell_m, spec.x_min + cols * cell_m, spec.y_max)
    for sheet_no in sheets_for_box(box, crs):
        try:
            path = ensure_sheet_descent(sheet_no, crs, cell_m)
        except Exception:  # noqa: BLE001 — 한 장이라도 못 지으면 프로젝트 길로
            logger.exception(
                "배수유역: 공용 하강장 %s 을 짓지 못해 프로젝트 길로 갑니다.", sheet_no
            )
            return None
        if path is None:
            continue
        with np.load(path) as data:
            x_min, y_max, _ = data["origin"].tolist()
            row_shift = (spec.y_max - y_max) / cell_m
            col_shift = (x_min - spec.x_min) / cell_m
            row0, col0 = int(round(row_shift)), int(round(col_shift))
            if abs(row_shift - row0) > 1e-6 or abs(col_shift - col0) > 1e-6:
                return None
            piece, piece_band = data["code"], data["band"]
            levels.update(data["levels"].tolist())
        r0, c0 = max(row0, 0), max(col0, 0)
        r1, c1 = min(row0 + piece.shape[0], rows), min(col0 + piece.shape[1], cols)
        if r0 >= r1 or c0 >= c1:
            continue
        src = (slice(r0 - row0, r1 - row0), slice(c0 - col0, c1 - col0))
        mine = piece[src] >= 0
        np.copyto(code[r0:r1, c0:c1], piece[src], where=mine)
        np.copyto(band[r0:r1, c0:c1], piece_band[src], where=mine)
    descent = _from_codes(spec, code, band, sorted(levels, reverse=True))
    logger.info(
        "[계측] 공용 하강장 이어 붙임 %.2fs · %d×%d", time.perf_counter() - started, rows, cols
    )
    return descent


def _from_codes(
    spec: GridSpec, code: np.ndarray, band: np.ndarray, levels: list[float]
) -> ContourDescent:
    """칸 코드 → 방향장. 격자 밖을 가리키는 칸은 제자리(`crop_descent` 와 같은 판정)."""
    rows, cols = spec.n_rows, spec.n_cols
    valid = code >= 0
    safe = np.where(valid, code, _CODE_SINK)
    target_row = np.arange(rows)[:, None] + _SHIFT_ROW[safe]
    target_col = np.arange(cols)[None, :] + _SHIFT_COL[safe]
    moves = (
        valid
        & (safe != _CODE_SINK)
        & (target_row >= 0)
        & (target_row < rows)
        & (target_col >= 0)
        & (target_col < cols)
    )
    flat = np.arange(spec.size, dtype=np.int64).reshape(rows, cols)
    receiver = np.where(moves, target_row * cols + target_col, flat)
    length = np.hypot(_SHIFT_ROW[safe], _SHIFT_COL[safe]) * spec.cell_m
    return ContourDescent(
        spec=spec,
        band_elevation=band,
        valid=valid,
        receiver=receiver.reshape(-1).astype(np.int32),
        step_length=np.where(moves, length, 0.0).astype(np.float32).reshape(-1),
        azimuth=np.where(valid, _CODE_AZIMUTH[safe], AZIMUTH_INVALID).astype(np.int16).reshape(-1),
        levels=levels,
    )