#!/usr/bin/env python3
"""
Batch planet generation — process all systems unattended.

Walks wiki/star-systems/, scaffolds body index.md files where missing,
then generates heightmap + globe + terrain data for every body.

Usage:
    python3 batch.py                          # full run: scaffold + generate
    python3 batch.py --scaffold-only          # just create body index.md files
    python3 batch.py --generate-only          # just render (bodies must exist)
    python3 batch.py --system GJ-144          # single system
    python3 batch.py --system GJ-144 --body GJ144d  # single body
    python3 batch.py --overrides sol.json     # per-body overrides
    python3 batch.py --dry-run                # validate data, don't generate

Skips:
    - GJ-0 (Sol) — manual overrides required, use --system GJ-0 explicitly
    - asteroid_belt, oort_cloud (no renderable surface)
    - Bodies that already have globe.png (use --force to regenerate)

Error handling:
    - Errors are printed to console and logged to /tmp/planet-gen-errors.log
    - Batch aborts if error rate exceeds 50% of attempted bodies
"""

import argparse
import json
import os
import sys
import time
import traceback
from datetime import datetime
from pathlib import Path

# Venv bootstrap
TOOLING_DIR = Path(__file__).resolve().parent
WORKTREE_ROOT = (TOOLING_DIR / ".." / "..").resolve()
_venv_python = WORKTREE_ROOT / ".venv" / "bin" / "python"
if _venv_python.exists() and Path(sys.executable).resolve() != _venv_python.resolve():
    os.execv(str(_venv_python), [str(_venv_python)] + sys.argv)

import yaml
import numpy as np

from body_definition_parser import parse_system
from planet_simulation import simulate
from render_heightmap import render_heightmap

# Systems to skip in batch mode (require manual handling)
SKIP_SYSTEMS = {"GJ-0"}

LOG_PATH = Path("/tmp/planet-gen-errors.log")


# ─────────────────────────────────────────────────────────────────────────────
# Logging
# ─────────────────────────────────────────────────────────────────────────────

def _log_error(system_id: str, body_id: str, error: str, tb: str):
    """Append error to log file with context."""
    with open(LOG_PATH, "a") as f:
        f.write(f"\n{'='*72}\n")
        f.write(f"  time:   {datetime.now().isoformat()}\n")
        f.write(f"  system: {system_id}\n")
        f.write(f"  body:   {body_id}\n")
        f.write(f"  error:  {error}\n")
        f.write(f"  traceback:\n{tb}\n")


# ─────────────────────────────────────────────────────────────────────────────
# System helpers
# ─────────────────────────────────────────────────────────────────────────────

def _read_system_name(index_md: Path) -> str:
    """Extract the system name from the # heading in index.md."""
    for line in index_md.read_text().splitlines()[:5]:
        if line.startswith("# "):
            return line[2:].strip()
    return index_md.parent.name


def _read_frontmatter(md_path: Path) -> dict:
    """Read YAML frontmatter from a body index.md."""
    content = md_path.read_text()
    if not content.startswith("---"):
        return {}
    try:
        end = content.index("---", 3)
        return yaml.safe_load(content[3:end]) or {}
    except (ValueError, yaml.YAMLError):
        return {}


# ─────────────────────────────────────────────────────────────────────────────
# Validation (dry-run)
# ─────────────────────────────────────────────────────────────────────────────

REQUIRED_FIELDS = ["id", "planet_class", "seed", "star", "orbit", "physical"]
REQUIRED_STAR = ["type", "luminosity_solar"]
REQUIRED_ORBIT = ["distance_au", "period_days", "axial_tilt_deg"]
REQUIRED_PHYSICAL = ["gravity_g", "atmosphere"]
REQUIRED_TERRAIN = ["land_fraction", "tectonics"]


