thermograph/backend/data/places.py

"""Typo-tolerant place-name index for search suggestions.

This module keeps a local index of world places (a GeoNames cities dump,
downloaded once into data/geonames/ and reused across restarts) so /suggest can
answer instantly, offline, and tolerate a single-letter typo — a substituted,
missing, or extra letter, or two adjacent letters swapped — anywhere in the
query, including the first character.

It also builds a vocabulary of place-name tokens so a typo'd word in a
multi-word query can be respelled against words the index knows ("pest
seattle" → "west seattle").

`suggest()` still takes an injected ``upstream`` callable so the module stays
free of the fetch layer, but /suggest now passes a no-op: autocomplete is served
purely from this index and makes no per-keystroke network call. (/geocode, the
explicit search, is the path that falls back to Nominatim on a local miss for
neighbourhoods/postcodes/tiny places the cities dump lacks — see web/app.py.)

The index loads in a background thread (start_loading()). Until it's ready —
or forever, if the download fails — search()/corrections() return None/empty
and /suggest simply returns no results, so the app never needs this data to
boot or serve.
"""
import bisect
import heapq
import os
import threading
import time
import unicodedata
import zipfile

import httpx

from core import audit
import paths

GEO_DIR = os.path.join(paths.DATA_DIR, "geonames")
GEONAMES_URL = "https://download.geonames.org/export/dump/"

# Reused connection instead of a bare httpx.get per call -- mirrors
# web/app.py's reused _frontend_client / data/climate.py's _client. Only a
# handful of calls ever happen (the dump files are cached forever, see _fetch),
# but a shared client costs nothing and keeps the pattern consistent everywhere
# this module talks HTTP.
_client = httpx.Client()

# Which GeoNames cities dump to index. cities1000 (~170k places, population
# ≥ 1000) is small enough to hold in memory and big enough to cover the tiny
# vacation towns people actually search for; set THERMOGRAPH_CITIES=cities5000
# (or cities15000) to trade coverage for a lighter footprint.
CITIES = os.environ.get("THERMOGRAPH_CITIES", "cities1000")

# Entry tuple layout: (norm_name, name, admin1, country, country_code, lat, lon, pop)
_POP = 7

_load_lock = threading.Lock()
_load_started = False

# Set once, atomically, by the loader thread: (entries, names, order, vocab)
# where `entries` is population-descending (so fuzzy scans can stop at the
# first matches found), `names`/`order` are the normalized names sorted
# alphabetically with their entry indices (for prefix bisection), and `vocab`
# maps each place-name token to the population of the biggest place using it.
_data = None


def start_loading() -> None:
    """Kick off the background index load, once per process (idempotent)."""
    global _load_started
    with _load_lock:
        if _load_started:
            return
        _load_started = True
    threading.Thread(target=_load, name="places-index", daemon=True).start()


def ready() -> bool:
    return _data is not None


def _fetch(name: str) -> str:
    """Path to a GeoNames dump file, downloading into data/geonames/ if absent.
    Cached forever — cities don't move; delete the folder to force a refresh."""
    path = os.path.join(GEO_DIR, name)
    if os.path.exists(path) and os.path.getsize(path) > 0:
        return path
    r = _client.get(GEONAMES_URL + name, timeout=120, follow_redirects=True)
    r.raise_for_status()
    tmp = path + ".part"
    with open(tmp, "wb") as f:
        f.write(r.content)
    os.replace(tmp, path)   # atomic: never leave a truncated file behind
    return path


def _load() -> None:
    t0 = time.monotonic()
    try:
        os.makedirs(GEO_DIR, exist_ok=True)
        # Admin-division and country display names ("US.WA" → Washington).
        admin1 = {}
        with open(_fetch("admin1CodesASCII.txt"), encoding="utf-8") as f:
            for line in f:
                cols = line.rstrip("\n").split("\t")
                if len(cols) >= 2:
                    admin1[cols[0]] = cols[1]
        countries = {}
        with open(_fetch("countryInfo.txt"), encoding="utf-8") as f:
            for line in f:
                if line.startswith("#"):
                    continue
                cols = line.split("\t")
                if len(cols) >= 5:
                    countries[cols[0]] = cols[4]
        entries = []
        with zipfile.ZipFile(_fetch(f"{CITIES}.zip")) as z, z.open(f"{CITIES}.txt") as f:
            for raw in f:
                cols = raw.decode("utf-8").rstrip("\n").split("\t")
                if len(cols) < 15:
                    continue
                name, ascii_name, cc, a1 = cols[1], cols[2], cols[8], cols[10]
                norm = _norm(ascii_name or name)
                if len(norm) < 2:
                    continue
                try:
                    lat, lon, pop = float(cols[4]), float(cols[5]), int(cols[14] or 0)
                except ValueError:
                    continue
                entries.append((norm, name, admin1.get(f"{cc}.{a1}"),
                                countries.get(cc), cc, lat, lon, pop))
        global _data
        _data = _build(entries)
    except Exception as e:  # noqa: BLE001 - suggestions degrade to the upstream geocoder
        audit.log_event("error", {"phase": "places_load", "error": repr(e),
                                  "seconds": round(time.monotonic() - t0, 1)})


