#!/usr/bin/env python3
"""
Planetary Screenshot Generator

Produces procedurally-rendered sphere images for each planet type
in the Settled Reach. Output used in:
  - wiki body-info-panel (360×360 container, sphere ~240px)
  - GTTR arrival window (diegetic implant UI, same asset)

Approach: numpy ray-sphere intersection + Lambertian shading + specular
+ procedural texture (octave-summed sine waves approximating noise).

No GPU required. Runs offline. Output is PNG at 512×512 (downscaled
to 360×360 for the panel; kept large for quality).

Usage:
    python3 generate_planets.py [--output-dir path]

Output:
    planet_temperate.png           — Earth-like: continents, ocean, clouds
    planet_temperate_terminator.png — Tidally locked: bright stripe, dark back
    planet_oceanic.png             — Water world: blue, archipelago dots
    planet_arid.png                — Desert/Mars: reddish-orange, dust storms
    planet_frozen.png              — Ice world: white, blue cracks
    planet_volcanic.png            — Volcanic: dark basalt, orange lava
    planet_barren.png              — Airless rocky: cratered grey
"""

import argparse
import math
import os
import numpy as np
from PIL import Image


SIZE = 512          # output pixel dimensions (square)
SPHERE_R = 0.92     # sphere radius in [-1, 1] space (slightly smaller than 1 = padding)


# ─────────────────────────────────────────────────────────────────────────────
# Procedural noise — octave sum of sine waves
# (approximates value noise without a noise library)
# ─────────────────────────────────────────────────────────────────────────────

def proc_noise(u: np.ndarray, v: np.ndarray, seed: float, octaves: int = 5) -> np.ndarray:
    """
    Returns values roughly in [-1, 1].
    u, v are 2D arrays of coordinates (e.g. lon/lat on sphere surface).
    seed shifts the pattern.
    """
    result = np.zeros_like(u, dtype=np.float32)
    amplitude = 1.0
    frequency = 1.0
    total_amp = 0.0
    for i in range(octaves):
        ph = seed + i * 3.7
        result += amplitude * (
            np.sin(frequency * u * 13.7 + ph) * np.cos(frequency * v * 8.1 + ph * 0.7) +
            np.cos(frequency * u * 7.3 - ph * 0.4) * np.sin(frequency * v * 11.3 + ph * 1.3)
        )
        total_amp += 2 * amplitude
        amplitude *= 0.5
        frequency *= 2.0
    return result / total_amp


# ─────────────────────────────────────────────────────────────────────────────
# Ray-sphere intersection
# ─────────────────────────────────────────────────────────────────────────────

def raytrace_sphere(size: int, r: float = 1.0):
    """
    Returns (hit_mask, nx, ny, nz, u, v) arrays of shape (size, size).
    - hit_mask: bool, True where the ray hits the sphere
    - nx/ny/nz: surface normals at hit points (unit length)
    - u/v: spherical UV coordinates in [0, 1]

    Camera at (0, 0, 3), looking toward origin. Sphere at origin.
    """
    # Pixel coordinates mapped to [-1, 1] square
    lin = np.linspace(-1, 1, size, dtype=np.float32)
    px, py = np.meshgrid(lin, -lin)   # y flipped so top = 1

    # Ray direction from camera
    oz = 3.0
    rdx, rdy, rdz = px, py, -oz * np.ones((size, size), dtype=np.float32)
    mag = np.sqrt(rdx**2 + rdy**2 + rdz**2)
    rdx /= mag; rdy /= mag; rdz /= mag

    # Ray-sphere: t^2 + 2t(o·d) + |o|^2 - r^2 = 0
    # o = (0, 0, oz), sphere center = (0, 0, 0)
    b = 2 * (oz * rdz)   # ox=oy=0
    c = oz**2 - r**2
    disc = b**2 - 4 * c
    hit = disc >= 0.0

    t_arr = np.where(hit, (-b - np.sqrt(np.maximum(disc, 0.0))) / 2.0, np.inf)

    # Hit position
    hx = rdx * t_arr
    hy = rdy * t_arr
    hz = oz + rdz * t_arr

    # Normals (outward) — normalise hit position since sphere at origin radius r
    norm = np.sqrt(hx**2 + hy**2 + hz**2)
    norm = np.where(hit, norm, 1.0)   # avoid /0 in miss pixels
    nx, ny, nz = hx / norm, hy / norm, hz / norm

    # Spherical UV: u = lon / 2π, v = lat / π + 0.5
    u_coord = (np.arctan2(nz, nx) / (2 * math.pi)) % 1.0
    v_coord = np.arcsin(np.clip(ny, -1, 1)) / math.pi + 0.5

    return hit, nx, ny, nz, u_coord, v_coord


