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 datetime import datetime
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 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": datetime.now().isoformat(timespec="seconds"),
        "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)