def _build(entries: list) -> tuple:
    entries.sort(key=lambda e: e[_POP], reverse=True)
    order = sorted(range(len(entries)), key=lambda i: entries[i][0])
    names = [entries[i][0] for i in order]
    vocab = {}
    for e in entries:
        for tok in e[0].split():
            if len(tok) >= 3 and vocab.get(tok, -1) < e[_POP]:
                vocab[tok] = e[_POP]
    return (entries, names, order, vocab)


def _norm(s: str) -> str:
    """Lowercased, accent-stripped, punctuation-flattened matching key, so
    "Coeur d'alene" finds Cœur d'Alene and "winston salem" Winston-Salem."""
    s = unicodedata.normalize("NFKD", s)
    s = "".join(c for c in s if not unicodedata.combining(c))
    for ch in ("'", "’", "."):
        s = s.replace(ch, "")
    for ch in ("-", ",", "/"):
        s = s.replace(ch, " ")
    return " ".join(s.casefold().split())


def _prefix_edit1(q: str, w: str) -> bool:
    """True when `q` is within one edit — a substituted, extra, or missing
    letter, or an adjacent swap — of some prefix of `w`. Prefix matching (not
    whole-name) so "chicgo" already suggests Chicago mid-typing."""
    n = len(q)
    i = 0
    m = min(n, len(w))
    while i < m and q[i] == w[i]:
        i += 1
    if i == n:
        return True                              # exact prefix, zero edits
    if w.startswith(q[i + 1:], i + 1):
        return True                              # one letter substituted
    if w.startswith(q[i + 1:], i):
        return True                              # one extra letter typed
    if w.startswith(q[i:], i + 1):
        return True                              # one letter missed
    return (i + 1 < n and i + 1 < len(w) and q[i] == w[i + 1] and q[i + 1] == w[i]
            and w.startswith(q[i + 2:], i + 2))  # adjacent letters swapped


def _within1(a: str, b: str) -> bool:
    """Whole-token Damerau-Levenshtein distance ≤ 1 (used for the vocabulary)."""
    la, lb = len(a), len(b)
    if abs(la - lb) > 1:
        return False
    i = 0
    m = min(la, lb)
    while i < m and a[i] == b[i]:
        i += 1
    if la == lb:
        if i == la:
            return True
        if a[i + 1:] == b[i + 1:]:
            return True                          # substitution
        return (i + 1 < la and a[i] == b[i + 1] and a[i + 1] == b[i]
                and a[i + 2:] == b[i + 2:])      # transposition
    s, l = (a, b) if la < lb else (b, a)
    return s[i:] == l[i + 1:]                    # insertion/deletion


def _result(e: tuple, match: str) -> dict:
    return {"name": e[1], "admin1": e[2], "country": e[3], "country_code": e[4],
            "lat": e[5], "lon": e[6], "population": e[_POP], "match": match}


def search(q: str, limit: int = 5) -> list[dict] | None:
    """Top `limit` places matching `q` as a (possibly typo'd) name prefix.

    Exact-prefix matches come first, population-descending; remaining slots are
    filled with names one edit away (only for queries of 4+ characters — with
    fewer, "one letter off" matches everything). Returns None while the index
    isn't loaded so the caller can fall back to the upstream geocoder.
    """
    data = _data
    if data is None:
        return None
    entries, names, order, _ = data
    qn = _norm(q)
    if len(qn) < 2:
        return []
    lo = bisect.bisect_left(names, qn)
    hi = bisect.bisect_left(names, qn + "\uffff", lo)
    top = heapq.nlargest(limit, range(lo, hi), key=lambda i: entries[order[i]][_POP])
    out = [_result(entries[order[i]], "prefix") for i in top]
    if len(out) < limit and len(qn) >= 4:
        need = limit - len(out)
        # A single edit can only touch one of the first two characters, so a
        # candidate's first two must overlap the query's — a cheap filter that
        # rejects ~85% of the index before the real per-name check. Scanning in
        # population order means we can stop at the first `need` hits.
        q0, q1 = qn[0], qn[1]
        for e in entries:
            w = e[0]
            if w[0] != q0 and w[0] != q1 and (len(w) < 2 or (w[1] != q0 and w[1] != q1)):
                continue
            if w.startswith(qn):
                continue                         # already counted in the prefix tier
            if _prefix_edit1(qn, w):
                out.append(_result(e, "fuzzy"))
                need -= 1
                if need == 0:
                    break
    return out


