"""GeoJSON exports: rivers, coastlines/lakes (marching squares), plate boundaries."""
from __future__ import annotations

import json
from pathlib import Path

import h3.api.basic_int as h3
import numpy as np

BOUNDARY_NAMES = {1: "convergent", 2: "divergent", 3: "transform"}
# marching-squares cases; corner bits TL=8 TR=4 BR=2 BL=1; edges T R B L
_CASES = {1: [("L", "B")], 2: [("B", "R")], 3: [("L", "R")], 4: [("T", "R")], 5: [("T", "R"), ("L", "B")],
          6: [("T", "B")], 7: [("L", "T")], 8: [("L", "T")], 9: [("T", "B")], 10: [("L", "T"), ("B", "R")],
          11: [("T", "R")], 12: [("L", "R")], 13: [("B", "R")], 14: [("L", "B")]}
_EDGE = {"T": (1, 0), "R": (2, 1), "B": (1, 2), "L": (0, 1)}   # doubled (dx, dy) inside a 2x2 block


def split_antimeridian(coords):
    parts, cur = [], [list(coords[0])]
    for a, b in zip(coords, coords[1:]):
        if abs(b[0] - a[0]) > 180.0:
            parts.append(cur)
            cur = []
        cur.append(list(b))
    parts.append(cur)
    return [p for p in parts if len(p) >= 2]


def _feature(geom_type, coords, props):
    return {"type": "Feature", "geometry": {"type": geom_type, "coordinates": coords}, "properties": props}


def river_lines(g, recv, river, strahler, discharge):
    """One LineString per run of equal Strahler order (runs share their junction vertex)."""
    n = g.n
    ar = np.arange(n)
    has_donor = np.zeros(n, bool)
    nr = river & (recv != ar)
    has_donor[recv[nr]] = True
    visited = np.zeros(n, bool)
    feats = []
    for s in np.flatnonzero(river & ~has_donor):
        cells, c = [int(s)], int(s)
        while True:
            visited[c] = True
            r = int(recv[c])
            if r == c:
                break
            cells.append(r)
            if not river[r] or visited[r]:
                break
            c = r
        body = cells[:-1] if len(cells) > 1 else cells
        start = 0
        for i in range(1, len(body) + 1):
            if i == len(body) or strahler[body[i]] != strahler[body[i - 1]]:
                run = body[start:i] + [cells[i] if i < len(cells) else body[-1]]
                coords = [[round(float(g.lon[k]), 4), round(float(g.lat[k]), 4)] for k in run]
                props = {"order": int(strahler[body[start]]), "discharge_km3_yr": round(float(discharge[body[i - 1]]), 3)}
                feats += [_feature("LineString", p, props) for p in split_antimeridian(coords)]
                start = i
    return feats


def boundary_features(g, plate, bnd_type):
    s, d = g.src, g.dst
    e = np.flatnonzero((plate[s] != plate[d]) & (s < d))
    lines = {}
    for k in e:
        a, b = int(g.ids[s[k]]), int(g.ids[d[k]])
        seg = [[round(lng, 4), round(lat, 4)] for lat, lng in h3.directed_edge_to_boundary(h3.cells_to_directed_edge(a, b))]
        if abs(seg[0][0] - seg[-1][0]) > 180:
            continue
        lines.setdefault(BOUNDARY_NAMES.get(int(bnd_type[s[k]]), "transform"), []).append(seg)
    return [_feature("MultiLineString", v, {"type": t}) for t, v in sorted(lines.items())]


def contours(mask):
    m = np.pad(np.asarray(mask, dtype=np.uint8), 1)
    code = m[:-1, :-1] * 8 + m[:-1, 1:] * 4 + m[1:, 1:] * 2 + m[1:, :-1]
    adj = {}
    for case, segs in _CASES.items():
        ys, xs = np.nonzero(code == case)
        for (e1, e2) in segs:
            for y, x in zip(ys.tolist(), xs.tolist()):
                p = (2 * x + _EDGE[e1][0], 2 * y + _EDGE[e1][1])
                q = (2 * x + _EDGE[e2][0], 2 * y + _EDGE[e2][1])
                adj.setdefault(p, []).append(q)
                adj.setdefault(q, []).append(p)
    seen, rings = set(), []
    for start in adj:
        if start in seen:
            continue
        ring, prev, cur = [start], None, start
        seen.add(start)
        while True:
            nxt = [q for q in adj[cur] if q != prev]
            if not nxt:
                break
            prev, cur = cur, nxt[0]
            ring.append(cur)
            if cur == start:
                break
            seen.add(cur)
        rings.append([(px / 2.0 - 1.0, py / 2.0 - 1.0) for px, py in ring])
    return rings