def _validate_body_def(bd: dict, body_dir: Path) -> list:
    """Validate a body definition. Returns list of error strings."""
    errors = []
    bid = bd.get("id", "?")

    for f in REQUIRED_FIELDS:
        if f not in bd:
            errors.append(f"{bid}: missing top-level field '{f}'")

    star = bd.get("star", {})
    for f in REQUIRED_STAR:
        if f not in star:
            errors.append(f"{bid}: missing star.{f}")

    orbit = bd.get("orbit", {})
    for f in REQUIRED_ORBIT:
        v = orbit.get(f)
        if v is None or v == "rand":
            errors.append(f"{bid}: orbit.{f} is {v!r} — must be resolved (not 'rand')")
        elif isinstance(v, (int, float)) and v <= 0:
            errors.append(f"{bid}: orbit.{f} = {v} — must be positive")

    phys = bd.get("physical", {})
    for f in REQUIRED_PHYSICAL:
        if f not in phys:
            errors.append(f"{bid}: missing physical.{f}")

    seed = bd.get("seed")
    if seed is None or not isinstance(seed, int):
        errors.append(f"{bid}: seed must be an integer, got {seed!r}")

    pclass = bd.get("planet_class", "")
    valid_classes = {"temperate", "oceanic", "forest", "arid", "frozen",
                     "volcanic", "barren", "gas_giant", "gas_giant_ringed",
                     "cold_arid", "hot_arid", "tropical", "boreal",
                     "temperate_terminator"}
    if pclass not in valid_classes:
        errors.append(f"{bid}: planet_class '{pclass}' not in {valid_classes}")

    is_gas = pclass in ("gas_giant", "gas_giant_ringed")
    terrain = bd.get("terrain", {})
    if not is_gas:
        for f in REQUIRED_TERRAIN:
            if f not in terrain:
                errors.append(f"{bid}: missing terrain.{f}")
        lf = terrain.get("land_fraction")
        if lf is not None and (lf < 0 or lf > 1):
            errors.append(f"{bid}: terrain.land_fraction = {lf} — must be [0, 1]")

    return errors


# ─────────────────────────────────────────────────────────────────────────────
# Scaffold
# ─────────────────────────────────────────────────────────────────────────────

def _scaffold_system(system_dir: Path, overrides: dict) -> list:
    """Scaffold body index.md files for one system. Returns body IDs created."""
    index_md = system_dir / "index.md"
    if not index_md.exists():
        return []

    from scaffold_bodies import _body_to_frontmatter, _body_prose

    body_defs = parse_system(str(index_md), overrides=overrides)
    bodies_dir = system_dir / "bodies"
    created = []

    for bd in body_defs:
        body_dir = bodies_dir / bd["id"]
        body_index = body_dir / "index.md"
        if body_index.exists():
            continue
        body_dir.mkdir(parents=True, exist_ok=True)
        fm = _body_to_frontmatter(bd)
        prose = _body_prose(bd, system_dir)
        body_index.write_text(f"---\n{fm}\n---\n\n{prose}")
        created.append(bd["id"])

    return created


# ─────────────────────────────────────────────────────────────────────────────
# Generate
# ─────────────────────────────────────────────────────────────────────────────

def _save_img(img, target: Path):
    """Write image directly — resume logic handles incomplete outputs."""
    img.save(str(target), format="PNG")


def _save(data, target: Path, save_fn):
    """Write data directly — resume logic handles incomplete outputs."""
    save_fn(data, str(target))


def _is_complete(body_dir: Path, is_gas: bool) -> bool:
    """Check if all expected outputs exist (for resume support)."""
    if not (body_dir / "globe.png").exists():
        return False
    if not is_gas:
        for f in ("heightmap.png", "terrain.npz", "markers.json"):
            if not (body_dir / f).exists():
                return False
    return True


