#!/usr/bin/env python3
"""
NavSea coastline subset builder v1.0

针对 GSI 海岸线数据的福冈 / 佐贺两包（C23-06_40、C23-06_41）生成
一个可复用的海岸线硬阻塞样例底库：

1. 解析 GML 曲线
2. 投影到 Web Mercator 米制坐标
3. 生成海岸线缓冲区
4. 生成 200m fishnet
5. 输出 SQLite + GeoJSON + JSON 元数据

说明：
- 这是 coastline 全国底库的样例切片，不是最终全国版。
- 先把福冈 / 佐贺海岸跑通，便于后续扩全国。
"""

from __future__ import annotations

import argparse
import datetime as dt
import json
import math
import sqlite3
import zipfile
from dataclasses import dataclass
from pathlib import Path
from typing import Iterable
import xml.etree.ElementTree as ET

from shapely.geometry import LineString, MultiLineString, box, mapping
from shapely.geometry.base import BaseGeometry
from shapely.prepared import prep
from shapely.ops import unary_union


RADIUS = 6378137.0
MAX_MERCATOR_LAT = 85.0511287798066
NS = {"gml": "http://www.opengis.net/gml/3.2"}


@dataclass(frozen=True)
class SourcePackage:
    path: Path
    prefecture: str
    code: str


def mercator_x(lon: float) -> float:
    return RADIUS * math.radians(lon)


def mercator_y(lat: float) -> float:
    lat = max(min(lat, MAX_MERCATOR_LAT), -MAX_MERCATOR_LAT)
    return RADIUS * math.log(math.tan(math.pi / 4.0 + math.radians(lat) / 2.0))


def mercator_point(lon: float, lat: float) -> tuple[float, float]:
    return mercator_x(lon), mercator_y(lat)


def align_floor(value: float, step: float) -> float:
    return math.floor(value / step) * step


def align_ceil(value: float, step: float) -> float:
    return math.ceil(value / step) * step


def parse_text_list(text: str) -> list[float]:
    return [float(part) for part in text.split() if part]


def detect_prefecture(meta_xml: str) -> str:
    title_start = meta_xml.find("<title>")
    if title_start == -1:
        return "unknown"
    title_end = meta_xml.find("</title>", title_start)
    if title_end == -1:
        return "unknown"
    return meta_xml[title_start + 7:title_end].strip()


def load_package(path: Path) -> SourcePackage:
    with zipfile.ZipFile(path) as zf:
        meta_name = next(name for name in zf.namelist() if "META" in name and name.endswith(".xml"))
        meta_xml = zf.read(meta_name).decode("shift_jis", errors="replace")
    prefecture = detect_prefecture(meta_xml)
    code = path.stem.replace("_GML", "")
    return SourcePackage(path=path, prefecture=prefecture, code=code)


def iter_coastline_lines(zip_path: Path) -> Iterable[tuple[str, str, LineString]]:
    """
    Yield (source_code, curve_id, projected_linestring).
    """
    with zipfile.ZipFile(zip_path) as zf:
        xml_name = next(
            name
            for name in zf.namelist()
            if name.endswith(".xml") and "META" not in name
        )
        root = ET.fromstring(zf.read(xml_name))

    for idx, curve in enumerate(root.findall(".//gml:Curve", NS), start=1):
        curve_id = curve.attrib.get("{http://www.opengis.net/gml/3.2}id", f"curve_{idx:06d}")
        coords: list[tuple[float, float]] = []
        for pos_list in curve.findall(".//gml:posList", NS):
            if not pos_list.text:
                continue
            values = parse_text_list(pos_list.text)
            if len(values) < 4 or len(values) % 2 != 0:
                continue
            for i in range(0, len(values), 2):
                lat = values[i]
                lon = values[i + 1]
                x, y = mercator_point(lon, lat)
                coords.append((x, y))
        if len(coords) < 2:
            continue
        yield xml_name, curve_id, LineString(coords)


def build_output_dir(path: Path) -> None:
    path.mkdir(parents=True, exist_ok=True)


