tuning/cache.py

"""Disk caches so the tuning suite stops hammering paid/rate-limited APIs.

- **Google (and NAIP) tiles**: `CachedImagery` implements the `DisplayImagery`
  contract (`sources.base.DisplayImagery.fetch_satellite_image`) and stores the exact
  returned PNG keyed by the five parameters that determine the pixels
  `(lat, lon, zoom, image_size, scale)`. On a miss it delegates to the real client.
  Because it stores the raw image, a cached tile is byte-identical to a live fetch —
  so running the pipeline over the cache preserves the estimate exactly.
- **Geocode + region**: `resolve_region` does several networked ArcGIS/Google calls
  per address. We memoize just the point (lat/lon/formatted address) to JSON so the
  label tool and repeat tuning runs skip it. (The full region object isn't cached —
  it's cheap to rebuild and holds live sessions.)

LiDAR is already disk-cached by `sources.lidar._stream_download`, so nothing to add.
"""

from __future__ import annotations

import json
import threading
from pathlib import Path

import requests
from PIL import Image

from lawn_estimator.config import DATA_DIR, Config
from lawn_estimator.sources.google import GoogleStaticMaps
from lawn_estimator.sources.imagery import NaipImagery

TUNING_DIR = DATA_DIR / "tuning"
TILE_CACHE_DIR = TUNING_DIR / "tile_cache"
GEOCODE_CACHE_PATH = TUNING_DIR / "geocode_cache.json"


def _build_client(imagery: str, config: Config, session: requests.Session):
    """The real imagery client the production pipeline would use for this source."""
    if imagery == "naip":
        return NaipImagery(session)
    if imagery == "google":
        if not config.google_maps_api_key:
            raise RuntimeError("imagery='google' needs GOOGLE_MAPS_API_KEY in your .env.")
        return GoogleStaticMaps(config.google_maps_api_key, session)
    raise ValueError(f"cache supports imagery in {{google, naip}}, got {imagery!r}")


class CachedImagery:
    """A `DisplayImagery` that memoizes tiles to disk. Wraps a real client; the on-disk
    PNG is the exact bytes the client returned, so cached == live for the pipeline."""

    def __init__(self, imagery: str, config: Config, session: requests.Session,
                 cache_dir: Path = TILE_CACHE_DIR, client=None):
        # `client` lets a caller wrap an arbitrary DisplayImagery (e.g. a region's
        # county leaf-off ortho) that _build_client doesn't know how to construct.
        self.imagery = imagery
        self._client = client if client is not None else _build_client(imagery, config, session)
        self._dir = cache_dir / imagery
        self._dir.mkdir(parents=True, exist_ok=True)
        self.hits = 0
        self.misses = 0

    def _key(self, latitude: float, longitude: float, zoom: int, image_size: int, scale: int) -> Path:
        return self._dir / f"{latitude:.6f}_{longitude:.6f}_z{zoom}_s{image_size}x{scale}.png"

    def fetch_satellite_image(
        self, latitude: float, longitude: float, zoom: int, image_size: int, scale: int
    ) -> Image.Image:
        path = self._key(latitude, longitude, zoom, image_size, scale)
        if path.exists():
            self.hits += 1
            return Image.open(path).convert("RGB")
        self.misses += 1
        image = self._client.fetch_satellite_image(
            latitude=latitude, longitude=longitude, zoom=zoom, image_size=image_size, scale=scale
        )
        tmp = path.with_suffix(".png.part")  # atomic write: never leave a half PNG
        image.save(tmp, "PNG")
        tmp.replace(path)
        return image


# ---------------------------------------------------------------------------
# Geocode memo — a tiny JSON keyed by the raw address string.
# ---------------------------------------------------------------------------
_GEO_LOCK = threading.Lock()


def _load_geocode_cache() -> dict:
    if GEOCODE_CACHE_PATH.exists():
        try:
            return json.loads(GEOCODE_CACHE_PATH.read_text(encoding="utf-8"))
        except (json.JSONDecodeError, OSError):
            return {}
    return {}


def cached_geocode(address: str) -> dict | None:
    """Return `{latitude, longitude, formatted_address, region}` if seen before."""
    return _load_geocode_cache().get(address)


def store_geocode(address: str, latitude: float, longitude: float,
                  formatted_address: str, region: str) -> None:
    with _GEO_LOCK:
        cache = _load_geocode_cache()
        cache[address] = {
            "latitude": latitude, "longitude": longitude,
            "formatted_address": formatted_address, "region": region,
        }
        GEOCODE_CACHE_PATH.parent.mkdir(parents=True, exist_ok=True)
        tmp = GEOCODE_CACHE_PATH.with_suffix(".json.part")
        tmp.write_text(json.dumps(cache, indent=2), encoding="utf-8")
        tmp.replace(GEOCODE_CACHE_PATH)