#!/usr/bin/env python3
from __future__ import annotations

import argparse
import datetime as dt
import json
import math
from pathlib import Path

import pymysql


DB_NAME = "navsea_japan_coast_grid"
DB_USER = "root"
DB_PASSWORD = "2chi9ks2"
DB_HOST = "localhost"
DB_SOCKET = "/tmp/mysql.sock"

DEFAULT_OUT_DIR = "src/pbf/coastline-mysql/japan_national"

DENSITY_OVERVIEW_FILE = "density_overview_grid.geojson"
CELL_SIZES = {
    "coast_200m": 200.0,
    "fish_port_20m": 20.0,
    "hazard_50m": 50.0,
}

RADIUS = 6378137.0
MAX_MERCATOR_LAT = 85.0511287798066


def mysql_connect(database: str, *, cursorclass=pymysql.cursors.Cursor):
    kwargs = {
        "host": DB_HOST,
        "user": DB_USER,
        "password": DB_PASSWORD,
        "database": database,
        "charset": "utf8mb4",
        "cursorclass": cursorclass,
    }
    if DB_SOCKET:
        kwargs["unix_socket"] = DB_SOCKET
    return pymysql.connect(**kwargs)


def rect_geometry(min_lon: float, min_lat: float, max_lon: float, max_lat: float) -> dict:
    return {
        "type": "Polygon",
        "coordinates": [
            [
                [min_lon, min_lat],
                [max_lon, min_lat],
                [max_lon, max_lat],
                [min_lon, max_lat],
                [min_lon, min_lat],
            ]
        ],
    }


def lon_from_mercator(x: float) -> float:
    return math.degrees(x / RADIUS)


def lat_from_mercator(y: float) -> float:
    return math.degrees(2.0 * math.atan(math.exp(y / RADIUS)) - math.pi / 2.0)


def mercator_bbox_to_lonlat(minx: float, miny: float, maxx: float, maxy: float) -> tuple[float, float, float, float]:
    return (
        lon_from_mercator(minx),
        lat_from_mercator(miny),
        lon_from_mercator(maxx),
        lat_from_mercator(maxy),
    )


def normalize_bbox(layer_name: str, min_lon: float, min_lat: float, max_lon: float, max_lat: float) -> tuple[float, float, float, float]:
    if layer_name == "coast_200m" and (
        abs(min_lon) > 180.0 or abs(max_lon) > 180.0 or abs(min_lat) > 90.0 or abs(max_lat) > 90.0
    ):
        return mercator_bbox_to_lonlat(min_lon, min_lat, max_lon, max_lat)
    return min_lon, min_lat, max_lon, max_lat


def fetch_stats(conn, where_sql: str, params: tuple) -> tuple[int, list[float] | None]:
    sql = f"""
        SELECT COUNT(*), MIN(min_lon), MIN(min_lat), MAX(max_lon), MAX(max_lat)
        FROM navsea_grid_cell
        WHERE {where_sql}
    """
    with conn.cursor() as cur:
        cur.execute(sql, params)
        row = cur.fetchone()
    count = int(row[0] or 0)
    if count == 0:
        return 0, None
    min_lon, min_lat, max_lon, max_lat = normalize_bbox(
        params[0] if params else "",
        float(row[1]),
        float(row[2]),
        float(row[3]),
        float(row[4]),
    )
    return count, [min_lon, min_lat, max_lon, max_lat]