def write_geojson(path: Path, features: list[dict]) -> None:
    payload = {"type": "FeatureCollection", "features": features}
    path.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8")


def open_sqlite(path: Path) -> sqlite3.Connection:
    if path.exists():
        path.unlink()
    conn = sqlite3.connect(path)
    conn.execute("PRAGMA journal_mode=OFF")
    conn.execute("PRAGMA synchronous=OFF")
    conn.execute("PRAGMA temp_store=MEMORY")
    conn.execute("PRAGMA cache_size=200000")
    conn.execute("PRAGMA locking_mode=EXCLUSIVE")
    conn.execute("PRAGMA foreign_keys=OFF")
    return conn


def create_schema(conn: sqlite3.Connection) -> None:
    conn.executescript(
        """
        CREATE TABLE meta (
          key TEXT PRIMARY KEY,
          value TEXT NOT NULL
        );

        CREATE TABLE coastline_feature (
          feature_id INTEGER PRIMARY KEY AUTOINCREMENT,
          source_code TEXT NOT NULL,
          prefecture TEXT NOT NULL,
          curve_id TEXT NOT NULL,
          geom_wkt TEXT NOT NULL,
          point_count INTEGER NOT NULL,
          minx REAL NOT NULL,
          miny REAL NOT NULL,
          maxx REAL NOT NULL,
          maxy REAL NOT NULL
        );

        CREATE TABLE grid_cell (
          cell_id TEXT PRIMARY KEY,
          source_code TEXT NOT NULL,
          prefecture TEXT NOT NULL,
          ix INTEGER NOT NULL,
          iy INTEGER NOT NULL,
          minx REAL NOT NULL,
          miny REAL NOT NULL,
          maxx REAL NOT NULL,
          maxy REAL NOT NULL,
          state INTEGER NOT NULL,
          state_name TEXT NOT NULL,
          source_mask INTEGER NOT NULL,
          blocked_permille INTEGER NOT NULL
        );

        CREATE INDEX idx_grid_cell_state ON grid_cell(state);
        CREATE INDEX idx_grid_cell_source ON grid_cell(source_code, prefecture);
        """
    )