def _generate_body_from_dir(body_dir: Path, hmap_w: int, hmap_h: int,
                            globe_size: int, force: bool) -> str:
    """
    Generate assets for one body.
    Returns: "generated", "skipped", or "error"
    """
    index_md = body_dir / "index.md"
    if not index_md.exists():
        return "skipped"

    bd = _read_frontmatter(index_md)
    if not bd or "id" not in bd or "planet_class" not in bd:
        return "skipped"

    body_id = bd["id"]
    name = bd.get("name") or body_id
    planet_class = bd.get("planet_class", "")
    is_gas = planet_class in ("gas_giant", "gas_giant_ringed")

    # Resume: skip if all outputs exist (unless --force)
    if not force and _is_complete(body_dir, is_gas):
        return "skipped"

    t0 = time.time()

    # Simulate
    terrain = simulate(bd)
    is_gas = not terrain  # re-check from actual simulation result

    # Heightmap — atomic write
    if not is_gas:
        hmap_img = render_heightmap(bd, terrain,
                                    out_w=hmap_w, out_h=hmap_h,
                                    chrome=False)
        _save_img(hmap_img, body_dir / "heightmap.png")

    # Globe — atomic write
    from planet_renderer import render_globe
    globe_img = render_globe(bd, terrain, size=globe_size)
    _save_img(globe_img, body_dir / "globe.png")

    # Terrain data — atomic write
    if not is_gas:
        save_dict = {}
        for key in ("elevation", "temperature", "moisture", "hillshade",
                    "biome", "surface_water", "river_grid"):
            if key in terrain:
                save_dict[key] = terrain[key]
        save_dict["sea_level"] = np.array([terrain["sea_level"]])
        np.savez_compressed(str(body_dir / "terrain.npz"), **save_dict)

        # Markers
        from generate import _build_markers
        markers = _build_markers(bd, terrain)
        _save(markers, body_dir / "markers.json",
                     lambda m, p: Path(p).write_text(json.dumps(m, indent=2)))

    elapsed = time.time() - t0
    kind = "gas" if is_gas else f"land={int((~terrain['surface_water']).sum())}"
    print(f"      {body_id:20s} ({name:20s}) {elapsed:5.1f}s  {kind}")
    return "generated"


# ─────────────────────────────────────────────────────────────────────────────
# Main
# ─────────────────────────────────────────────────────────────────────────────

