"""Forward-geocode houses through Yandex Geocoder API to backfill lat/lon
and canonical address, plus optional reverse audit of already-geocoded houses.

Phase 2-3 of Forgejo issue #582. Two modes (mutually exclusive):

  1. Backfill (default) — for the ~4141 rows WHERE lat IS NULL OR lon IS NULL:
     forward-geocode `houses.address` → snap to a Yandex `house`-precision
     point, UPDATE houses with the new lat/lon + canonical address payload,
     and write an `address_mismatch_audit` row with `audit_status='backfill'`.

  2. Audit-only (--audit-only) — for the ~4452 rows that already have coords:
     reverse-geocode (lat, lon) → snapped point + canonical address, compute
     ST_Distance vs stored coords, write an `address_mismatch_audit` row with
     status 'ok' (≤50m) or 'mismatch' (>50m). Does NOT touch houses.

Design choices:
- **Per-row SAVEPOINT** (`db.begin_nested()`): a single Yandex/PostGIS error
  must not nuke the entire batch. Per backend.md, never use bare rollback
  inside a loop.
- **Resumable** via UNIQUE (house_id, audit_batch). Re-running the same
  --batch label skips already-processed houses, so a partial run can be
  picked up after CAPTCHA / network blip / 25k/day quota hit.
- **Precision filter**: backfill skips matches with precision in
  ('street', 'other', 'range', 'near', None) — those are too imprecise for
  comparable-listing spatial queries and would silently degrade matching
  recall. The audit row still records what Yandex returned for forensics.
- **Rate limit**: 50ms between calls (~20 req/sec, well under Yandex's
  25 req/sec service limit). Backfill mode runs single-threaded.
- **Daily quota**: 4141 backfill + 4452 audit ≈ 8.6k requests. Free Geocoder
  tier is 25k/day → comfortable buffer for retries.

Usage:
    YANDEX_GEOCODER_API_KEY=xxx \\
    DATABASE_URL=postgresql+psycopg://... \\
    python -m scripts.backfill_house_coords --batch 2026-05-27_backfill

    # Audit-only on the 4452 already-geocoded houses
    python -m scripts.backfill_house_coords --batch 2026-05-27_audit \\
        --audit-only --limit 500

Outputs:
- Backfill mode: UPDATE rows in `houses`, INSERT rows in
  `address_mismatch_audit` with status 'backfill' / 'no_match' / 'imprecise'.
- Audit mode: INSERT rows in `address_mismatch_audit` with status 'ok' /
  'mismatch' / 'no_match' / 'error'.
- Per-batch progress is logged every 25 rows.
"""

from __future__ import annotations

import argparse
import asyncio
import json
import logging
import os
from dataclasses import dataclass
from datetime import date
from pathlib import Path
from typing import Any

import httpx
from sqlalchemy import text
from sqlalchemy.orm import Session

# Allow running both as `python -m scripts.backfill_house_coords` (preferred)
# and as a stand-alone file. Mirrors the audit_address_mismatch import dance.
try:
    from app.core.db import SessionLocal  # type: ignore[import-not-found]
except ImportError:  # pragma: no cover — fallback for adhoc invocation
    import sys

    sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
    from app.core.db import SessionLocal

try:
    from scripts._yandex_reverse import (  # type: ignore[import-not-found]
        YandexReverseResult,
        forward_via_api,
        reverse_via_api,
    )
except ImportError:
    from _yandex_reverse import (  # type: ignore[no-redef]
        YandexReverseResult,
        forward_via_api,
        reverse_via_api,
    )

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s %(levelname)s %(name)s %(message)s",
)
logger = logging.getLogger("backfill_house_coords")

# Yandex Geocoder service limits per docs (as of 2026-05):
#   - 25k requests/day free tier
#   - 25 requests/sec sustained
# 50ms between calls = ~20 req/sec, leaving headroom for connection ramp-up.
_REQUEST_DELAY_S = 0.05