def main() -> None:
    parser = argparse.ArgumentParser(description="Build Fukuoka/Saga coastline hard-blocked grid")
    parser.add_argument(
        "--input",
        dest="inputs",
        action="append",
        default=[],
        help="输入海岸线 GML zip；可重复指定。默认使用 40/41 两个包。",
    )
    parser.add_argument(
        "--output-dir",
        default="out/coastline/fukuoka_saga",
        help="输出目录",
    )
    parser.add_argument("--cell-size-m", type=float, default=200.0, help="栅格尺寸，默认 200m")
    parser.add_argument("--buffer-m", type=float, default=50.0, help="海岸线缓冲距离，默认 50m")
    parser.add_argument(
        "--margin-cells",
        type=int,
        default=2,
        help="输出范围外扩的格子数，默认 2 个格子",
    )
    parser.add_argument(
        "--write-geojson",
        action="store_true",
        help="同时输出 coastlines.geojson 和 grid.geojson",
    )
    args = parser.parse_args()

    project_root = Path(__file__).resolve().parent.parent
    if args.inputs:
        input_paths = [Path(item) for item in args.inputs]
    else:
        input_paths = [
            project_root / "coastline" / "C23-06_40_GML.zip",
            project_root / "coastline" / "C23-06_41_GML.zip",
        ]

    packages: list[SourcePackage] = []
    for path in input_paths:
        if not path.exists():
            raise SystemExit(f"missing input: {path}")
        packages.append(load_package(path))

    out_dir = Path(args.output_dir)
    build_output_dir(out_dir)

    all_lines: list[LineString] = []
    line_features: list[dict] = []
    package_stats: list[dict] = []

    for package in packages:
        line_count = 0
        point_count = 0
        minx = miny = float("inf")
        maxx = maxy = float("-inf")
        for source_name, curve_id, line in iter_coastline_lines(package.path):
            coords = list(line.coords)
            point_count += len(coords)
            line_count += 1
            xs = [p[0] for p in coords]
            ys = [p[1] for p in coords]
            minx = min(minx, min(xs))
            miny = min(miny, min(ys))
            maxx = max(maxx, max(xs))
            maxy = max(maxy, max(ys))
            all_lines.append(line)
            if args.write_geojson:
                line_features.append(
                    {
                        "type": "Feature",
                        "properties": {
                            "source_code": package.code,
                            "prefecture": package.prefecture,
                            "curve_id": curve_id,
                            "source_xml": source_name,
                            "point_count": len(coords),
                        },
                        "geometry": mapping(line),
                    }
                )

        package_stats.append(
            {
                "source_code": package.code,
                "prefecture": package.prefecture,
                "zip": str(package.path),
                "line_count": line_count,
                "point_count": point_count,
                "bbox_mercator": [minx, miny, maxx, maxy],
            }
        )

    if not all_lines:
        raise SystemExit("no coastline curves were parsed")

    coastline_geom = MultiLineString([line.coords[:] for line in all_lines])
    coastline_buffer = coastline_geom.buffer(args.buffer_m, cap_style=2, join_style=2)
    prepared_buffer = prep(coastline_buffer)

    coast_bounds = coastline_geom.bounds
    minx = align_floor(coast_bounds[0] - args.buffer_m - args.margin_cells * args.cell_size_m, args.cell_size_m)
    miny = align_floor(coast_bounds[1] - args.buffer_m - args.margin_cells * args.cell_size_m, args.cell_size_m)
    maxx = align_ceil(coast_bounds[2] + args.buffer_m + args.margin_cells * args.cell_size_m, args.cell_size_m)
    maxy = align_ceil(coast_bounds[3] + args.buffer_m + args.margin_cells * args.cell_size_m, args.cell_size_m)

    sqlite_path = out_dir / "fukuoka_saga_coast_grid.sqlite"
    conn = open_sqlite(sqlite_path)
    create_schema(conn)

    build_time = dt.datetime.now().isoformat(timespec="seconds")
    source_scope = "福冈/佐贺"
    source_code = "C23-06_40+41"
    meta_rows = [
        ("build_time", build_time),
        ("cell_size_m", str(args.cell_size_m)),
        ("buffer_m", str(args.buffer_m)),
        ("margin_cells", str(args.margin_cells)),
        ("source_scope", source_scope),
        ("source_code", source_code),
        ("source_count", str(len(packages))),
        ("source_packages", json.dumps([stats["zip"] for stats in package_stats], ensure_ascii=False)),
        ("coastline_line_count", str(len(all_lines))),
        ("coastline_bounds_mercator", json.dumps(coast_bounds, ensure_ascii=False)),
        ("grid_bounds_mercator", json.dumps([minx, miny, maxx, maxy], ensure_ascii=False)),
        ("notes", "GSI coastline subset for Fukuoka/Saga, generated from C23-06_40/41"),
    ]
    conn.executemany("INSERT INTO meta(key, value) VALUES (?, ?)", meta_rows)
    conn.executemany(
        """
        INSERT INTO coastline_feature(
          source_code, prefecture, curve_id, geom_wkt, point_count, minx, miny, maxx, maxy
        ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)
        """,
        [
            (
                package.code,
                package.prefecture,
                curve_id,
                line.wkt,
                len(line.coords),
                line.bounds[0],
                line.bounds[1],
                line.bounds[2],
                line.bounds[3],
            )
            for package in packages
            for _, curve_id, line in iter_coastline_lines(package.path)
        ],
    )
    conn.commit()

    insert_rows = []
    grid_features = []
    blocked_count = 0
    candidate_count = 0
    total_cells = 0

    ix0 = int(round(minx / args.cell_size_m))
    iy0 = int(round(miny / args.cell_size_m))
    ix1 = int(round(maxx / args.cell_size_m))
    iy1 = int(round(maxy / args.cell_size_m))

    for ix in range(ix0, ix1):
        cell_minx = ix * args.cell_size_m
        cell_maxx = cell_minx + args.cell_size_m
        for iy in range(iy0, iy1):
            cell_miny = iy * args.cell_size_m
            cell_maxy = cell_miny + args.cell_size_m
            cell = box(cell_minx, cell_miny, cell_maxx, cell_maxy)
            intersects = prepared_buffer.intersects(cell)
            state = 2 if intersects else 1
            state_name = "HARD_BLOCKED" if intersects else "NAVIGABLE_CANDIDATE"
            blocked_permille = 1000 if intersects else 0
            source_mask = 1
            cell_id = f"{ix}:{iy}"
            insert_rows.append(
                (
                    cell_id,
                    source_code,
                    source_scope,
                    ix,
                    iy,
                    cell_minx,
                    cell_miny,
                    cell_maxx,
                    cell_maxy,
                    state,
                    state_name,
                    source_mask,
                    blocked_permille,
                )
            )
            total_cells += 1
            if intersects:
                blocked_count += 1
            else:
                candidate_count += 1
            if args.write_geojson:
                grid_features.append(
                    {
                        "type": "Feature",
                        "properties": {
                            "cell_id": cell_id,
                            "ix": ix,
                            "iy": iy,
                            "state": state,
                            "state_name": state_name,
                            "source_mask": source_mask,
                            "blocked_permille": blocked_permille,
                        },
                        "geometry": mapping(cell),
                    }
                )
            if len(insert_rows) >= 5000:
                conn.executemany(
                    """
                    INSERT INTO grid_cell(
                      cell_id, source_code, prefecture, ix, iy, minx, miny, maxx, maxy,
                      state, state_name, source_mask, blocked_permille
                    ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                    """,
                    insert_rows,
                )
                conn.commit()
                print(f"已写入格网 {total_cells} 个，blocked={blocked_count} candidate={candidate_count}")
                insert_rows.clear()

    if insert_rows:
        conn.executemany(
            """
            INSERT INTO grid_cell(
              cell_id, source_code, prefecture, ix, iy, minx, miny, maxx, maxy,
              state, state_name, source_mask, blocked_permille
            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
            """,
            insert_rows,
        )
        conn.commit()

    conn.execute("ANALYZE")
    conn.commit()

    grid_meta = {
        "build_time": build_time,
        "input_packages": [stats["zip"] for stats in package_stats],
        "package_stats": package_stats,
        "cell_size_m": args.cell_size_m,
        "buffer_m": args.buffer_m,
        "margin_cells": args.margin_cells,
        "grid_bounds_mercator": [minx, miny, maxx, maxy],
        "coastline_feature_count": len(all_lines),
        "grid_cell_count": total_cells,
        "blocked_count": blocked_count,
        "candidate_count": candidate_count,
    }

    meta_path = out_dir / "fukuoka_saga_coast_grid.meta.json"
    meta_path.write_text(json.dumps(grid_meta, ensure_ascii=False, indent=2), encoding="utf-8")

    if args.write_geojson:
        coastline_geojson = out_dir / "fukuoka_saga_coastline.geojson"
        write_geojson(coastline_geojson, line_features)
        grid_geojson = out_dir / "fukuoka_saga_grid.geojson"
        write_geojson(grid_geojson, grid_features)

    conn.close()

    print("完成")
    print(f"  SQLite: {sqlite_path}")
    print(f"  元数据: {meta_path}")
    if args.write_geojson:
        print(f"  海岸线GeoJSON: {out_dir / 'fukuoka_saga_coastline.geojson'}")
        print(f"  栅格GeoJSON: {out_dir / 'fukuoka_saga_grid.geojson'}")
    print(f"  coastline_features={len(all_lines)}")
    print(f"  grid_cells={total_cells}")
    print(f"  blocked={blocked_count}")
    print(f"  candidate={candidate_count}")


if __name__ == "__main__":
    main()