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

#!/usr/bin/env python3
"""
import_heightmaps.py — Import terrain elevation grids into atlas_body_heightmaps.

For each inhabited body with a terrain_reference, simulates the terrain via
planet_simulation.simulate() and stores the float32 LE elevation BLOB plus
sea_level metadata in atlas_body_heightmaps (#906, D-202).

The BLOB format matches the Rust loader spec (D-202):
  - float32 little-endian, row-major
  - width × height values, each in [0.0, 1.0]
  - width = GRID_W (512), height = GRID_H (256)

Incremental: bodies that already have a row in atlas_body_heightmaps are
skipped unless --force is passed.

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

Exit codes:
    0  completed (possibly with skipped or errored bodies)
    1  fatal error (missing DB, schema error)
"""

import argparse
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 numpy as np
import sqlite3

from atlas_common import (
    GRID_W,
    GRID_H,
    DB_PATH,
    ensure_atlas_schema,
    load_body_def,
    query_inhabited_bodies,
)
from planet_simulation import simulate


def _elevation_to_blob(elevation: np.ndarray) -> bytes:
    """Convert a float32 elevation grid to a little-endian BLOB."""
    arr = elevation.astype("<f4")  # float32 LE, explicit
    assert arr.shape == (GRID_H, GRID_W), (
        f"elevation shape {arr.shape} does not match expected ({GRID_H}, {GRID_W})"
    )
    return arr.tobytes()


def import_body(
    body_info: dict,
    conn: sqlite3.Connection,
    force: bool,
    dry_run: bool,
    verbose: bool,
) -> dict:
    """Import heightmap BLOB for one body.

    Returns a dict with:
        status:  'imported' | 'skipped' | 'gas_giant' | 'error'
        message: detail (on error or skip)
    """
    body_id = body_info["body_id"]
    terrain_ref = body_info["terrain_reference"]

    # Incremental check — skip if already imported
    if not force:
        existing = conn.execute(
            "SELECT 1 FROM atlas_body_heightmaps WHERE body_id = ?", (body_id,)
        ).fetchone()
        if existing:
            return {"status": "skipped", "message": "already imported"}

    body_dir = REPO_ROOT / Path(terrain_ref).parent
    if not body_dir.exists():
        return {"status": "error", "message": f"body_dir not found: {body_dir}"}

    bd = load_body_def(body_dir)
    if not bd:
        return {"status": "error", "message": f"no body definition found in {body_dir}"}

    try:
        terrain = simulate(bd)
    except Exception as exc:
        return {"status": "error", "message": f"simulate() failed: {exc}"}

    if not terrain:
        return {"status": "gas_giant"}

    elevation = terrain.get("elevation")
    if elevation is None:
        return {"status": "error", "message": "terrain dict missing 'elevation' key"}

    sea_level = float(terrain.get("sea_level", 0.0))
    blob = _elevation_to_blob(elevation)

    if verbose:
        land_pct = float(np.mean(elevation >= sea_level)) * 100
        print(f"    {body_id}: {GRID_W}x{GRID_H} grid, sea_level={sea_level:.3f}, "
              f"land={land_pct:.1f}%, blob={len(blob)} bytes")

    if not dry_run:
        conn.execute(
            """INSERT INTO atlas_body_heightmaps
                   (body_id, width, height, data, sea_level, imported_at)
               VALUES (?, ?, ?, ?, ?, datetime('now'))
               ON CONFLICT(body_id) DO UPDATE SET
                   width = excluded.width,
                   height = excluded.height,
                   data = excluded.data,
                   sea_level = excluded.sea_level,
                   imported_at = excluded.imported_at""",
            (body_id, GRID_W, GRID_H, blob, sea_level),
        )

    return {"status": "imported"}


def main() -> None:
    parser = argparse.ArgumentParser(
        description="Import terrain elevation BLOBs into atlas_body_heightmaps (#906)"
    )
    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("--force", action="store_true",
                        help="Re-import even if a row already exists")
    parser.add_argument("--dry-run", action="store_true",
                        help="Simulate 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  Heightmap BLOB Import (#906)")
    print(f"  DB: {db_path}")
    if args.dry_run:
        print(f"  Mode: DRY RUN (no DB writes)")
    if args.force:
        print(f"  Force: enabled (will overwrite existing rows)")
    print()

    conn = sqlite3.connect(str(db_path))
    conn.execute("PRAGMA foreign_keys=ON")
    ensure_atlas_schema(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\n")

    t_total = time.time()
    n_imported = 0
    n_skipped = 0
    n_gas = 0
    n_errors = 0

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

        result = import_body(body_info, conn, args.force, args.dry_run, args.verbose)
        elapsed = time.time() - t0
        status = result["status"]

        if status == "imported":
            n_imported += 1
            print(f"  [{i+1}/{len(bodies)}] {body_id:20s} imported ({elapsed:.1f}s)")
        elif status == "skipped":
            n_skipped += 1
            if args.verbose:
                print(f"  [{i+1}/{len(bodies)}] {body_id:20s} skipped (already imported)")
        elif status == "gas_giant":
            n_gas += 1
            if args.verbose:
                print(f"  [{i+1}/{len(bodies)}] {body_id:20s} gas giant — no surface")
        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_total = time.time() - t_total
    print(f"\n  Done in {elapsed_total:.1f}s")
    print(f"  imported={n_imported}  skipped={n_skipped}  "
          f"gas_giant={n_gas}  errors={n_errors}")

    if n_errors > 0:
        print(f"\n  {n_errors} error(s) — check output above", file=sys.stderr)


if __name__ == "__main__":
    main()