settled-reach/tooling/planet-gen/import_city_names.py

#!/usr/bin/env python3
"""
import_city_names.py — Populate atlas_city_names from wiki markers.json content.

For each inhabited body, reads city records from markers.json and inserts rows
into atlas_city_names with:
  - name, kind, population from markers.json
  - economic_role from bodies table
  - corp_id from corporations.headquarters_body cross-reference (#909)

Incremental: clears and reimports all rows for each body on every run (the
table has no stable local IDs — city identity is name × body_id). Use --body
to restrict to a single body.

Usage:
    tooling/planet-gen/import_city_names.py
    tooling/planet-gen/import_city_names.py --body GJ380c
    tooling/planet-gen/import_city_names.py --dry-run

Exit codes:
    0  completed
    1  fatal error (missing DB, schema error)
"""

import argparse
import json
import sys
import time
from pathlib import Path

TOOLING_DIR = Path(__file__).resolve().parent
REPO_ROOT = (TOOLING_DIR / ".." / "..").resolve()

_venv_python = REPO_ROOT / ".venv" / "bin" / "python"
if _venv_python.exists() and Path(sys.executable).resolve() != _venv_python.resolve():
    import os
    os.execv(str(_venv_python), [str(_venv_python)] + sys.argv)

import sqlite3

from generate_atlas import (
    DB_PATH,
    ensure_atlas_schema,
    query_inhabited_bodies,
)


def _build_hq_index(conn: sqlite3.Connection) -> dict[str, str]:
    """Build a mapping of body_id → corp_id for all corp HQ locations."""
    rows = conn.execute(
        "SELECT headquarters_body, corp_id FROM corporations "
        "WHERE headquarters_body IS NOT NULL"
    ).fetchall()
    index: dict[str, str] = {}
    for body_id, corp_id in rows:
        # If multiple corps have the same HQ body, take the first (alphabetical
        # corp_id for determinism). This is unlikely but safe.
        if body_id not in index:
            index[body_id] = corp_id
    return index


def _load_city_records(body_dir: Path) -> list[dict]:
    """Load named city records from markers.json. Returns empty list if none."""
    markers_path = body_dir / "markers.json"
    if not markers_path.exists():
        return []
    try:
        markers = json.loads(markers_path.read_text())
    except json.JSONDecodeError:
        return []
    return [
        c for c in (markers.get("cities") or [])
        if c.get("name") and isinstance(c["name"], str) and c["name"].strip()
    ]


def import_body_cities(
    body_info: dict,
    conn: sqlite3.Connection,
    hq_index: dict[str, str],
    dry_run: bool,
    verbose: bool,
) -> dict:
    """Import atlas_city_names rows for one body.

    Returns a dict with:
        status:   'imported' | 'no_cities' | 'error'
        imported: count of rows written
        message:  detail (on error)
    """
    body_id = body_info["body_id"]
    terrain_ref = body_info["terrain_reference"]
    economic_role = body_info.get("economic_role") or "unknown"
    corp_id = hq_index.get(body_id)

    body_dir = REPO_ROOT / Path(terrain_ref).parent
    cities = _load_city_records(body_dir)

    if not cities:
        return {"status": "no_cities", "imported": 0}

    if verbose:
        print(f"    {body_id}: {len(cities)} cities, economic_role={economic_role}"
              + (f", corp_hq={corp_id}" if corp_id else ""))

    if not dry_run:
        # Full rebuild for this body: delete existing rows, re-insert.
        conn.execute("DELETE FROM atlas_city_names WHERE body_id = ?", (body_id,))

        for city in cities:
            name = city["name"].strip()
            kind = city.get("kind") or "city"
            population = int(city.get("population") or 0)
            # Only set corp_id on the capital city of a corp HQ body.
            city_corp_id = corp_id if (kind == "capital" and corp_id) else None

            conn.execute(
                """INSERT INTO atlas_city_names
                       (body_id, name, kind, economic_role, population, corp_id,
                        reserved, updated_at)
                   VALUES (?, ?, ?, ?, ?, ?, 0, datetime('now'))""",
                (body_id, name, kind, economic_role, population, city_corp_id),
            )

    return {"status": "imported", "imported": len(cities)}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Populate atlas_city_names from wiki markers.json (#908, #909)"
    )
    parser.add_argument("--db", default=str(DB_PATH), help="Path to systems.db")
    parser.add_argument("--body", help="Process only this body_id")
    parser.add_argument("--dry-run", action="store_true",
                        help="Read and validate without writing to DB")
    parser.add_argument("--verbose", action="store_true",
                        help="Print per-body detail")
    args = parser.parse_args()

    db_path = Path(args.db)
    if not db_path.exists():
        print(f"error: {db_path} not found", file=sys.stderr)
        sys.exit(1)

    print(f"\n  City Names Import (#908 + #909)")
    print(f"  DB: {db_path}")
    if args.dry_run:
        print(f"  Mode: DRY RUN (no DB writes)")
    print()

    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA foreign_keys=ON")
    ensure_atlas_schema(conn)

    hq_index = _build_hq_index(conn)

    bodies = query_inhabited_bodies(conn)
    if args.body:
        bodies = [b for b in bodies if b["body_id"] == args.body]
        if not bodies:
            print(f"error: body '{args.body}' not found or has no terrain_reference",
                  file=sys.stderr)
            conn.close()
            sys.exit(1)

    print(f"  {len(bodies)} inhabited bodies with terrain_reference")
    print(f"  {len(hq_index)} corp HQ body mappings\n")

    t_total = time.time()
    n_imported = 0
    n_no_cities = 0
    n_errors = 0
    total_rows = 0

    for i, body_info in enumerate(bodies):
        body_id = body_info["body_id"]

        result = import_body_cities(body_info, conn, hq_index, args.dry_run, args.verbose)
        status = result["status"]

        if status == "imported":
            n_imported += 1
            total_rows += result["imported"]
            if args.verbose:
                print(f"  [{i+1}/{len(bodies)}] {body_id:20s} {result['imported']} cities")
        elif status == "no_cities":
            n_no_cities += 1
        elif status == "error":
            n_errors += 1
            print(f"  [{i+1}/{len(bodies)}] {body_id:20s} ERROR: {result.get('message', '')}")

    if not args.dry_run:
        conn.commit()

    conn.close()

    elapsed = time.time() - t_total
    print(f"\n  Done in {elapsed:.1f}s")
    print(f"  bodies_with_cities={n_imported}  no_cities={n_no_cities}  "
          f"errors={n_errors}  total_rows={total_rows}")


if __name__ == "__main__":
    main()