# ─────────────────────────────────────────────────────────────────────────────
# Lighting
# ─────────────────────────────────────────────────────────────────────────────

def light(nx, ny, nz, lx=-0.6, ly=0.5, lz=-0.4, shininess=20.0):
    """
    Diffuse + specular from a fixed star direction.
    Returns diffuse (float array) and specular (float array).
    """
    lmag = math.sqrt(lx**2 + ly**2 + lz**2)
    lx /= lmag; ly /= lmag; lz /= lmag

    diffuse = np.clip(nx * lx + ny * ly + nz * lz, 0.0, 1.0)

    # Reflect: r = 2(n·l)n - l
    ndotl = nx * lx + ny * ly + nz * lz
    rx = 2 * ndotl * nx - lx
    ry = 2 * ndotl * ny - ly
    rz = 2 * ndotl * nz - lz

    # View direction: toward camera at (0,0,3) — for normalized normals ~(0,0,1) approx
    vz = 1.0
    spec = np.clip(rx * 0 + ry * 0 + rz * vz, 0.0, 1.0) ** shininess
    return diffuse, spec


def lerp_color(a, b, t):
    """Blend between two RGB tuples by t (0-1 scalar or array)."""
    t = np.clip(t, 0, 1)
    if hasattr(t, '__len__'):
        t = t[..., np.newaxis]
    return np.array(a) * (1 - t) + np.array(b) * t


# ─────────────────────────────────────────────────────────────────────────────
# Planet type renderers
# ─────────────────────────────────────────────────────────────────────────────

