-- The durable cache. This IS the system of record.
CREATE TABLE geocode_cache (
  normalized_address text PRIMARY KEY,
  lat                double precision,
  lng                double precision,
  access_lat         double precision,   -- road-network entry point
  access_lng         double precision,
  label              text,
  query_score        real,               -- 0..1
  field_score        jsonb,              -- per-component confidence
  house_number_type  text,               -- 'PA' | 'interpolated'
  here_id            text,
  map_release        text,
  geocoded_at        timestamptz NOT NULL DEFAULT now()
);

-- Idempotency. Survives worker restarts.
CREATE TABLE batch_jobs (
  chunk_id     int PRIMARY KEY,
  job_id       text UNIQUE,          -- persisted BEFORE the first poll
  status       text NOT NULL,        -- pending|submitted|running|succeeded|failed
  submitted_at timestamptz,
  row_ids      bigint[] NOT NULL     -- deterministic membership
);

CREATE TABLE exception_queue (
  raw_address text PRIMARY KEY,
  reason      text,
  query_score real,
  queued_at   timestamptz DEFAULT now()
);

-- Normalize at WRITE time, then deduplicate. Order matters:
-- deduplicating before normalization catches only exact string matches.

WITH normalized AS (
  SELECT
    id,
    trim(regexp_replace(
      lower(
        regexp_replace(raw_address, '\s+', ' ', 'g')
      ),
      '\m(st|ave|blvd|rd|apt)\M',
      CASE                                  -- expand abbreviations
        WHEN true THEN 'street' END,
      'g'
    )) AS norm
  FROM source_addresses
)
INSERT INTO work_queue (normalized_address, source_ids)
SELECT norm, array_agg(id)
FROM normalized
WHERE NOT EXISTS (                          -- skip what we already have
  SELECT 1 FROM geocode_cache c WHERE c.normalized_address = normalized.norm
)
GROUP BY norm;                              -- ← the deduplication

-- Report this BEFORE promising a timeline.
SELECT count(*) AS billable_addresses FROM work_queue;

import os, time, requests, psycopg2

BASE = "https://batch.search.hereapi.com/v7/batch/jobs"
API_KEY = os.environ["HERE_API_KEY"]
SERVICE_GEOCODE = "hrn:here:service::olp-here:search-geocode-7"

CHUNK_SIZE = 50_000
MAX_CONCURRENT_JOBS = 3      # read HERE's limits page. It sets your parallelism.


class EntitlementError(Exception): ...


def make_chunks(conn):
    """
    DETERMINISTIC. Stable ordering + fixed size means a restarted process
    reconstructs the same chunks. A random shuffle does not.
    """
    with conn.cursor() as cur:
        cur.execute("""
            INSERT INTO batch_jobs (chunk_id, status, row_ids)
            SELECT
              (row_number() OVER (ORDER BY normalized_address) - 1) / %s AS chunk_id,
              'pending',
              array_agg(id ORDER BY id)
            FROM work_queue
            GROUP BY chunk_id
            ON CONFLICT (chunk_id) DO NOTHING
        """, (CHUNK_SIZE,))
        conn.commit()


def submit_next(conn) -> str | None:
    """
    Respects the concurrency limit. 429 on start means QUEUE — not retry.
    Firing retries in a tight loop achieves nothing except log volume.
    """
    with conn.cursor() as cur:
        cur.execute("SELECT count(*) FROM batch_jobs WHERE status IN ('submitted','running')")
        if cur.fetchone()[0] >= MAX_CONCURRENT_JOBS:
            return None

        cur.execute("SELECT chunk_id, row_ids FROM batch_jobs "
                    "WHERE status='pending' ORDER BY chunk_id LIMIT 1 "
                    "FOR UPDATE SKIP LOCKED")
        row = cur.fetchone()
        if not row:
            return None
        chunk_id, row_ids = row

    payload = build_payload(conn, row_ids)

    r = requests.post(BASE, params={
        "serviceHrn": SERVICE_GEOCODE,
        "inputDelimiter": "|", "outputDelimiter": "|",
        "outputColumns": "|".join([
            "position", "title", "id",
            "addressLabel", "addressPostalCode", "addressCountryCode",
        ]),
        "outputType": "csv",
        "startJob": "true",
        "apiKey": API_KEY,
    }, data=payload.encode("utf-8"),
       headers={"Content-Type": "text/plain", "Accept": "application/json"},
       timeout=180)

    if r.status_code == 403:
        raise EntitlementError("Batch not entitled")
    if r.status_code == 429:
        return None                     # concurrency limit. Queue. Do not hammer.
    r.raise_for_status()

    job_id = r.json()["id"]

    # ── PERSIST BEFORE POLLING. ────────────────────────────────────────────
    # If the worker dies here, the job is RUNNING on HERE's side regardless.
    # Resubmitting a 50k-record chunk bills 50k records again.
    with conn.cursor() as cur:
        cur.execute("UPDATE batch_jobs SET job_id=%s, status='submitted', "
                    "submitted_at=now() WHERE chunk_id=%s", (job_id, chunk_id))
        conn.commit()

    return job_id