def main():
    parser = argparse.ArgumentParser(
        description="Batch planet generation — all systems unattended")
    parser.add_argument("--system", help="Process only this system (dir name, e.g. GJ-144)")
    parser.add_argument("--body", help="Process only this body (requires --system)")
    parser.add_argument("--scaffold-only", action="store_true")
    parser.add_argument("--generate-only", action="store_true")
    parser.add_argument("--overrides", help="Per-body overrides JSON")
    parser.add_argument("--force", action="store_true",
                        help="Regenerate even if globe.png exists")
    parser.add_argument("--dry-run", action="store_true",
                        help="Validate body definitions without generating")
    parser.add_argument("--verify-determinism", type=int, metavar="N", default=0,
                        help="Run N random bodies twice and verify identical output")
    parser.add_argument("--heightmap-size", default="1024x512")
    parser.add_argument("--globe-size", type=int, default=512)
    args = parser.parse_args()

    hw, hh = args.heightmap_size.lower().split("x")
    hmap_w, hmap_h = int(hw), int(hh)

    wiki_systems = WORKTREE_ROOT / "wiki" / "star-systems"
    if not wiki_systems.exists():
        print(f"error: {wiki_systems} not found", file=sys.stderr)
        sys.exit(1)

    overrides = {}
    if args.overrides:
        with open(args.overrides) as f:
            overrides = json.load(f)

    # Collect system directories
    if args.system:
        system_dirs = [wiki_systems / args.system]
        if not system_dirs[0].exists():
            print(f"error: system {args.system} not found", file=sys.stderr)
            sys.exit(1)
    else:
        system_dirs = sorted([
            d for d in wiki_systems.iterdir()
            if d.is_dir() and (d / "index.md").exists()
        ])

    # Clear error log
    if LOG_PATH.exists():
        LOG_PATH.unlink()

    t_total = time.time()
    total_scaffolded = 0
    total_generated = 0
    total_skipped = 0
    total_errors = 0
    total_attempted = 0
    total_valid = 0
    total_invalid = 0
    error_rate_threshold = 0.50

    print(f"\n  Planet Generator — Batch Mode")
    print(f"  Systems: {len(system_dirs)}")
    print(f"  Heightmap: {hmap_w}×{hmap_h}  Globe: {args.globe_size}×{args.globe_size}")
    if args.dry_run:
        print(f"  Mode: DRY RUN (validation only)")
    print()

    for system_dir in system_dirs:
        system_id = system_dir.name

        # Skip Sol in batch mode (needs manual overrides)
        if system_id in SKIP_SYSTEMS and not args.system:
            print(f"  {system_id} — skipped (manual)")
            continue

        system_name = _read_system_name(system_dir / "index.md")

        # Count bodies before starting
        bodies_dir_check = system_dir / "bodies"
        if bodies_dir_check.exists():
            n_bodies = sum(1 for d in bodies_dir_check.iterdir()
                          if d.is_dir() and (d / "index.md").exists())
        else:
            # Peek at the system table to estimate body count
            try:
                defs = parse_system(str(system_dir / "index.md"), overrides=overrides)
                n_bodies = len(defs)
            except Exception:
                n_bodies = 0

        print(f"  {system_id} — {system_name} ({n_bodies} bodies)")

        # ── Scaffold ─────────────────────────────────────────────────────
        if not args.generate_only and not args.dry_run:
            try:
                created = _scaffold_system(system_dir, overrides)
                if created:
                    total_scaffolded += len(created)
                    for bid in created:
                        print(f"      scaffolded {bid}")
            except Exception as e:
                tb = traceback.format_exc()
                print(f"    SCAFFOLD ERROR: {e}")
                _log_error(system_id, "*", str(e), tb)

        # ── Validate / Generate ──────────────────────────────────────────
        bodies_dir = system_dir / "bodies"
        if not bodies_dir.exists():
            if not args.scaffold_only and not args.dry_run:
                # Try scaffold first if bodies dir doesn't exist
                try:
                    _scaffold_system(system_dir, overrides)
                except Exception:
                    pass
            if not bodies_dir.exists():
                continue

        body_dirs = sorted(d for d in bodies_dir.iterdir() if d.is_dir())
        if args.body:
            body_dirs = [d for d in body_dirs if d.name == args.body]

        for body_dir in body_dirs:
            index_md = body_dir / "index.md"
            if not index_md.exists():
                continue

            bd = _read_frontmatter(index_md)
            if not bd or "id" not in bd:
                continue

            body_id = bd.get("id", "?")
            body_name = bd.get("name") or body_id

            if args.dry_run:
                # Validate only
                errors = _validate_body_def(bd, body_dir)
                if errors:
                    total_invalid += 1
                    print(f"      {body_id:20s} ({body_name:20s}) INVALID")
                    for err in errors:
                        print(f"        - {err}")
                else:
                    total_valid += 1
                    print(f"      {body_id:20s} ({body_name:20s}) ok")
                continue

            if args.scaffold_only:
                continue

            # Generate with error handling
            total_attempted += 1
            try:
                result = _generate_body_from_dir(
                    body_dir, hmap_w, hmap_h, args.globe_size, args.force)
                if result == "generated":
                    total_generated += 1
                elif result == "skipped":
                    total_skipped += 1
                elif result == "error":
                    total_errors += 1
            except Exception as e:
                total_errors += 1
                tb = traceback.format_exc()
                print(f"      {body_id:20s} ({body_name:20s}) ERROR: {e}")
                _log_error(system_id, body_id, str(e), tb)

            # Circuit breaker: abort if error rate is too high
            if total_attempted >= 10 and total_errors / total_attempted > error_rate_threshold:
                print(f"\n  ABORT: error rate {total_errors}/{total_attempted} "
                      f"({total_errors/total_attempted*100:.0f}%) exceeds "
                      f"{error_rate_threshold*100:.0f}% threshold")
                print(f"  Check {LOG_PATH} for details")
                sys.exit(1)

    elapsed = time.time() - t_total

    print(f"\n  Batch complete: {elapsed:.0f}s")
    if args.dry_run:
        print(f"    valid:      {total_valid}")
        print(f"    invalid:    {total_invalid}")
    else:
        print(f"    scaffolded: {total_scaffolded}")
        print(f"    generated:  {total_generated}")
        print(f"    skipped:    {total_skipped}")
        print(f"    errors:     {total_errors}")
    if total_errors > 0:
        print(f"    error log:  {LOG_PATH}")

    # ── Determinism verification ─────────────────────────────────────────
    if args.verify_determinism > 0:
        _verify_determinism(wiki_systems, args.verify_determinism,
                            hmap_w, hmap_h, args.globe_size)