def render_planet(hit, nx, ny, nz, u, v, surface_fn, cloud_fn=None,
                  bg=(4, 6, 10), star_lx=-0.55, star_ly=0.45, star_lz=0.70,
                  has_atmosphere=True):
    """
    Core render loop. surface_fn(u, v, nx, ny, nz) → RGB float [0..1].
    cloud_fn(u, v) → alpha float [0..1] or None.
    Returns PIL Image (RGBA).
    """
    H, W = hit.shape
    rgb = np.zeros((H, W, 4), dtype=np.float32)

    # Background (space)
    rgb[..., 0] = bg[0] / 255.0
    rgb[..., 1] = bg[1] / 255.0
    rgb[..., 2] = bg[2] / 255.0
    rgb[..., 3] = 1.0

    # Star field
    rng = np.random.default_rng(42)
    star_mask = rng.random((H, W)) < 0.002
    star_bright = rng.uniform(0.4, 1.0, (H, W))
    rgb[~hit & star_mask, 0] = star_bright[~hit & star_mask]
    rgb[~hit & star_mask, 1] = star_bright[~hit & star_mask]
    rgb[~hit & star_mask, 2] = star_bright[~hit & star_mask]

    if not hit.any():
        arr = (np.clip(rgb, 0, 1) * 255).astype(np.uint8)
        return Image.fromarray(arr)

    # Work only on hit pixels to avoid NaN propagation from miss areas
    h_idx = np.where(hit)
    nx_h = nx[h_idx]; ny_h = ny[h_idx]; nz_h = nz[h_idx]
    u_h = u[h_idx]; v_h = v[h_idx]

    # Surface color (compute on hit pixels)
    surf_full = surface_fn(u, v, nx, ny, nz)   # full grid for simplicity
    surf_h = surf_full[h_idx]   # (N, 3)

    # Lighting on hit pixels only
    lmag = math.sqrt(star_lx**2 + star_ly**2 + star_lz**2)
    lx = star_lx / lmag; ly = star_ly / lmag; lz = star_lz / lmag

    diff_h = np.clip(nx_h * lx + ny_h * ly + nz_h * lz, 0.0, 1.0)

    ndotl = nx_h * lx + ny_h * ly + nz_h * lz
    rx_h = 2 * ndotl * nx_h - lx
    ry_h = 2 * ndotl * ny_h - ly
    rz_h = 2 * ndotl * nz_h - lz
    spec_h = np.clip(rz_h, 0.0, 1.0) ** 25  # view dir = +z approx

    # Ambient + diffuse + spec — generous ambient for readability (dark side not black)
    ambient = 0.22
    lit_h = surf_h * (ambient + 0.78 * diff_h[:, np.newaxis]) + spec_h[:, np.newaxis] * 0.30

    # Cloud layer
    if cloud_fn is not None:
        cloud_full = cloud_fn(u, v)
        cloud_h = cloud_full[h_idx]
        cloud_alpha_h = np.clip(cloud_h * 0.9, 0, 1)
        cloud_rgb_h = np.ones((len(h_idx[0]), 3)) * 0.96
        cloud_lit_h = cloud_rgb_h * (ambient + 0.85 * diff_h[:, np.newaxis])
        lit_h = lit_h * (1 - cloud_alpha_h[:, np.newaxis]) + cloud_lit_h * cloud_alpha_h[:, np.newaxis]

    # Atmosphere rim glow (blue edge)
    if has_atmosphere:
        rim_h = np.abs(nz_h)   # dot with view +z
        rim_glow_h = (1 - rim_h) ** 5 * 0.7
        atmo_c = np.array([0.35, 0.60, 1.0])
        lit_h = lit_h + rim_glow_h[:, np.newaxis] * atmo_c * 0.35

    lit_h = np.clip(lit_h, 0, 1)

    out = np.zeros((H * W, 3), dtype=np.float32)
    out[np.ravel_multi_index(h_idx, (H, W))] = lit_h
    out = out.reshape(H, W, 3)

    rgb[hit, 0] = out[hit, 0]
    rgb[hit, 1] = out[hit, 1]
    rgb[hit, 2] = out[hit, 2]
    rgb[hit, 3] = 1.0

    arr = (np.clip(rgb, 0, 1) * 255).astype(np.uint8)
    return Image.fromarray(arr)


# ─────────────────────────────────────────────────────────────────────────────
# Planet type definitions
# ─────────────────────────────────────────────────────────────────────────────

