Aislo/B04_PreProcess/B04_PreProcess_Engine_Watershed_Export.py

"""배수유역 단계별 검증 산출물을 영구저장소에 남긴다.

기능을 하나씩 붙일 때마다 그 단계의 결과를 파일로 남겨 사람이 QGIS 등으로 직접 열어
대조할 수 있게 하는 것이 목적이다(2026-07-31 사용자 지시). 새 단계를 추가할 때는
`STAGES`에 이름을 하나 더 넣고 `write_stage()`를 호출하면 된다 — 파일명 규칙과 매니페스트
갱신은 여기서 일괄로 처리한다.

저장 위치: `storage/{회사}/{사용자}/{프로젝트}/B04_PreProcess/drainage/`
  - `{단계번호}_{단계이름}.geojson` — WGS84 FeatureCollection, 피처마다 `kind` 속성
  - `manifest.json` — 지금까지 남긴 단계 목록과 요약값
"""

from __future__ import annotations

import json
import logging
from collections.abc import Sequence
from pathlib import Path
from typing import Any, Callable

import numpy as np
from shapely.geometry.base import BaseGeometry

from common_util.common_util_storage import resolve_stored_project_path
from common_util.common_util_time import utc_now_text
from config.config_system import DRAINAGE_CACHE_DIRNAME

logger = logging.getLogger(__name__)

# 단계 이름 → 파일 접두 번호. 순서대로 읽으면 파이프라인 진행 순서가 된다.
STAGES: dict[str, str] = {
    "primary_region": "01",
    "flow_direction": "02",
    # B05가 읽어 세부유역을 나누는 데 필요한 최소 배열·기하. 화살표·표고는 넣지 않는다.
    "road_routing": "03",
}

_MANIFEST_FILENAME = "manifest.json"
LonLat = Callable[[float, float], tuple[float, float]]


def drainage_dir(stored_path: str) -> Path:
    return (
        Path(resolve_stored_project_path(stored_path)) / "B04_PreProcess" / DRAINAGE_CACHE_DIRNAME
    )


def write_stage(
    stored_path: str,
    stage: str,
    layers: dict[str, Sequence[BaseGeometry | tuple[BaseGeometry, dict[str, Any]]]],
    properties: dict[str, Any],
    to_lonlat: LonLat,
) -> str | None:
    """한 단계의 기하 산출물을 GeoJSON으로 저장하고 매니페스트를 갱신한다.

    `layers`는 {레이어이름: 사업지 CRS(m) 기하 목록}이며 레이어 이름이 피처 `kind` 속성이
    된다. 기하 대신 `(기하, 속성dict)` 짝을 넣으면 그 속성이 피처에 함께 실린다.
    좌표는 여기서 WGS84로 바꾼다 — 저장 파일은 어떤 도구로 열어도 바로 보여야 한다.
    """
    prefix = STAGES.get(stage)
    if prefix is None:
        logger.warning("배수유역: 등록되지 않은 저장 단계 '%s' — 저장을 건너뜁니다.", stage)
        return None

    features: list[dict[str, Any]] = []
    counts: dict[str, int] = {}
    for kind, entries in layers.items():
        for index, entry in enumerate(entries):
            # 항목은 기하 하나이거나 (기하, 속성) 짝이다 — 관 누가거리처럼 붙일 값이 있을 때 쓴다.
            geometry, extra = entry if isinstance(entry, tuple) else (entry, None)
            feature = _to_feature(kind, index, geometry, to_lonlat, extra)
            if feature is not None:
                features.append(feature)
        counts[kind] = len(entries)

    filename = f"{prefix}_{stage}.geojson"
    directory = drainage_dir(stored_path)
    target = directory / filename
    document = {
        "type": "FeatureCollection",
        "crs": {"type": "name", "properties": {"name": "urn:ogc:def:crs:OGC:1.3:CRS84"}},
        "properties": {**properties, "counts": counts},
        "features": features,
    }
    try:
        directory.mkdir(parents=True, exist_ok=True)
        with target.open("w", encoding="utf-8") as file:
            json.dump(document, file, ensure_ascii=False)
    except OSError:
        logger.warning("배수유역: %s 저장 실패 (%s)", stage, target)
        return None

    _update_manifest(directory, stage, filename, {**properties, "counts": counts})
    logger.info("배수유역: %s 저장 — %s (피처 %d개)", stage, target, len(features))
    return str(target)


