#!/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):

    reach atlas planet import-heightmaps            # full bake
    reach atlas planet import-heightmaps --limit 3  # smoke test
    reach atlas planet import-heightmaps --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

from tooling.core import config, console
from tooling.core.errors import ReachError

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

PLANET_DIR = Path(__file__).resolve().parent
REPO_ROOT = config.repo_root()

from tooling.domains.atlas.planet.body_definition_parser import parse_system  # noqa: E402
from tooling.domains.atlas.planet.planet_simulation import simulate  # noqa: E402
from tooling.domains.atlas.planet.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(argv: list[str] | None = None) -> 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():
        raise ReachError(f"{db_path} not found", fix="make regen-db, or point --db at an existing systems.db")

    console.event("  Heightmap bake (#963, D-202 amended)")
    if args.dry_run:
        console.event("  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)
        console.event(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]

    console.event(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
            console.event(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:
            console.event(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:
                console.event(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
            console.event(f"  [{i+1}/{len(rows)}] {body_id:18s} ERROR: {res.get('message','')}")

    console.event(f"  Done in {time.time()-t0:.0f}s — baked={n_baked} gas_giant={n_gas} errors={n_err}")
    console.event("  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()