def planet_temperate(out_dir):
    """
    Temperate / Earth-like. Continents (amber-brown), ocean (blue), polar caps.
    Character: "amber continental shelves" per Kallast wiki.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 2, v * 3, seed=1.1)
    n2 = proc_noise(u * 5, v * 7, seed=2.3, octaves=3)

    def surface(u, v, nx, ny, nz):
        h = n * 0.7 + n2 * 0.3
        lat = (v - 0.5) * 2   # [-1, 1]
        polar = np.abs(lat) ** 3
        land = h > 0.08
        # Ocean: rich blue, deepens with depth
        ocean_t = np.clip((h + 0.6) * 0.9, 0, 1)
        ocean_c = lerp_color((8, 28, 80), (35, 100, 175), ocean_t) / 255.0
        # Land: amber grain belt → hillside → highland
        land_t = np.clip((h - 0.08) / 0.55, 0, 1)
        land_c = lerp_color((155, 165, 75), (105, 90, 65), land_t) / 255.0
        # Forest band at mid elevation
        forest_band = np.clip((land_t - 0.2) * 5, 0, 1) * np.clip((0.6 - land_t) * 5, 0, 1)
        forest_c = np.array([55, 100, 50]) / 255.0
        land_c = land_c * (1 - forest_band[..., np.newaxis]) + forest_c * forest_band[..., np.newaxis]
        base = np.where(land[..., np.newaxis], land_c, ocean_c)
        # Polar ice caps
        ice_c = np.array([0.88, 0.91, 0.97])
        polar_blend = np.clip((polar - 0.55) * 5, 0, 1)
        base = base * (1 - polar_blend[..., np.newaxis]) + ice_c * polar_blend[..., np.newaxis]
        return base

    def clouds(u, v):
        c1 = proc_noise(u * 3, v * 2, seed=9.1, octaves=4)
        return np.clip((c1 + 0.2) * 1.5, 0, 1) * 0.5

    img = render_planet(hit, nx, ny, nz, u, v, surface, clouds)
    img.save(os.path.join(out_dir, "planet_temperate.png"))
    print(f"  planet_temperate.png")


def planet_temperate_terminator(out_dir):
    """
    Tidally locked. One face scorched, one face frozen, habitable terminator band.
    Character: "terminator-band settlement" per Feldmark, Caparica wikis.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 2, v * 2, seed=3.3)

    def surface(u, v, nx, ny, nz):
        # Longitude position: 0=day face, 0.5=night face
        lon = (u - 0.25) % 1.0   # shift so day face is center
        day = np.clip(1 - lon * 2, 0, 1)
        night = np.clip(lon * 2 - 1, 0, 1)
        term = 1 - day - night
        h = n * 0.5

        day_c = np.array([0.80, 0.55, 0.25])    # scorched orange-gold
        night_c = np.array([0.15, 0.20, 0.30])  # frozen dark blue
        term_c = lerp_color((85, 130, 90), (130, 160, 100), np.clip(h, 0, 1)) / 255.0

        base = (day_c * day[..., np.newaxis] +
                night_c * night[..., np.newaxis] +
                term_c * term[..., np.newaxis])
        return np.clip(base, 0, 1)

    def clouds(u, v):
        c = proc_noise(u * 2.5, v * 4, seed=11.1, octaves=3)
        lon = (u - 0.25) % 1.0
        term_weight = np.clip(1 - np.abs(lon - 0.5) * 4, 0, 1)
        return np.clip((c + 0.3) * 0.6, 0, 1) * term_weight

    img = render_planet(hit, nx, ny, nz, u, v, surface, clouds)
    img.save(os.path.join(out_dir, "planet_temperate_terminator.png"))
    print(f"  planet_temperate_terminator.png")