def export_layer(
    conn,
    *,
    path: Path,
    where_sql: str,
    params: tuple,
) -> int:
    path.parent.mkdir(parents=True, exist_ok=True)
    sql = f"""
        SELECT layer_name, cell_id, row_idx, col_idx, cell_size_m, state_name, class_name,
               source_name, min_lon, min_lat, max_lon, max_lat
        FROM navsea_grid_cell
        WHERE {where_sql}
        ORDER BY row_idx, col_idx
    """
    count = 0
    with conn.cursor(pymysql.cursors.SSCursor) as cur, path.open("w", encoding="utf-8") as fh:
        cur.execute(sql, params)
        fh.write('{"type":"FeatureCollection","features":[\n')
        first = True
        for row in cur:
            (
                layer_name,
                cell_id,
                row_idx,
                col_idx,
                cell_size_m,
                state_name,
                class_name,
                source_name,
                min_lon,
                min_lat,
                max_lon,
                max_lat,
            ) = row
            min_lon, min_lat, max_lon, max_lat = normalize_bbox(
                row[0],
                float(min_lon),
                float(min_lat),
                float(max_lon),
                float(max_lat),
            )
            feature = {
                "type": "Feature",
                "properties": {
                    "layer_name": layer_name,
                    "cell_id": cell_id,
                    "row": int(row_idx),
                    "col": int(col_idx),
                    "cell_size_m": float(cell_size_m),
                    "state_name": state_name,
                    "class_name": class_name,
                    "source_name": source_name,
                },
                "geometry": rect_geometry(min_lon, min_lat, max_lon, max_lat),
            }
            if not first:
                fh.write(",\n")
            fh.write(json.dumps(feature, ensure_ascii=False))
            first = False
            count += 1
        fh.write("\n]}\n")
    return count


def _range_for_overlap(min_value: float, max_value: float, fine_size_m: float) -> tuple[int, int]:
    start = int(math.floor(min_value / fine_size_m))
    end = int(math.floor((max_value - 1e-9) / fine_size_m))
    return start, end


def _overlaps_finer_cells(row_idx: int, col_idx: int, coarse_size_m: float, fine_size_m: float, fine_keys: set[tuple[int, int]]) -> bool:
    min_x = col_idx * coarse_size_m
    min_y = row_idx * coarse_size_m
    max_x = min_x + coarse_size_m
    max_y = min_y + coarse_size_m
    row_start, row_end = _range_for_overlap(min_y, max_y, fine_size_m)
    col_start, col_end = _range_for_overlap(min_x, max_x, fine_size_m)
    for fine_row in range(row_start, row_end + 1):
        for fine_col in range(col_start, col_end + 1):
            if (fine_row, fine_col) in fine_keys:
                return True
    return False


