#!/usr/bin/env python3
"""
import_heightmaps.py — bake per-body canonical elevation assets (D-202 amended, #963).

This is the build-time elevation producer. For each NON-Sol inhabited body it
runs planet_simulation.simulate() at the canonical 1024×512 grid and writes two
files into the body's wiki directory:

  reliefmap.png   — clean color hypsometric render (display / Atlas), 1024×512
  heightmap.png   — 16-bit GRAYSCALE elevation, the canonical cascade input,
                    with sea_level stored in a PNG tEXt chunk

It also performs the one-time rename of the legacy color `heightmap.png`
(which was a relief render) → `reliefmap.png` across ALL bodies, including Sol
(GJ-0) — Sol's terrain is real-geography (sol_import) and is NOT re-simulated,
only its display file is renamed.

No systems.db writes: the PNG is the store (the atlas_body_heightmaps BLOB
table is dropped). Run with uv (numpy/scipy/Pillow):

    uv run python tooling/planet-gen/import_heightmaps.py            # full bake
    uv run python tooling/planet-gen/import_heightmaps.py --limit 3  # smoke test
    uv run python tooling/planet-gen/import_heightmaps.py --dry-run

Exit codes: 0 = completed (possibly with per-body errors), 1 = fatal.
"""

import argparse
import glob
import sqlite3
import sys
import time
from pathlib import Path

import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo

TOOLING_DIR = Path(__file__).resolve().parent
REPO_ROOT = (TOOLING_DIR / ".." / "..").resolve()
sys.path.insert(0, str(TOOLING_DIR))

from body_definition_parser import parse_system  # noqa: E402
from planet_simulation import simulate  # noqa: E402
from render_heightmap import render_heightmap  # noqa: E402

DB_PATH = REPO_ROOT / "server" / "data" / "systems.db"
WIKI_SYSTEMS = REPO_ROOT / "wiki" / "star-systems"


def rename_legacy_heightmaps(dry_run: bool) -> int:
    """One-time: legacy color `heightmap.png` (a relief render) → `reliefmap.png`,
    for every body. Bodies the bake re-simulates get their `reliefmap.png`
    overwritten with a fresh 1024 render afterwards. Idempotent: skips bodies
    that already have `reliefmap.png`."""
    n = 0
    for hm in sorted(glob.glob(str(WIKI_SYSTEMS / "*" / "bodies" / "*" / "heightmap.png"))):
        relief = Path(hm).with_name("reliefmap.png")
        if relief.exists():
            continue  # already migrated (or freshly baked)
        if not dry_run:
            Path(hm).rename(relief)
        n += 1
    return n


def bake_body(bd: dict, body_dir: Path, dry_run: bool) -> dict:
    """Simulate one body at 1024×512 and write reliefmap.png + 16-bit heightmap.png."""
    try:
        terrain = simulate(bd)
    except Exception as exc:  # noqa: BLE001
        return {"status": "error", "message": f"simulate failed: {exc}"}
    if not terrain or "elevation" not in terrain:
        return {"status": "gas_giant"}

    elevation = np.clip(np.asarray(terrain["elevation"], dtype=np.float32), 0.0, 1.0)
    sea_level = float(terrain.get("sea_level", 0.0))

    if not dry_run:
        # Clean color relief (no painted features — cascade/Atlas overlay those).
        relief = render_heightmap(bd, terrain, render_mode="cartographic", chrome=False)
        relief.save(body_dir / "reliefmap.png")
        # Canonical 16-bit grayscale elevation; sea_level rides in a tEXt chunk.
        u16 = (elevation * 65535.0).astype(np.uint16)
        meta = PngInfo()
        meta.add_text("sea_level", f"{sea_level:.6f}")
        Image.fromarray(u16).save(body_dir / "heightmap.png", pnginfo=meta)

    return {"status": "baked", "shape": elevation.shape, "sea_level": sea_level}


def main() -> None:
    ap = argparse.ArgumentParser(description="Bake canonical heightmap/reliefmap assets (#963)")
    ap.add_argument("--db", default=str(DB_PATH))
    ap.add_argument("--body", help="Bake only this body_id")
    ap.add_argument("--limit", type=int, help="Bake at most N bodies (smoke test)")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--skip-rename", action="store_true", help="Skip the legacy rename pass")
    args = ap.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("\n  Heightmap bake (#963, D-202 amended)")
    if args.dry_run:
        print("  Mode: DRY RUN (no files written)")

    # 1. Universal rename of the legacy color heightmap.png → reliefmap.png.
    if not args.skip_rename:
        n_renamed = rename_legacy_heightmaps(args.dry_run)
        print(f"  Renamed legacy heightmap.png → reliefmap.png: {n_renamed} bodies")

    # 2. Bake non-Sol inhabited bodies.
    conn = sqlite3.connect(str(db_path))
    rows = conn.execute(
        "SELECT body_id, terrain_reference FROM bodies "
        "WHERE inhabited=1 AND terrain_reference IS NOT NULL AND population>0 "
        "AND system_id != 'GJ 0' ORDER BY body_id"
    ).fetchall()
    conn.close()
    if args.body:
        rows = [r for r in rows if r[0] == args.body]
    if args.limit:
        rows = rows[: args.limit]

    print(f"  Baking {len(rows)} non-Sol inhabited bodies\n")

    parsed: dict[str, list] = {}
    t0 = time.time()
    n_baked = n_gas = n_err = 0
    for i, (body_id, tref) in enumerate(rows):
        body_dir = REPO_ROOT / Path(tref).parent
        sys_index = str(REPO_ROOT / Path(tref.split("/bodies/")[0]) / "index.md")
        try:
            defs = parsed.get(sys_index) or parse_system(sys_index)
            parsed[sys_index] = defs
        except Exception as exc:  # noqa: BLE001
            print(f"  [{i+1}/{len(rows)}] {body_id:18s} ERROR parse_system: {exc}")
            n_err += 1
            continue
        bd = next((d for d in defs if d.get("id") == body_id), None)
        if bd is None:
            print(f"  [{i+1}/{len(rows)}] {body_id:18s} ERROR: not in system defs")
            n_err += 1
            continue
        ts = time.time()
        res = bake_body(bd, body_dir, args.dry_run)
        dt = time.time() - ts
        if res["status"] == "baked":
            n_baked += 1
            if (i + 1) % 25 == 0 or args.limit:
                print(f"  [{i+1}/{len(rows)}] {body_id:18s} baked {res['shape']} "
                      f"sea={res['sea_level']:.3f} ({dt:.1f}s)")
        elif res["status"] == "gas_giant":
            n_gas += 1
        else:
            n_err += 1
            print(f"  [{i+1}/{len(rows)}] {body_id:18s} ERROR: {res.get('message','')}")

    print(f"\n  Done in {time.time()-t0:.0f}s — baked={n_baked} gas_giant={n_gas} errors={n_err}")
    print("  Stage: git add wiki/star-systems  (reliefmap.png renames + heightmap.png)")
    if n_err:
        sys.exit(0)  # per-body errors are non-fatal; reported above


if __name__ == "__main__":
    main()