# Precision values we ACCEPT for backfill — anything else means Yandex didn't
# resolve to a specific building, and writing the result back into houses
# would degrade matching recall.
#   `exact`  → match found at the exact address (best case)
#   `number` → house number matched, but unit/entrance unspecified (acceptable)
# `near` / `range` / `street` / `other` / None → skipped (logged for analysis).
_BACKFILL_OK_PRECISION = frozenset({"exact", "number"})

# Audit threshold per issue #582 — distances above this flag a "mismatch"
# (the row still goes in the audit table, just with status='mismatch' for
# the report SQL to bucket separately).
_MISMATCH_DISTANCE_M = 50.0


# ---------------------------------------------------------------------------
# Domain types
# ---------------------------------------------------------------------------


@dataclass
class HouseRow:
    """One house from the source query — minimal fields needed for geocode."""

    id: int
    address: str
    lat: float | None
    lon: float | None


# ---------------------------------------------------------------------------
# Source-row queries
# ---------------------------------------------------------------------------


def _select_houses_without_coords(db: Session, limit: int | None) -> list[HouseRow]:
    """Pull houses needing forward geocode (lat IS NULL OR lon IS NULL).

    Skips rows with empty address — there's nothing to geocode there, they
    need a separate cleanup pass.
    """
    sql = (
        "SELECT id, address, lat, lon "
        "FROM houses "
        "WHERE (lat IS NULL OR lon IS NULL) "
        "  AND address IS NOT NULL "
        "  AND length(trim(address)) > 0 "
        "ORDER BY id"
    )
    if limit is not None:
        sql += " LIMIT CAST(:limit AS int)"
        rows = db.execute(text(sql), {"limit": limit}).mappings().all()
    else:
        rows = db.execute(text(sql)).mappings().all()
    return [
        HouseRow(id=r["id"], address=r["address"], lat=r["lat"], lon=r["lon"]) for r in rows
    ]


def _select_houses_with_coords(db: Session, limit: int | None) -> list[HouseRow]:
    """Pull houses needing reverse audit (both lat AND lon present)."""
    sql = (
        "SELECT id, address, lat, lon "
        "FROM houses "
        "WHERE lat IS NOT NULL "
        "  AND lon IS NOT NULL "
        "  AND address IS NOT NULL "
        "  AND length(trim(address)) > 0 "
        "ORDER BY id"
    )
    if limit is not None:
        sql += " LIMIT CAST(:limit AS int)"
        rows = db.execute(text(sql), {"limit": limit}).mappings().all()
    else:
        rows = db.execute(text(sql)).mappings().all()
    return [
        HouseRow(id=r["id"], address=r["address"], lat=r["lat"], lon=r["lon"]) for r in rows
    ]


def _already_processed_ids(db: Session, batch: str) -> set[int]:
    """house_ids already in address_mismatch_audit for this batch → skip set."""
    rows = db.execute(
        text("SELECT house_id FROM address_mismatch_audit WHERE audit_batch = CAST(:b AS text)"),
        {"b": batch},
    ).all()
    return {r[0] for r in rows}


# ---------------------------------------------------------------------------
# Distance helper — PostGIS, lon/lat order
# ---------------------------------------------------------------------------


def _distance_meters(
    db: Session, olat: float, olon: float, slat: float, slon: float
) -> float | None:
    """Great-circle distance (meters) via PostGIS geography type.

    Lifted from `audit_address_mismatch.py` to keep the two scripts using
    the same authority for distance computation. ST_MakePoint takes lon
    first per PostGIS convention.
    """
    row = db.execute(
        text(
            "SELECT ST_Distance("
            "  ST_SetSRID(ST_MakePoint(CAST(:olon AS double precision), "
            "                          CAST(:olat AS double precision)), 4326)::geography, "
            "  ST_SetSRID(ST_MakePoint(CAST(:slon AS double precision), "
            "                          CAST(:slat AS double precision)), 4326)::geography"
            ") AS m"
        ),
        {"olat": olat, "olon": olon, "slat": slat, "slon": slon},
    ).first()
    if row is None or row[0] is None:
        return None
    return float(row[0])


# ---------------------------------------------------------------------------
# DB writers
# ---------------------------------------------------------------------------


