geo_tools/app/io/writers.py

"""
geo_tools.io.writers
~~~~~~~~~~~~~~~~~~~~~
统一的矢量数据写出接口，支持：
- Shapefile (.shp)
- GeoJSON (.geojson / .json)
- GeoPackage (.gpkg)       ← 支持追加图层
- File Geodatabase (.gdb)  ← 通过 fiona OpenFileGDB 驱动
- FlatGeobuf (.fgb)
- CSV（含 WKT 列）
"""

from __future__ import annotations

from pathlib import Path
from typing import Any, Literal

import geopandas as gpd

from app.utils.logger import get_logger

logger = get_logger(__name__)


# ── 主入口 ──────────────────────────────────────────────────────────────────────

def write_vector(
    gdf: gpd.GeoDataFrame,
    path: str | Path,
    layer: str | None = None,
    driver: str | None = None,
    encoding: str = "utf-8",
    mode: Literal["w", "a"] = "w",
    **kwargs: Any,
) -> Path:
    """统一的矢量数据写出入口，自动识别格式。

    Parameters
    ----------
    gdf:
        待写出的 GeoDataFrame。
    path:
        目标路径（文件或 `.gdb` 目录）。
    layer:
        图层名（GPKG、GDB 多图层格式时使用）。
    driver:
        强制指定 fiona 驱动名（通常不需要，自动推断）。
    encoding:
        字段编码，Shapefile 导出中文时常需 ``"gbk"``。
    mode:
        ``"w"`` 覆盖写出，``"a"`` 追加图层（GPKG / GDB 支持）。

    Returns
    -------
    Path
        实际写出的路径。
    """
    path = Path(path)
    suffix = path.suffix.lower()

    if suffix == ".csv":
        return _write_csv_vector(gdf, path)

    # 自动推断驱动
    if driver is None:
        driver = _infer_driver(path)

    # 确保父目录存在（GDB 是目录，其父目录要存在）
    if suffix == ".gdb":
        path.parent.mkdir(parents=True, exist_ok=True)
    else:
        path.parent.mkdir(parents=True, exist_ok=True)

    write_kwargs: dict[str, Any] = {
        "driver": driver,
        "encoding": encoding,
        "mode": mode,
        **kwargs,
    }
    if layer is not None:
        write_kwargs["layer"] = layer

    logger.info(
        "写出矢量数据：%s（驱动：%s，图层：%s，模式：%s，要素数：%d）",
        path, driver, layer, mode, len(gdf),
    )
    gdf.to_file(str(path), **write_kwargs)
    logger.info("写出完成：%s", path)
    return path


# ── GDB 专用 ────────────────────────────────────────────────────────────────────

def write_gdb(
    gdf: gpd.GeoDataFrame,
    gdb_path: str | Path,
    layer: str,
    mode: Literal["w", "a"] = "w",
    encoding: str = "utf-8",
    **kwargs: Any,
) -> Path:
    """将 GeoDataFrame 写出到 Esri File Geodatabase（.gdb）中。

    Parameters
    ----------
    gdf:
        待写出的 GeoDataFrame。
    gdb_path:
        目标 ``.gdb`` 目录路径（不存在时自动创建）。
    layer:
        图层名称（必填）。
    mode:
        ``"w"`` 覆盖图层；``"a"`` 向已有 GDB 追加图层。

    Notes
    -----
    写出 GDB 依赖 fiona 的 ``OpenFileGDB``（写）或 ``FileGDB``（需 ESRI 驱动）支持。
    当前 fiona >= 1.9 的 ``OpenFileGDB`` 驱动已支持创建和写出，无需额外安装。
    """
    gdb_path = Path(gdb_path)
    if not layer:
        raise ValueError("写出 GDB 必须指定 layer 参数。")

    gdb_path.parent.mkdir(parents=True, exist_ok=True)

    logger.info("写出到 GDB：%s >> 图层 %r（模式：%s）", gdb_path.name, layer, mode)
    gdf.to_file(
        str(gdb_path),
        layer=layer,
        driver="OpenFileGDB",
        mode=mode,
        encoding=encoding,
        **kwargs,
    )
    logger.info("GDB 写出完成：%s >> %s", gdb_path, layer)
    return gdb_path


# ── GPKG 专用 ───────────────────────────────────────────────────────────────────

def write_gpkg(
    gdf: gpd.GeoDataFrame,
    gpkg_path: str | Path,
    layer: str,
    mode: Literal["w", "a"] = "w",
    **kwargs: Any,
) -> Path:
    """将 GeoDataFrame 写出为 GeoPackage 中的一个图层。

    Parameters
    ----------
    gpkg_path:
        目标 ``.gpkg`` 文件路径（不存在时自动创建）。
    layer:
        图层名称（必填）。
    mode:
        ``"w"`` 覆盖；``"a"`` 向已有 GPKG 追加图层。
    """
    gpkg_path = Path(gpkg_path)
    gpkg_path.parent.mkdir(parents=True, exist_ok=True)
    gdf.to_file(str(gpkg_path), layer=layer, driver="GPKG", mode=mode, **kwargs)
    logger.info("GPKG 写出完成：%s >> %s", gpkg_path, layer)
    return gpkg_path


# ── CSV 写出 ────────────────────────────────────────────────────────────────────

def _write_csv_vector(gdf: gpd.GeoDataFrame, path: Path, **kwargs: Any) -> Path:
    """将 GeoDataFrame 写出为含 WKT 几何列的 CSV。"""
    path.parent.mkdir(parents=True, exist_ok=True)
    df = gdf.copy()
    df["geometry"] = df["geometry"].apply(lambda g: g.wkt if g is not None else None)
    df.to_csv(path, index=False, encoding="utf-8-sig", **kwargs)  # utf-8-sig 兼容 Excel
    logger.info("CSV 写出完成：%s", path)
    return path


def write_csv(gdf: gpd.GeoDataFrame, path: str | Path, **kwargs: Any) -> Path:
    """将 GeoDataFrame 写出为含 WKT 几何列的 CSV（公开接口）。"""
    return _write_csv_vector(gdf, Path(path), **kwargs)


# ── 工具函数 ─────────────────────────────────────────────────────────────────────

def _infer_driver(path: Path) -> str:
    """根据文件扩展名推断 fiona 驱动。"""
    _EXT_TO_DRIVER: dict[str, str] = {
        ".shp": "ESRI Shapefile",
        ".geojson": "GeoJSON",
        ".json": "GeoJSON",
        ".gpkg": "GPKG",
        ".gdb": "OpenFileGDB",
        ".kml": "KML",
        ".fgb": "FlatGeobuf",
        ".gml": "GML",
        ".dxf": "DXF",
    }
    suffix = path.suffix.lower()
    if path.is_dir() and suffix == ".gdb":
        return "OpenFileGDB"
    driver = _EXT_TO_DRIVER.get(suffix)
    if driver is None:
        raise ValueError(
            f"无法自动推断 fiona 驱动，未知扩展名：{suffix!r}。"
            f"请显式传入 driver=... 参数。"
        )
    return driver