def _to_feature(
    kind: str,
    index: int,
    geometry: BaseGeometry,
    to_lonlat: LonLat,
    extra: dict[str, Any] | None = None,
) -> dict[str, Any] | None:
    coordinates = _to_lonlat_coords(geometry, to_lonlat)
    if coordinates is None:
        return None
    return {
        "type": "Feature",
        "properties": {"kind": kind, "index": index, **(extra or {})},
        "geometry": {"type": geometry.geom_type, "coordinates": coordinates},
    }


def _to_lonlat_coords(geometry: BaseGeometry, to_lonlat: LonLat) -> Any:
    """shapely 기하를 WGS84 GeoJSON 좌표 배열로 바꾼다."""
    if geometry.is_empty:
        return None
    kind = geometry.geom_type
    if kind == "Point":
        return list(to_lonlat(geometry.x, geometry.y))
    if kind == "LineString":
        return [list(to_lonlat(x, y)) for x, y in geometry.coords]
    if kind == "Polygon":
        return [
            [list(to_lonlat(x, y)) for x, y in ring.coords]
            for ring in (geometry.exterior, *geometry.interiors)
        ]
    if kind in {"MultiPoint", "MultiLineString", "MultiPolygon", "GeometryCollection"}:
        parts = [_to_lonlat_coords(part, to_lonlat) for part in geometry.geoms]
        return [part for part in parts if part is not None]
    return None


def write_grid_arrays(
    stored_path: str, stage: str, spec: Any, arrays: dict[str, Any], summary: dict[str, Any]
) -> str | None:
    """격자 크기의 배열들을 `.npz`로 남긴다(셀 마스크·흐름 방향·도달 여부 등).

    셀이 수십만 개라 GeoJSON 폴리곤으로는 못 남긴다. 격자 원점·셀 크기와 배열만 저장하면
    어느 셀이 어떤 값이었는지 그대로 복원된다. 요약값은 manifest에도 기록한다.
    """
    prefix = STAGES.get(stage)
    if prefix is None or not arrays:
        return None
    directory = drainage_dir(stored_path)
    target = directory / f"{prefix}_{stage}.npz"
    try:
        directory.mkdir(parents=True, exist_ok=True)
        np.savez_compressed(
            target,
            x_min=spec.x_min,
            y_max=spec.y_max,
            cell_m=spec.cell_m,
            n_rows=spec.n_rows,
            n_cols=spec.n_cols,
            **arrays,
        )
    except OSError:
        logger.warning("배수유역: %s 배열 저장 실패 (%s)", stage, target)
        return None
    _update_manifest(directory, f"{stage}_arrays", target.name, summary)
    logger.info("배수유역: %s 배열 저장 — %s (%s)", stage, target, ", ".join(arrays))
    return str(target)


def _update_manifest(
    directory: Path, stage: str, filename: str, properties: dict[str, Any]
) -> None:
    """지금까지 남긴 단계 목록을 한 파일에 모아 둔다 — 무엇이 저장돼 있는지 한눈에 본다."""
    manifest_path = directory / _MANIFEST_FILENAME
    manifest: dict[str, Any] = {}
    if manifest_path.exists():
        try:
            with manifest_path.open("r", encoding="utf-8") as file:
                loaded = json.load(file)
            if isinstance(loaded, dict):
                manifest = loaded
        except (OSError, json.JSONDecodeError):
            logger.warning("배수유역: manifest를 읽지 못해 새로 만듭니다 (%s).", manifest_path)
    manifest[stage] = {
        "file": filename,
        "saved_at": utc_now_text(),
        "properties": properties,
    }
    try:
        with manifest_path.open("w", encoding="utf-8") as file:
            json.dump(manifest, file, ensure_ascii=False, indent=2)
    except OSError:
        logger.warning("배수유역: manifest 저장 실패 (%s).", manifest_path)