def export_density_overview(conn, *, path: Path) -> tuple[int, dict]:
    path.parent.mkdir(parents=True, exist_ok=True)

    kept_counts = {
        "20": 0,
        "50": 0,
        "200": 0,
    }
    kept_bbox = [float("inf"), float("inf"), float("-inf"), float("-inf")]
    keys_20: set[tuple[int, int]] = set()
    keys_50: set[tuple[int, int]] = set()

    def write_feature(fh, feature: dict, first: bool) -> bool:
        if not first:
            fh.write(",\n")
        fh.write(json.dumps(feature, ensure_ascii=False))
        return False

    with conn.cursor(pymysql.cursors.SSCursor) as cur, path.open("w", encoding="utf-8") as fh:
        fh.write('{"type":"FeatureCollection","features":[\n')
        first = True

        # 20m: always keep
        cur.execute(
            """
            SELECT layer_name, cell_id, row_idx, col_idx, cell_size_m, state_name, class_name,
                   source_name, min_lon, min_lat, max_lon, max_lat
            FROM navsea_grid_cell
            WHERE layer_name=%s
            ORDER BY row_idx, col_idx
            """,
            ("fish_port_20m",),
        )
        for row in cur:
            (
                layer_name,
                cell_id,
                row_idx,
                col_idx,
                cell_size_m,
                state_name,
                class_name,
                source_name,
                min_lon,
                min_lat,
                max_lon,
                max_lat,
            ) = row
            keys_20.add((int(row_idx), int(col_idx)))
            kept_counts["20"] += 1
            min_lon, min_lat, max_lon, max_lat = normalize_bbox(
                layer_name,
                float(min_lon),
                float(min_lat),
                float(max_lon),
                float(max_lat),
            )
            kept_bbox[0] = min(kept_bbox[0], min_lon)
            kept_bbox[1] = min(kept_bbox[1], min_lat)
            kept_bbox[2] = max(kept_bbox[2], max_lon)
            kept_bbox[3] = max(kept_bbox[3], max_lat)
            feature = {
                "type": "Feature",
                "properties": {
                    "layer_name": layer_name,
                    "cell_id": cell_id,
                    "row": int(row_idx),
                    "col": int(col_idx),
                    "cell_size_m": float(cell_size_m),
                    "state_name": state_name,
                    "class_name": class_name,
                    "source_name": source_name,
                    "density_key": "20",
                    "density_level": 3,
                    "density_name": "高密度",
                },
                "geometry": rect_geometry(min_lon, min_lat, max_lon, max_lat),
            }
            first = write_feature(fh, feature, first)

        # 50m: keep only when no 20m cell overlaps this area
        cur.execute(
            """
            SELECT layer_name, cell_id, row_idx, col_idx, cell_size_m, state_name, class_name,
                   source_name, min_lon, min_lat, max_lon, max_lat
            FROM navsea_grid_cell
            WHERE layer_name=%s
            ORDER BY row_idx, col_idx
            """,
            ("hazard_50m",),
        )
        for row in cur:
            (
                layer_name,
                cell_id,
                row_idx,
                col_idx,
                cell_size_m,
                state_name,
                class_name,
                source_name,
                min_lon,
                min_lat,
                max_lon,
                max_lat,
            ) = row
            row_idx_i = int(row_idx)
            col_idx_i = int(col_idx)
            if _overlaps_finer_cells(row_idx_i, col_idx_i, CELL_SIZES["hazard_50m"], CELL_SIZES["fish_port_20m"], keys_20):
                continue
            keys_50.add((row_idx_i, col_idx_i))
            kept_counts["50"] += 1
            min_lon, min_lat, max_lon, max_lat = normalize_bbox(
                layer_name,
                float(min_lon),
                float(min_lat),
                float(max_lon),
                float(max_lat),
            )
            kept_bbox[0] = min(kept_bbox[0], min_lon)
            kept_bbox[1] = min(kept_bbox[1], min_lat)
            kept_bbox[2] = max(kept_bbox[2], max_lon)
            kept_bbox[3] = max(kept_bbox[3], max_lat)
            feature = {
                "type": "Feature",
                "properties": {
                    "layer_name": layer_name,
                    "cell_id": cell_id,
                    "row": row_idx_i,
                    "col": col_idx_i,
                    "cell_size_m": float(cell_size_m),
                    "state_name": state_name,
                    "class_name": class_name,
                    "source_name": source_name,
                    "density_key": "50",
                    "density_level": 2,
                    "density_name": "中密度",
                },
                "geometry": rect_geometry(min_lon, min_lat, max_lon, max_lat),
            }
            first = write_feature(fh, feature, first)

        # 200m: keep only when neither 20m nor 50m cell overlaps this area
        cur.execute(
            """
            SELECT layer_name, cell_id, row_idx, col_idx, cell_size_m, state_name, class_name,
                   source_name, min_lon, min_lat, max_lon, max_lat
            FROM navsea_grid_cell
            WHERE layer_name=%s
            ORDER BY row_idx, col_idx
            """,
            ("coast_200m",),
        )
        for row in cur:
            (
                layer_name,
                cell_id,
                row_idx,
                col_idx,
                cell_size_m,
                state_name,
                class_name,
                source_name,
                min_lon,
                min_lat,
                max_lon,
                max_lat,
            ) = row
            row_idx_i = int(row_idx)
            col_idx_i = int(col_idx)
            if _overlaps_finer_cells(row_idx_i, col_idx_i, CELL_SIZES["coast_200m"], CELL_SIZES["fish_port_20m"], keys_20):
                continue
            if _overlaps_finer_cells(row_idx_i, col_idx_i, CELL_SIZES["coast_200m"], CELL_SIZES["hazard_50m"], keys_50):
                continue
            kept_counts["200"] += 1
            min_lon, min_lat, max_lon, max_lat = normalize_bbox(
                layer_name,
                float(min_lon),
                float(min_lat),
                float(max_lon),
                float(max_lat),
            )
            kept_bbox[0] = min(kept_bbox[0], min_lon)
            kept_bbox[1] = min(kept_bbox[1], min_lat)
            kept_bbox[2] = max(kept_bbox[2], max_lon)
            kept_bbox[3] = max(kept_bbox[3], max_lat)
            feature = {
                "type": "Feature",
                "properties": {
                    "layer_name": layer_name,
                    "cell_id": cell_id,
                    "row": row_idx_i,
                    "col": col_idx_i,
                    "cell_size_m": float(cell_size_m),
                    "state_name": state_name,
                    "class_name": class_name,
                    "source_name": source_name,
                    "density_key": "200",
                    "density_level": 1,
                    "density_name": "低密度",
                },
                "geometry": rect_geometry(min_lon, min_lat, max_lon, max_lat),
            }
            first = write_feature(fh, feature, first)

        fh.write("\n]}\n")

    total = kept_counts["20"] + kept_counts["50"] + kept_counts["200"]
    if total == 0:
        raise RuntimeError("density overview export produced no features")
    summary = {
        "count": total,
        "by_density": kept_counts,
        "bbox": None if kept_bbox[0] == float("inf") else kept_bbox,
    }
    return total, summary


