"""Load greenspace/water polygons and subtract them from postcode boundaries."""

from pathlib import Path

import polars as pl
from shapely import make_valid, wkb
from shapely.geometry import MultiPolygon, Polygon
from shapely.strtree import STRtree

from .geometry import _SNAP_GRID, _poly_valid, safe_difference, safe_intersection, safe_union


def load_greenspace(path: Path) -> tuple[STRtree, list]:
    """Load greenspace parquet and build an STRtree spatial index.

    Geometries are repaired with ``make_valid`` on load: an invalid park/lake
    polygon would make the per-postcode ``intersects`` predicate (and the exact
    difference path) liable to raise mid-merge, hours into a build. Empty
    geometries are dropped.

    Returns:
        (tree, geoms) where tree is a Shapely STRtree and geoms is
        the list of geometries indexed by the tree.
    """
    df = pl.read_parquet(path)
    geoms = []
    for raw in df["geometry"].to_list():
        geom = wkb.loads(raw)
        if not geom.is_valid:
            geom = make_valid(geom)
        if not geom.is_empty:
            geoms.append(geom)
    tree = STRtree(geoms)
    return tree, geoms


MAX_REMOVAL_FRACTION = 0.9  # Keep original if >90% would be removed

# Greenspace that merely trims the edge of a postcode is fine, but greenspace
# that CROSSES it (a river, a strip of parkland, a golf course running through a
# village) splits the postcode into a MultiPolygon -- one the map then draws as
# several disconnected pieces. When subtraction disconnects a postcode, re-add
# the postcode's OWN removed land along the narrowest necks (a morphological
# closing clipped to the original footprint) so the parts stay joined by a thin
# bridge. Parts left more than ~2x this width apart (a genuinely wide barrier)
# stay split. Because the bridge is the postcode's own land, no address moves and
# only a thin sliver of green is kept back.
_RECONNECT_BRIDGE_M = 25.0


def _reconnect_split(
    result: Polygon | MultiPolygon, postcode_geom: Polygon | MultiPolygon
) -> Polygon | MultiPolygon:
    """Re-join postcode parts that greenspace subtraction pulled apart by re-adding
    the narrow removed necks (within the original postcode), leaving wide barriers
    intact."""
    if result.geom_type != "MultiPolygon":
        return result
    closed = result.buffer(_RECONNECT_BRIDGE_M).buffer(-_RECONNECT_BRIDGE_M)
    if not closed.is_valid:
        closed = make_valid(closed)
    # The closing material that lies inside the original postcode but outside the
    # subtraction result == the thin green necks linking the parts. The exact
    # overlay path can leave line/point debris (coincident edges) that is
    # zero-area but NOT is_empty; `_poly_valid` strips it to polygons only, so the
    # is_empty guard works and the union can never return a GeometryCollection
    # (which `to_wgs84_geojson_multi` would silently truncate to a single part).
    bridges = _poly_valid(
        safe_difference(safe_intersection(closed, postcode_geom), result), _SNAP_GRID
    )
    if bridges.is_empty:
        return result
    return _poly_valid(safe_union([result, bridges]), _SNAP_GRID)


def subtract_greenspace(
    postcode_geom: Polygon | MultiPolygon,
    tree: STRtree,
    geoms: list,
) -> Polygon | MultiPolygon:
    """Subtract park/water polygons that overlap the postcode geometry.

    Uses the STRtree for fast candidate lookup, then subtracts the union
    of intersecting greenspace from the postcode polygon. If subtraction
    would remove >90% of the area, keeps the original (the postcode
    genuinely covers that land, e.g. churchyards, riverside addresses).

    If the subtraction disconnects the postcode (greenspace crossing it),
    :func:`_reconnect_split` re-adds the narrowest removed necks so the postcode
    stays a single piece rather than shipping as scattered fragments.
    """
    candidate_idxs = tree.query(postcode_geom)
    if len(candidate_idxs) == 0:
        return postcode_geom

    # Collect geometries that actually intersect (not just bbox overlap)
    intersecting = []
    for idx in candidate_idxs:
        g = geoms[idx]
        if g.intersects(postcode_geom):
            intersecting.append(g)

    if not intersecting:
        return postcode_geom

    green_union = safe_union(intersecting)
    result = safe_difference(postcode_geom, green_union)

    if result.is_empty:
        return postcode_geom

    # Don't over-trim postcodes that genuinely cover green/water areas
    original_area = postcode_geom.area
    if original_area > 0 and result.area / original_area < (1 - MAX_REMOVAL_FRACTION):
        return postcode_geom

    return _reconnect_split(result, postcode_geom)