def _update_house_coords(
    db: Session,
    *,
    house_id: int,
    lat: float,
    lon: float,
    payload: dict[str, Any],
) -> None:
    """UPDATE houses SET lat/lon + merge yandex_geocode into raw_payload.

    The `houses_set_geom_trg` BEFORE UPDATE trigger (009_houses.sql) maintains
    `geom` automatically when lat/lon change, so we don't need to set geom
    explicitly here. `raw_payload || jsonb_build_object(...)` is the idiomatic
    psycopg-safe way to merge — single ALTER, no read-modify-write race.
    """
    db.execute(
        text(
            "UPDATE houses "
            "   SET lat = CAST(:lat AS double precision), "
            "       lon = CAST(:lon AS double precision), "
            "       raw_payload = COALESCE(raw_payload, '{}'::jsonb) "
            "                     || jsonb_build_object('yandex_geocode', "
            "                                           CAST(:payload AS jsonb)) "
            " WHERE id = CAST(:id AS bigint)"
        ),
        {"id": house_id, "lat": lat, "lon": lon, "payload": json.dumps(payload)},
    )


def _insert_audit_row(
    db: Session,
    *,
    house_id: int,
    batch: str,
    original_address: str | None,
    original_lat: float | None,
    original_lon: float | None,
    snapped_address: str | None,
    snapped_lat: float | None,
    snapped_lon: float | None,
    distance_m: float | None,
    audit_status: str,
    error_message: str | None,
    raw_payload: dict[str, Any] | None,
) -> None:
    """INSERT … ON CONFLICT DO NOTHING into address_mismatch_audit.

    Same column shape as `audit_address_mismatch._insert_audit_row` but the
    `district` and `street_differs` fields are left NULL — backfill/audit
    here doesn't have a stratification basis and we let the report SQL
    derive district at query time if needed (via Yandex address parse).

    Caller wraps in `begin_nested()` per backend.md SAVEPOINT pattern.
    """
    db.execute(
        text(
            "INSERT INTO address_mismatch_audit ("
            "  house_id, audit_batch, district,"
            "  original_address, original_lat, original_lon,"
            "  snapped_address, snapped_lat, snapped_lon,"
            "  distance_m, street_differs,"
            "  audit_status, error_message, raw_payload"
            ") VALUES ("
            "  CAST(:house_id AS bigint), CAST(:batch AS text), NULL,"
            "  :original_address, :original_lat, :original_lon,"
            "  :snapped_address, :snapped_lat, :snapped_lon,"
            "  :distance_m, NULL,"
            "  CAST(:audit_status AS text), :error_message,"
            "  CAST(:raw_payload AS jsonb)"
            ") ON CONFLICT (house_id, audit_batch) DO NOTHING"
        ),
        {
            "house_id": house_id,
            "batch": batch,
            "original_address": original_address,
            "original_lat": original_lat,
            "original_lon": original_lon,
            "snapped_address": snapped_address,
            "snapped_lat": snapped_lat,
            "snapped_lon": snapped_lon,
            "distance_m": distance_m,
            "audit_status": audit_status,
            "error_message": error_message,
            "raw_payload": json.dumps(raw_payload) if raw_payload is not None else None,
        },
    )


# ---------------------------------------------------------------------------
# Backfill loop (forward geocode, lat IS NULL houses)
# ---------------------------------------------------------------------------


def _classify_backfill_status(res: YandexReverseResult | None) -> str:
    """Translate a forward-geocode result into an audit_status value.

    'backfill' — Yandex returned a precise hit, lat/lon will be written.
    'imprecise' — match returned but precision is too low (street/other/...).
    'no_match' — Yandex returned an empty featureMember.
    'error' — handled by the caller's exception branch.
    """
    if res is None or res.address is None:
        return "no_match"
    if res.precision not in _BACKFILL_OK_PRECISION:
        return "imprecise"
    if res.snapped_lat is None or res.snapped_lon is None:
        return "no_match"
    return "backfill"