def planet_oceanic(out_dir):
    """
    Ocean world. Mostly water, scattered archipelagos.
    Character: Caparica — aquaculture domes visible, terminator farming.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 4, v * 5, seed=5.5, octaves=4)

    def surface(u, v, nx, ny, nz):
        h = n
        land = h > 0.55   # very little land — archipelago only
        ocean_deep_c = np.array([10, 30, 70]) / 255.0
        ocean_shallow_c = np.array([30, 80, 130]) / 255.0
        ocean_t = np.clip((h + 0.5) * 0.8, 0, 1)
        ocean_c = ocean_deep_c * (1 - ocean_t[..., np.newaxis]) + ocean_shallow_c * ocean_t[..., np.newaxis]
        land_c = np.array([80, 120, 70]) / 255.0
        base = np.where(land[..., np.newaxis], land_c, ocean_c)
        # Polar ice
        lat = np.abs(v - 0.5) * 2
        ice = np.clip((lat - 0.7) * 5, 0, 1)
        base = base * (1 - ice[..., np.newaxis]) + np.array([0.85, 0.90, 0.95]) * ice[..., np.newaxis]
        return base

    def clouds(u, v):
        c = proc_noise(u * 2, v * 3, seed=22.2, octaves=4)
        return np.clip((c + 0.3) * 0.7, 0, 1) * 0.7

    img = render_planet(hit, nx, ny, nz, u, v, surface, clouds)
    img.save(os.path.join(out_dir, "planet_oceanic.png"))
    print(f"  planet_oceanic.png")


def planet_arid(out_dir):
    """
    Arid / desert / Mars-analog. Reddish-orange dust, no permanent surface water.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 2, v * 3, seed=7.7, octaves=4)
    n2 = proc_noise(u * 8, v * 10, seed=8.2, octaves=2)

    def surface(u, v, nx, ny, nz):
        h = n * 0.6 + n2 * 0.4
        base_t = np.clip((h + 0.5) * 0.9, 0, 1)
        dark_c = np.array([100, 50, 35]) / 255.0
        light_c = np.array([190, 130, 80]) / 255.0
        base = dark_c * (1 - base_t[..., np.newaxis]) + light_c * base_t[..., np.newaxis]
        # Dust storm wisps
        storm = proc_noise(u * 6, v * 2, seed=13.1, octaves=3)
        storm_alpha = np.clip((storm + 0.6) * 0.3, 0, 1)
        dust_c = np.array([0.78, 0.60, 0.45])
        base = base * (1 - storm_alpha[..., np.newaxis]) + dust_c * storm_alpha[..., np.newaxis]
        # Thin polar cap
        lat = np.abs(v - 0.5) * 2
        ice = np.clip((lat - 0.85) * 8, 0, 1)
        base = base * (1 - ice[..., np.newaxis]) + np.array([0.90, 0.88, 0.88]) * ice[..., np.newaxis]
        return np.clip(base, 0, 1)

    # No cloud layer (thin atmosphere)
    img = render_planet(hit, nx, ny, nz, u, v, surface, cloud_fn=None)
    img.save(os.path.join(out_dir, "planet_arid.png"))
    print(f"  planet_arid.png")


def planet_frozen(out_dir):
    """
    Ice world. White/blue, glacial features, ice caps extend to equator.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 3, v * 4, seed=11.1, octaves=5)
    n2 = proc_noise(u * 7, v * 9, seed=12.3, octaves=3)

    def surface(u, v, nx, ny, nz):
        h = n * 0.6 + n2 * 0.4
        # Ice everywhere — variation between white and blue-grey
        base_t = np.clip((h + 0.3) * 0.9, 0, 1)
        deep_ice = np.array([0.55, 0.65, 0.80])
        snow_c = np.array([0.90, 0.92, 0.96])
        base = deep_ice * (1 - base_t[..., np.newaxis]) + snow_c * base_t[..., np.newaxis]
        # Exposed rock patches at mid-latitude
        lat = np.abs(v - 0.5) * 2
        rock_zone = np.clip((0.4 - lat) * 3, 0, 1) * np.clip(n2 + 0.1, 0, 1)
        rock_c = np.array([0.35, 0.32, 0.30])
        base = base * (1 - rock_zone[..., np.newaxis]) + rock_c * rock_zone[..., np.newaxis]
        return base

    def clouds(u, v):
        c = proc_noise(u * 2, v * 2, seed=30.0, octaves=3)
        return np.clip((c + 0.4) * 0.4, 0, 1) * 0.3

    img = render_planet(hit, nx, ny, nz, u, v, surface, clouds)
    img.save(os.path.join(out_dir, "planet_frozen.png"))
    print(f"  planet_frozen.png")


def planet_volcanic(out_dir):
    """
    Active volcanic world. Dark basalt surface with orange/red lava flows.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 3, v * 4, seed=15.5, octaves=4)
    n2 = proc_noise(u * 10, v * 12, seed=16.7, octaves=3)

    def surface(u, v, nx, ny, nz):
        h = n * 0.7 + n2 * 0.3
        # Dark basalt base
        basalt_t = np.clip((h + 0.5) * 0.6, 0, 1)
        dark_c = np.array([25, 22, 20]) / 255.0
        mid_c = np.array([65, 55, 50]) / 255.0
        base = dark_c * (1 - basalt_t[..., np.newaxis]) + mid_c * basalt_t[..., np.newaxis]
        # Lava flows: bright orange channels in low-elevation areas
        lava_t = np.clip((-h - 0.1) * 3, 0, 1)
        lava_c = np.array([0.95, 0.45, 0.10])
        base = base + lava_c * lava_t[..., np.newaxis]
        # Volcano glow spots
        g = proc_noise(u * 15, v * 15, seed=17.3, octaves=2)
        glow = np.clip((g - 0.7) * 5, 0, 1)
        glow_c = np.array([1.0, 0.6, 0.2])
        base = base + glow_c * glow[..., np.newaxis] * 0.5
        return np.clip(base, 0, 1)

    def clouds(u, v):
        # Volcanic haze — yellowish-brown sulfur clouds
        c = proc_noise(u * 4, v * 3, seed=55.0, octaves=3)
        return np.clip((c + 0.5) * 0.6, 0, 1) * 0.4

    # Warmer star (volcanic world might be inner orbit — stronger illumination)
    img = render_planet(hit, nx, ny, nz, u, v, surface, clouds,
                        star_lx=-0.55, star_ly=0.40, star_lz=0.73)
    img.save(os.path.join(out_dir, "planet_volcanic.png"))
    print(f"  planet_volcanic.png")