def _signed_area(ring):
    a = np.asarray(ring, dtype=np.float64)
    return 0.5 * float(np.sum(a[:-1, 0] * a[1:, 1] - a[1:, 0] * a[:-1, 1]))


def _point_in_ring(pt, ring):
    a = np.asarray(ring, dtype=np.float64)
    x1, y1, x2, y2 = a[:-1, 0], a[:-1, 1], a[1:, 0], a[1:, 1]
    cross = (y1 > pt[1]) != (y2 > pt[1])
    xi = x1 + (pt[1] - y1) * (x2 - x1) / np.where(y2 != y1, y2 - y1, 1e-300)
    return bool(np.count_nonzero(cross & (pt[0] < xi)) % 2)


def _orient(ring, ccw):
    return ring if (_signed_area(ring) > 0) == ccw else ring[::-1]


def _clip(ring, left, x0=180.0):
    """Sutherland–Hodgman clip of a closed ring to x ≤ x0 (left) or x ≥ x0."""
    inside = (lambda p: p[0] <= x0) if left else (lambda p: p[0] >= x0)
    cut = lambda p, q: [x0, p[1] + (x0 - p[0]) / (q[0] - p[0]) * (q[1] - p[1])]
    pts, out = ring[:-1], []
    for i in range(len(pts)):
        cur, prev = pts[i], pts[i - 1]
        if inside(cur):
            if not inside(prev):
                out.append(cut(prev, cur))
            out.append(cur)
        elif inside(prev):
            out.append(cut(prev, cur))
    return out + [out[0]] if len(out) >= 3 else []


def _split_seam(rings):
    """rings[0] exterior + holes in continuous longitude; split at +180 into ≤2 polygons in [−180, 180]."""
    if max(p[0] for p in rings[0]) <= 180.0:
        return [rings]
    polys = []
    for left in (True, False):
        ext = _clip(rings[0], left)
        if len(ext) < 4:
            continue
        holes = [h for h in (_clip(r, left) for r in rings[1:]) if len(h) >= 4]
        part = [ext] + holes
        if not left:
            part = [[[x - 360.0, y] for x, y in r] for r in part]
        polys.append(part)
    return polys


def _round(poly):
    return [[[round(x, 4), round(y, 4)] for x, y in r] for r in poly]


def contour_features(mask, kind):
    """Land/lake outlines as RFC 7946 (Multi)Polygons: CCW exteriors, CW holes, split at the antimeridian."""
    H, W = mask.shape
    col = int(np.argmin(mask.sum(axis=0)))
    rings = []
    for ring in contours(np.roll(mask, -col, axis=1)):
        pts = [[(x + col + 0.5) / W * 360.0 - 180.0, max(-90.0, min(90.0, 90.0 - (y + 0.5) / H * 180.0))]
               for x, y in ring]
        if len(pts) >= 4 and pts[0] == pts[-1]:
            rings.append(pts)
    depth = [sum(_point_in_ring(r[0], o) for j, o in enumerate(rings) if j != i) for i, r in enumerate(rings)]
    feats = []
    for i, r in enumerate(rings):
        if depth[i] % 2:
            continue
        holes = [_orient(rings[j], False) for j in range(len(rings))
                 if depth[j] == depth[i] + 1 and _point_in_ring(rings[j][0], r)]
        polys = [_round(p) for p in _split_seam([_orient(r, True)] + holes)]
        if len(polys) == 1:
            feats.append(_feature("Polygon", polys[0], {"kind": kind}))
        elif polys:
            feats.append(_feature("MultiPolygon", polys, {"kind": kind}))
    return feats


def _write(path: Path, feats) -> None:
    path.write_text(json.dumps({"type": "FeatureCollection", "features": feats}, separators=(",", ":")))


def write_all(ctx, out_dir: Path, land_raster, lake_raster) -> None:
    g, d = ctx.grid, ctx.data
    gdir = out_dir / "geo"
    gdir.mkdir(parents=True, exist_ok=True)
    _write(gdir / "rivers.geojson", river_lines(g, d["recv"], d["river"], d["strahler"], d["discharge_km3_yr"]))
    _write(gdir / "plate_boundaries.geojson", boundary_features(g, d["plate"], d["bnd_type"]))
    step = max(1, land_raster.shape[1] // 2048)
    _write(gdir / "coast.geojson", contour_features(land_raster[::step, ::step], "land"))
    _write(gdir / "lakes.geojson", contour_features(lake_raster[::step, ::step] & land_raster[::step, ::step], "lake"))