def _suggest_score(r: dict, exact: bool) -> int:
    """Prominence rank with an exactness boost — big enough that a real place
    beats a same-size typo match, small enough that a hamlet spelled exactly
    like the typo can't outrank a metropolis one edit away ("Seatle", a
    village, must not beat Seattle)."""
    pop = r.get("population") or 0
    return pop * 4 + 1 if exact else pop


def _suggest_merge(local: list, upstream: tuple, limit: int) -> list:
    """Blend local-index and upstream results into one top-`limit` list.
    Everything the user's spelling matches exactly counts as exact; only the
    local index's typo-tolerant matches take the fuzzy (unboosted) score."""
    scored = [(r, _suggest_score(r, r.get("match") != "fuzzy")) for r in local]
    scored += [(r, _suggest_score(r, True)) for r in upstream]
    scored.sort(key=lambda t: t[1], reverse=True)
    out, seen = [], set()
    for r, _ in scored:
        key = ((r.get("name") or "").casefold(), r.get("admin1") or "",
               r.get("country_code") or "")
        if key not in seen:
            seen.add(key)
            out.append(r)
            if len(out) == limit:
                break
    return out


def suggest(q: str, upstream, limit: int = 5) -> tuple[list[dict], str | None]:
    """The whole suggestion policy: the top `limit` places for a (possibly
    typo'd) query prefix, plus the respelling that produced them (or None).

    The local index answers instantly and tolerates a single-letter typo;
    ``upstream(q) -> iterable`` (an injected geocoder lookup, so this module
    stays free of the fetch layer) fills remaining slots with what the index
    doesn't know (neighbourhoods, postcodes). The upstream is consulted unless
    the local index already answered convincingly: full slots and a substantial
    place matching the spelling as typed. Fuzzy hits don't count as convincing —
    they're guesses, and when the query is an alternate name the index doesn't
    know ("münchen" prefix-matches only small towns; Munich lives upstream),
    neither those towns nor a coincidental fuzzy big-city hit may suppress the
    real answer. When both come up empty, one query token is respelled against
    known place-name tokens and retried ("pest seattle" → "west seattle").

    An upstream failure degrades to whatever the index found — it propagates
    only when there is nothing at all to serve (index still loading, no
    correction hits)."""
    q = q.strip()
    if len(q) < 2:
        return [], None
    local = search(q, limit)     # None while the index loads
    results = list(local or [])
    upstream_error = None
    prominent = max((r.get("population") or 0 for r in results
                     if r.get("match") == "prefix"), default=0)
    if len(results) < limit or prominent < 100_000:
        try:
            results = _suggest_merge(results, tuple(upstream(q)), limit)
        except Exception as e:  # noqa: BLE001 - local results (if any) still serve
            upstream_error = e
    corrected = None
    if not results and len(q) >= 4:
        for phrase in corrections(q):
            hits = search(phrase, limit) or []
            if not hits:
                try:
                    hits = list(upstream(phrase))
                except Exception:  # noqa: BLE001 - a failed probe just means no hits
                    hits = []
            if hits:
                results, corrected = hits[:limit], phrase
                break
    if not results and local is None and upstream_error is not None:
        raise upstream_error
    return results[:limit], corrected


def corrections(q: str, max_phrases: int = 3) -> list[str]:
    """Candidate respellings of `q` with one token replaced by a known
    place-name token a single edit away — "pest seattle" → "west seattle".

    Ranked by how likely the swap is: tokens the vocabulary has never seen get
    corrected first, and replacements are ordered by the population of the
    biggest place using them ("west" over "wesh"). Callers verify candidates by
    actually searching, so a wrong guess only costs one lookup.
    """
    data = _data
    if data is None:
        return []
    vocab = data[3]
    toks = _norm(q).split()
    cands: dict[str, tuple] = {}
    for i, t in enumerate(toks):
        if len(t) < 3:
            continue                             # respelling 1-2 letter tokens is noise
        unknown = t not in vocab
        t0, t1 = t[0], t[1]
        for v, vpop in vocab.items():
            if v[0] != t0 and v[0] != t1 and (len(v) < 2 or (v[1] != t0 and v[1] != t1)):
                continue                         # same first-two-chars filter as search()
            if abs(len(v) - len(t)) > 1 or v == t or not _within1(t, v):
                continue
            phrase = " ".join(toks[:i] + [v] + toks[i + 1:])
            score = (unknown, vpop)
            if cands.get(phrase, (False, -1)) < score:
                cands[phrase] = score
    ranked = sorted(cands, key=cands.get, reverse=True)
    return ranked[:max_phrases]