async def _run_backfill_mode(
    db: Session, sample: list[HouseRow], batch: str, api_key: str
) -> int:
    """Forward-geocode each house, UPDATE coords on precise hits, audit-log all."""
    processed = 0
    updated = 0
    n_imprecise = 0
    n_no_match = 0
    n_error = 0

    async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
        for i, row in enumerate(sample, start=1):
            status = "backfill"
            err: str | None = None
            res: YandexReverseResult | None = None
            try:
                res = await forward_via_api(row.address, api_key, client=client)
            except httpx.HTTPError as e:
                status = "error"
                err = f"http_error: {e!s}"
                n_error += 1
            except Exception as e:  # pragma: no cover — defensive
                status = "error"
                err = f"unhandled: {e!s}"
                n_error += 1

            if status != "error":
                status = _classify_backfill_status(res)
                if status == "imprecise":
                    n_imprecise += 1
                elif status == "no_match":
                    n_no_match += 1

            try:
                with db.begin_nested():
                    if status == "backfill" and res is not None and res.snapped_lat is not None:
                        # safe: status='backfill' guarantees snapped_lat/lon non-None.
                        assert res.snapped_lon is not None
                        _update_house_coords(
                            db,
                            house_id=row.id,
                            lat=res.snapped_lat,
                            lon=res.snapped_lon,
                            payload={
                                "address": res.address,
                                "precision": res.precision,
                                "kind": res.kind,
                                "batch": batch,
                                "source": "yandex_geocoder_api",
                            },
                        )
                        updated += 1
                    _insert_audit_row(
                        db,
                        house_id=row.id,
                        batch=batch,
                        original_address=row.address,
                        original_lat=row.lat,
                        original_lon=row.lon,
                        snapped_address=res.address if res else None,
                        snapped_lat=res.snapped_lat if res else None,
                        snapped_lon=res.snapped_lon if res else None,
                        distance_m=None,
                        audit_status=status,
                        error_message=err,
                        raw_payload=res.raw if res else None,
                    )
                # Per-row commit so resume picks up exactly where we crashed.
                db.commit()
                processed += 1
            except Exception as e:
                db.rollback()
                logger.warning("backfill insert failed for house_id=%s: %s", row.id, e)

            if i % 25 == 0:
                logger.info(
                    "backfill progress %d/%d updated=%d imprecise=%d no_match=%d error=%d",
                    i,
                    len(sample),
                    updated,
                    n_imprecise,
                    n_no_match,
                    n_error,
                )

            # Yandex 25 req/sec → 50ms between calls is plenty of headroom.
            if i < len(sample):
                await asyncio.sleep(_REQUEST_DELAY_S)

    logger.info(
        "backfill done: processed=%d updated=%d imprecise=%d no_match=%d error=%d",
        processed,
        updated,
        n_imprecise,
        n_no_match,
        n_error,
    )
    return processed


# ---------------------------------------------------------------------------
# Audit-only loop (reverse geocode, lat IS NOT NULL houses)
# ---------------------------------------------------------------------------