def _verify_determinism(wiki_systems: Path, n_samples: int,
                        hmap_w: int, hmap_h: int, globe_size: int):
    """Run N bodies twice, verify outputs are bit-identical."""
    import hashlib
    import tempfile
    import shutil

    print(f"\n  Determinism verification — {n_samples} samples")

    # Collect all body dirs that have been generated
    all_body_dirs = []
    for system_dir in wiki_systems.iterdir():
        bodies = system_dir / "bodies"
        if bodies.exists():
            for bd in bodies.iterdir():
                if bd.is_dir() and (bd / "globe.png").exists():
                    all_body_dirs.append(bd)

    if not all_body_dirs:
        print("    no generated bodies to verify")
        return

    rng = np.random.default_rng(42)
    samples = rng.choice(len(all_body_dirs), min(n_samples, len(all_body_dirs)),
                         replace=False)

    passed = 0
    failed = 0
    for idx in samples:
        body_dir = all_body_dirs[idx]
        bd = _read_frontmatter(body_dir / "index.md")
        if not bd:
            continue
        body_id = bd["id"]

        # Generate into a temp dir
        tmp_dir = Path(tempfile.mkdtemp(prefix=f"detcheck_{body_id}_"))
        tmp_body = tmp_dir / body_id
        tmp_body.mkdir()

        # Copy index.md so the generator can read it
        shutil.copy2(body_dir / "index.md", tmp_body / "index.md")

        try:
            _generate_body_from_dir(tmp_body, hmap_w, hmap_h, globe_size, force=True)
        except Exception as e:
            print(f"    {body_id}: generation failed — {e}")
            shutil.rmtree(tmp_dir)
            failed += 1
            continue

        # Compare file hashes
        all_match = True
        for fname in ("globe.png", "heightmap.png", "terrain.npz", "markers.json"):
            orig = body_dir / fname
            rerun = tmp_body / fname
            if not orig.exists() and not rerun.exists():
                continue
            if not orig.exists() or not rerun.exists():
                print(f"    {body_id}: {fname} — missing in {'original' if not orig.exists() else 'rerun'}")
                all_match = False
                continue
            h1 = hashlib.sha256(orig.read_bytes()).hexdigest()[:16]
            h2 = hashlib.sha256(rerun.read_bytes()).hexdigest()[:16]
            if h1 != h2:
                print(f"    {body_id}: {fname} — MISMATCH (orig={h1} rerun={h2})")
                all_match = False

        if all_match:
            passed += 1
        else:
            failed += 1

        shutil.rmtree(tmp_dir)

    print(f"    passed: {passed}  failed: {failed}")
    if failed > 0:
        print(f"    WARNING: non-deterministic output detected!")


if __name__ == "__main__":
    main()