Invalidating Tile Caches After a Source Update

from __future__ import annotations

import math
from dataclasses import dataclass
from typing import Iterator

from pyproj import Transformer
from shapely.geometry import box
from shapely.ops import transform as shapely_transform

# Web Mercator extent, the grid nearly every tile pyramid uses.
WORLD = 20037508.342789244
TILE_PX = 256

# How far a resampling kernel reaches beyond the pixel it is producing, in
# source pixels. Bilinear touches 2x2, cubic 4x4, lanczos 6x6 — so a tile whose
# edge sits one pixel outside the changed area still changes.
KERNEL_REACH = {"nearest": 0, "bilinear": 1, "cubic": 2, "cubicspline": 2, "lanczos": 3}


@dataclass(frozen=True)
class Tile:
    z: int
    x: int
    y: int

    @property
    def key(self) -> str:
        return f"tiles/{self.z}/{self.x}/{self.y}.png"


def resolution_at(z: int) -> float:
    """Metres per pixel at this zoom, in Web Mercator."""
    return (2 * WORLD) / (TILE_PX * (2 ** z))


def tiles_for_extent(minx: float, miny: float, maxx: float, maxy: float, z: int) -> Iterator[Tile]:
    """Every tile at zoom z whose footprint intersects the extent (EPSG:3857)."""
    span = (2 * WORLD) / (2 ** z)
    x0 = max(0, int((minx + WORLD) // span))
    x1 = min(2 ** z - 1, int((maxx + WORLD) // span))
    # Tile y counts downward from the top of the world in the XYZ scheme.
    y0 = max(0, int((WORLD - maxy) // span))
    y1 = min(2 ** z - 1, int((WORLD - miny) // span))
    for x in range(x0, x1 + 1):
        for y in range(y0, y1 + 1):
            yield Tile(z, x, y)


def affected_tiles(
    changed_wkt: str, source_crs: str, min_zoom: int = 0, max_zoom: int = 16,
    resampling: str = "bilinear",
) -> set[Tile]:
    """Every tile that must be rebuilt after `changed_wkt` was republished."""
    from shapely import from_wkt

    geom = from_wkt(changed_wkt)
    if source_crs != "EPSG:3857":
        transformer = Transformer.from_crs(source_crs, "EPSG:3857", always_xy=True)
        geom = shapely_transform(transformer.transform, geom)

    reach = KERNEL_REACH.get(resampling, 1)
    out: set[Tile] = set()
    for z in range(min_zoom, max_zoom + 1):
        # Buffer in METRES at this zoom: the kernel's reach is in pixels, and a
        # pixel is a different size at every level. Forgetting this leaves seams
        # at low zooms while looking correct at high ones.
        buffered = geom.buffer(reach * resolution_at(z))
        out.update(tiles_for_extent(*buffered.bounds, z=z))
    return out


def invalidate(s3, bucket: str, tiles: set[Tile], cdn=None, distribution_id: str = "") -> int:
    """Delete from origin, then ask the CDN to forget. Order matters."""
    keys = [{"Key": t.key} for t in sorted(tiles, key=lambda t: (t.z, t.x, t.y))]
    for batch in (keys[i:i + 1000] for i in range(0, len(keys), 1000)):
        s3.delete_objects(Bucket=bucket, Delete={"Objects": batch, "Quiet": True})

    if cdn is not None and keys:
        # Path-based invalidation is metered and slow; a small number of wildcard
        # paths beats ten thousand exact ones when the set is contiguous.
        paths = _collapse_to_wildcards([k["Key"] for k in keys])
        cdn.create_invalidation(
            DistributionId=distribution_id,
            InvalidationBatch={
                "Paths": {"Quantity": len(paths), "Items": paths},
                "CallerReference": f"tiles-{hash(frozenset(tiles)) & 0xffffffff}",
            },
        )
    return len(keys)

def test_buffer_adds_a_ring() -> None:
    footprint = box(1_000_000, 7_000_000, 1_010_000, 7_010_000).wkt
    plain = affected_tiles(footprint, "EPSG:3857", 12, 12, resampling="nearest")
    buffered = affected_tiles(footprint, "EPSG:3857", 12, 12, resampling="lanczos")
    assert buffered > plain, "lanczos must invalidate more than nearest"


def test_tile_y_is_not_flipped() -> None:
    """A northern extent must map to a LOW y in the XYZ scheme."""
    north = box(0, 8_000_000, 10_000, 8_010_000).wkt
    south = box(0, -8_010_000, 10_000, -8_000_000).wkt
    ny = min(t.y for t in affected_tiles(north, "EPSG:3857", 8, 8))
    sy = min(t.y for t in affected_tiles(south, "EPSG:3857", 8, 8))
    assert ny < sy, "y axis inverted — northern tiles must have smaller y"


def test_every_zoom_is_covered() -> None:
    footprint = box(1_000_000, 7_000_000, 1_001_000, 7_001_000).wkt
    tiles = affected_tiles(footprint, "EPSG:3857", 0, 14)
    assert {t.z for t in tiles} == set(range(15))

by_zoom = Counter(t.z for t in tiles)
logger.info("invalidating %d tiles: %s", len(tiles), dict(sorted(by_zoom.items())))
if by_zoom[3] > 16:
    raise ValueError("low-zoom count implausible — check the source CRS")