async def _run_audit_mode(
    db: Session, sample: list[HouseRow], batch: str, api_key: str
) -> int:
    """Reverse-geocode each house, compute distance, audit-log status/mismatch."""
    processed = 0
    n_ok = 0
    n_mismatch = 0
    n_no_match = 0
    n_error = 0

    async with httpx.AsyncClient(timeout=httpx.Timeout(10.0)) as client:
        for i, row in enumerate(sample, start=1):
            # Type-narrow: audit mode only feeds rows with non-null coords.
            assert row.lat is not None and row.lon is not None
            status = "ok"
            err: str | None = None
            res: YandexReverseResult | None = None
            try:
                res = await reverse_via_api(row.lat, row.lon, api_key, client=client)
            except httpx.HTTPError as e:
                status = "error"
                err = f"http_error: {e!s}"
                n_error += 1
            except Exception as e:  # pragma: no cover — defensive
                status = "error"
                err = f"unhandled: {e!s}"
                n_error += 1

            distance = None
            if res is not None and status == "ok":
                if res.address is None:
                    status = "no_match"
                    n_no_match += 1
                else:
                    if res.snapped_lat is not None and res.snapped_lon is not None:
                        distance = _distance_meters(
                            db, row.lat, row.lon, res.snapped_lat, res.snapped_lon
                        )
                        if distance is not None and distance > _MISMATCH_DISTANCE_M:
                            status = "mismatch"
                            n_mismatch += 1
                        else:
                            n_ok += 1
                    else:
                        n_ok += 1

            try:
                with db.begin_nested():
                    _insert_audit_row(
                        db,
                        house_id=row.id,
                        batch=batch,
                        original_address=row.address,
                        original_lat=row.lat,
                        original_lon=row.lon,
                        snapped_address=res.address if res else None,
                        snapped_lat=res.snapped_lat if res else None,
                        snapped_lon=res.snapped_lon if res else None,
                        distance_m=distance,
                        audit_status=status,
                        error_message=err,
                        raw_payload=res.raw if res else None,
                    )
                db.commit()
                processed += 1
            except Exception as e:
                db.rollback()
                logger.warning("audit insert failed for house_id=%s: %s", row.id, e)

            if i % 25 == 0:
                logger.info(
                    "audit progress %d/%d ok=%d mismatch=%d no_match=%d error=%d",
                    i,
                    len(sample),
                    n_ok,
                    n_mismatch,
                    n_no_match,
                    n_error,
                )

            if i < len(sample):
                await asyncio.sleep(_REQUEST_DELAY_S)

    logger.info(
        "audit done: processed=%d ok=%d mismatch=%d no_match=%d error=%d",
        processed,
        n_ok,
        n_mismatch,
        n_no_match,
        n_error,
    )
    return processed


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


def _parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    """argparse setup, factored out for testability."""
    p = argparse.ArgumentParser(
        description=(
            "Phase 2-3 of issue #582 — backfill houses.lat/lon via Yandex forward "
            "geocode, or audit already-geocoded houses via reverse geocode."
        ),
    )
    p.add_argument(
        "--batch",
        default=f"{date.today().isoformat()}_backfill",
        help="Audit batch label. Same batch re-run skips already-processed houses.",
    )
    p.add_argument(
        "--audit-only",
        action="store_true",
        help=(
            "Run reverse-geocode audit on houses WITH coords instead of forward "
            "backfill on houses WITHOUT coords. Does not modify the houses table."
        ),
    )
    p.add_argument(
        "--limit",
        type=int,
        default=None,
        help=(
            "Optional cap on source-row count. Useful for canary runs "
            "(e.g. --limit 100 before letting the full 4k loose)."
        ),
    )
    return p.parse_args(argv)


async def main(argv: list[str] | None = None) -> int:
    """CLI entry point. Returns the number of rows processed this run."""
    args = _parse_args(argv)
    api_key = os.environ.get("YANDEX_GEOCODER_API_KEY")
    if not api_key:
        raise SystemExit(
            "YANDEX_GEOCODER_API_KEY is required — forward geocode is API-only."
        )

    mode = "audit" if args.audit_only else "backfill"
    logger.info(
        "starting batch=%s mode=%s limit=%s",
        args.batch,
        mode,
        args.limit if args.limit is not None else "all",
    )

    db = SessionLocal()
    try:
        if args.audit_only:
            sample = _select_houses_with_coords(db, args.limit)
        else:
            sample = _select_houses_without_coords(db, args.limit)
        logger.info("loaded source rows: %d", len(sample))

        done = _already_processed_ids(db, args.batch)
        if done:
            logger.info(
                "resuming batch %s: %d rows already processed",
                args.batch,
                len(done),
            )
        remaining = [s for s in sample if s.id not in done]
        if not remaining:
            logger.info("nothing to do — batch %s is complete for the loaded sample", args.batch)
            return 0

        if args.audit_only:
            n = await _run_audit_mode(db, remaining, args.batch, api_key)
        else:
            n = await _run_backfill_mode(db, remaining, args.batch, api_key)

        logger.info("done: processed=%d batch=%s mode=%s", n, args.batch, mode)
        return n
    finally:
        db.close()


if __name__ == "__main__":  # pragma: no cover
    asyncio.run(main())