def planet_barren(out_dir):
    """
    Airless barren world. Cratered grey-brown, no atmosphere glow.
    """
    hit, nx, ny, nz, u, v = raytrace_sphere(SIZE, SPHERE_R)
    n = proc_noise(u * 3, v * 4, seed=20.0, octaves=4)
    n2 = proc_noise(u * 12, v * 14, seed=21.3, octaves=2)

    def surface(u, v, nx, ny, nz):
        h = n * 0.6 + n2 * 0.4
        t = np.clip((h + 0.5) * 0.8, 0, 1)
        dark_c = np.array([50, 45, 42]) / 255.0
        light_c = np.array([140, 130, 120]) / 255.0
        base = dark_c * (1 - t[..., np.newaxis]) + light_c * t[..., np.newaxis]
        return base

    # No clouds, no atmosphere glow — hard terminator
    img = render_planet(hit, nx, ny, nz, u, v, surface, cloud_fn=None,
                        has_atmosphere=False)
    img.save(os.path.join(out_dir, "planet_barren.png"))
    print(f"  planet_barren.png")


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

PLANET_TYPES = [
    ("temperate",              planet_temperate),
    ("temperate_terminator",   planet_temperate_terminator),
    ("oceanic",                planet_oceanic),
    ("arid",                   planet_arid),
    ("frozen",                 planet_frozen),
    ("volcanic",               planet_volcanic),
    ("barren",                 planet_barren),
]


def main():
    parser = argparse.ArgumentParser(description="Generate procedural planet screenshots")
    parser.add_argument("--output-dir", default="client/assets/planets",
                        help="Directory for output PNG files")
    parser.add_argument("--type", choices=[p[0] for p in PLANET_TYPES],
                        help="Render only one planet type")
    args = parser.parse_args()

    os.makedirs(args.output_dir, exist_ok=True)
    print(f"Output dir: {args.output_dir}")
    print(f"Rendering {SIZE}×{SIZE}px spheres…\n")

    targets = PLANET_TYPES if not args.type else [(t, fn) for t, fn in PLANET_TYPES if t == args.type]

    for ptype, fn in targets:
        fn(args.output_dir)

    print(f"\nDone. {len(targets)} planet type(s) written to {args.output_dir}/")
    print("For wiki/GTTR display: scale to 240×240 within the 360×360 panel container.")


if __name__ == "__main__":
    main()