def main() -> None:
    parser = argparse.ArgumentParser(description="Export NavSea MySQL grid layers to static GeoJSON assets")
    parser.add_argument("--db-name", default=DB_NAME)
    parser.add_argument("--out-dir", default=DEFAULT_OUT_DIR)
    args = parser.parse_args()

    project_root = Path(__file__).resolve().parent.parent
    out_dir = Path(args.out_dir)
    if not out_dir.is_absolute():
        out_dir = project_root / out_dir
    out_dir.mkdir(parents=True, exist_ok=True)

    layers = {
        "coast_200m": {
            "file": "coast_200m_grid.geojson",
            "where": "layer_name=%s",
            "params": ("coast_200m",),
            "cell_size_m": 200.0,
        },
        "fish_port_20m": {
            "file": "fish_port_20m_grid.geojson",
            "where": "layer_name=%s",
            "params": ("fish_port_20m",),
            "cell_size_m": 20.0,
        },
        "hazard_50m": {
            "file": "hazard_50m_grid.geojson",
            "where": "layer_name=%s",
            "params": ("hazard_50m",),
            "cell_size_m": 50.0,
        },
    }

    manifest_layers: dict[str, dict] = {}
    conn = mysql_connect(args.db_name)
    try:
        for layer_name, spec in layers.items():
            out_path = out_dir / spec["file"]
            count = export_layer(conn, path=out_path, where_sql=spec["where"], params=spec["params"])
            stat_count, bbox = fetch_stats(conn, spec["where"], spec["params"])
            if count != stat_count:
                raise RuntimeError(f"{layer_name} export count mismatch: {count} != {stat_count}")
            manifest_layers[layer_name] = {
                "count": count,
                "bbox_lonlat": bbox,
                "cell_size_m": spec["cell_size_m"],
                "export_file": str(out_path.relative_to(project_root)),
            }
            print(f"{layer_name}: {count} -> {out_path}")

        density_path = out_dir / DENSITY_OVERVIEW_FILE
        density_count, density_summary = export_density_overview(conn, path=density_path)
        manifest_layers["density_overview"] = {
            "count": density_count,
            "bbox_lonlat": density_summary["bbox"],
            "cell_size_m": None,
            "export_file": str(density_path.relative_to(project_root)),
            "by_density": density_summary["by_density"],
        }
        print(f"density_overview: {density_count} -> {density_path}")
    finally:
        conn.close()

    manifest = {
        "database": args.db_name,
        "generated_at": dt.datetime.now().isoformat(timespec="seconds"),
        "layers": manifest_layers,
    }
    manifest_path = out_dir / "manifest.json"
    manifest_path.write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
    print(f"manifest: {manifest_path}")


if __name__ == "__main__":
    main()