def build_payload(conn, row_ids) -> str:
    with conn.cursor() as cur:
        cur.execute("SELECT id, normalized_address, country FROM work_queue "
                    "WHERE id = ANY(%s) ORDER BY id", (row_ids,))
        lines = ["recId|q|country"]
        lines += [f"{r[0]}|{r[1]}|{r[2]}" for r in cur.fetchall()]
    return "\n".join(lines)

def poll_all(conn):
    """Exponential backoff. A 40-minute job does not need checking every 2s."""
    with conn.cursor() as cur:
        cur.execute("SELECT chunk_id, job_id FROM batch_jobs "
                    "WHERE status IN ('submitted','running')")
        active = cur.fetchall()

    for chunk_id, job_id in active:
        r = requests.get(f"{BASE}/{job_id}", params={"apiKey": API_KEY}, timeout=30)
        r.raise_for_status()
        status = r.json()["status"]

        if status == "success":
            _collect(conn, chunk_id, job_id)
        elif status in ("failed", "cancelled"):
            _mark(conn, chunk_id, status)
        else:
            _mark(conn, chunk_id, "running")


def _collect(conn, chunk_id: int, job_id: str):
    # ── /errors: HTTP 204 means ZERO ERRORS. Not a failure. ───────────────
    e = requests.get(f"{BASE}/{job_id}/errors", params={"apiKey": API_KEY}, timeout=60)
    if e.status_code == 204:
        errors = None                   # celebrate
    else:
        e.raise_for_status()
        errors = e.text
        _queue_exceptions(conn, errors) # partial success. Do NOT discard the rest.

    # ── /result: HTTP 404 means NOT READY. Poll status first. ─────────────
    path = f"/tmp/{job_id}.csv"
    res = requests.get(f"{BASE}/{job_id}/result", params={"apiKey": API_KEY},
                       stream=True, timeout=600)
    if res.status_code == 404:
        raise RuntimeError("Results not ready — status was success?")
    res.raise_for_status()

    # Stream. One million records do not fit in memory.
    with open(path, "wb") as f:
        for chunk in res.iter_content(chunk_size=1 << 20):
            f.write(chunk)

    _merge_into_cache(conn, path)
    requests.delete(f"{BASE}/{job_id}", params={"apiKey": API_KEY}, timeout=30)
    _mark(conn, chunk_id, "succeeded")

MIN_QUERY_SCORE = 0.90     # calibrate against YOUR ground truth. Do not port.


def _merge_into_cache(conn, csv_path: str):
    """
    Successes merge. Sub-threshold results go to a human.

    A pipeline that treats geocoding as binary — succeeded or failed —
    will persist low-confidence fallbacks as truth. That corruption is
    silent, permanent, and propagates into every routing decision and
    every analytical aggregate downstream.
    """
    good, review = [], []
    for row in stream_csv(csv_path):
        if row["query_score"] >= MIN_QUERY_SCORE:
            good.append(row)
        else:
            review.append(row)

    with conn.cursor() as cur:
        cur.executemany("""
            INSERT INTO geocode_cache
              (normalized_address, lat, lng, access_lat, access_lng, label,
               query_score, field_score, house_number_type, here_id, geocoded_at)
            VALUES (%(norm)s,%(lat)s,%(lng)s,%(alat)s,%(alng)s,%(label)s,
                    %(query_score)s,%(field_score)s,%(hn_type)s,%(here_id)s,now())
            ON CONFLICT (normalized_address) DO UPDATE
              SET lat=EXCLUDED.lat, lng=EXCLUDED.lng,
                  query_score=EXCLUDED.query_score, geocoded_at=now()
        """, good)

        cur.executemany("""
            INSERT INTO exception_queue (raw_address, reason, query_score)
            VALUES (%(norm)s, 'below_confidence_threshold', %(query_score)s)
            ON CONFLICT (raw_address) DO NOTHING
        """, review)
        conn.commit()

-- Backfill the source table from the cache.
UPDATE source_addresses s
SET lat = c.lat, lng = c.lng,
    access_lat = c.access_lat, access_lng = c.access_lng,
    geocoded_at = c.geocoded_at
FROM geocode_cache c, work_queue w
WHERE w.normalized_address = c.normalized_address
  AND s.id = ANY(w.source_ids);