"""Terrain export for games: freeze a square of the map, exactly as the viewer shows it, into engine-neutral files. Generated terrain."""
from __future__ import annotations

import json
import math
import multiprocessing
import os
import shutil
import tempfile
from datetime import date
from pathlib import Path

import numpy as np
from PIL import Image
from scipy.ndimage import map_coordinates

import worldgen_path  # noqa: F401  (mapgen on sys.path)
import refine
import rivers as RV
import serve
import tiles as T
from mapgen import render as RN
from mapgen.sphere import latlon_to_xyz

FORMAT_VERSION = 1
WATER = ["land", "sea", "lake", "river"]


def samples_for(size_km: float, res_m: float) -> int:
    """The smallest 2^n + 1 covering size_km at res_m (Unity, Godot Terrain3D and most tools take these sizes)."""
    need = size_km * 1000.0 / res_m + 1
    n = 1
    while 2 ** n + 1 < need:
        n += 1
    return 2 ** n + 1


def grid_latlon(lat0, lon0, n, res_m, radius_km, rows=None):
    """Latitude/longitude of the n × n samples (or of rows a:b): azimuthal equidistant about (lat0, lon0), x east,
    y north, row 0 north."""
    c = (n - 1) / 2
    x = (np.arange(n) - c) * res_m
    r = np.arange(n) if rows is None else np.arange(*rows)
    y = (c - r) * res_m
    X, Y = np.meshgrid(x, y)
    d = np.hypot(X, Y) / (radius_km * 1000.0)                  # angular distance from the centre
    b = np.arctan2(X, Y)                                        # bearing
    la0, lo0 = math.radians(lat0), math.radians(lon0)
    la = np.arcsin(np.sin(la0) * np.cos(d) + np.cos(la0) * np.sin(d) * np.cos(b))
    lo = lo0 + np.arctan2(np.sin(b) * np.sin(d) * np.cos(la0), np.cos(d) - np.sin(la0) * np.sin(la))
    return np.degrees(la), (np.degrees(lo) + 180.0) % 360.0 - 180.0


def to_local(lat0, lon0, lat, lon, radius_km):
    """Local x east, y north (m) of points (the inverse of grid_latlon)."""
    la0, lo0 = math.radians(lat0), math.radians(lon0)
    la, lo = np.radians(lat), np.radians(lon)
    cd = np.clip(np.sin(la0) * np.sin(la) + np.cos(la0) * np.cos(la) * np.cos(lo - lo0), -1, 1)
    d = np.arccos(cd)
    b = np.arctan2(np.sin(lo - lo0) * np.cos(la), np.cos(la0) * np.sin(la) - np.sin(la0) * np.cos(la) * np.cos(lo - lo0))
    r = d * radius_km * 1000.0
    return r * np.sin(b), r * np.cos(b)


_SRC = None                                                     # the tile source (forked workers inherit it)


def _tile(args):
    return tile_data(_SRC, *args)


def tile_data(src, z, x, y) -> dict:
    """One tile's export layers (256 × 256): heights, water classes, biome, ground, landform, unshaded colours."""
    f = src.fields(z, x, y, 1, ("z", "water", "cat"))
    c = lambda a: np.asarray(a)[1:-1, 1:-1]
    lake, river = c(f["lake"]).astype(bool), c(f["river"]).astype(bool)
    flat = np.ones(lake.shape)                                  # unshaded colours: the export shades its own grid
    rgb = RN.relief_rgb(c(f["z"]), flat, c(f["holdridge"]), c(f["ground"]), c(f["ice"]), lake | c(f["river_draw"]),
                        ~c(f["water"]))
    water = np.where(c(f["water"]), 1, np.where(lake, 2, np.where(river, 3, 0))).astype(np.uint8)
    return {"z": c(f["z"]).astype(np.float32), "water": water, "biome": c(f["holdridge"]).astype(np.uint8),
            "ground": c(f["ground"]).astype(np.uint8), "landform": c(f["landform"]).astype(np.uint8),
            "rgb": np.asarray(rgb, dtype=np.uint8)}


def _rivers(src, rs, lat0, lon0, half_m, res_m):
    """River centrelines inside the square (local metres), world rivers where the world shows, refined ones inside."""
    R = src.R
    centre = latlon_to_xyz(lat0, lon0).reshape(3)
    radius = half_m * math.sqrt(2) / 1000.0
    nets = [(src.river_net, False)] + ([(rs.net, True)] if not rs.empty and rs.net is not None else [])
    out = []
    for net, refined in nets:
        if net.tree is None:
            continue
        for s in net.candidates(centre, radius):
            p, _ = net.segment_points(s, centre, radius, res_m / 1000.0)
            if not len(p):
                continue
            lat = np.degrees(np.arcsin(np.clip(p[:, 2], -1, 1)))
            lon = np.degrees(np.arctan2(p[:, 1], p[:, 0]))
            x, y = to_local(lat0, lon0, lat, lon, R)
            keep = (np.abs(x) <= half_m) & (np.abs(y) <= half_m)
            if not rs.empty:
                w = rs.weight(p)
                keep &= (w >= 0.5) if refined else (w < 0.5)
            runs = np.split(np.arange(len(p)), np.where(~keep)[0])
            q = (net.half_w[s] / 0.004) ** 2 / 31.7                # back from the hydraulic width (rivers.half_width_km)
            for run in runs:
                run = run[keep[run]]
                if len(run) >= 2:
                    out.append({"points": [[round(float(a), 1), round(float(b), 1)] for a, b in zip(x[run], y[run])],
                                "width_m": round(float(2000.0 * net.half_w[s]), 1), "discharge_km3_yr": round(float(q), 3),
                                "refined": refined})
    out.sort(key=lambda r: (r["points"][0], r["width_m"]))
    return out


def _hillshade(h, res_m, exag):
    gy, gx = np.gradient(np.asarray(h, dtype=np.float32), np.float32(res_m))
    dzdx, dzdn = gx * exag, -gy * exag                          # rows run north → south
    a, b = np.radians(315.0), np.radians(45.0)
    L = (np.sin(a) * np.cos(b), np.cos(a) * np.cos(b), np.sin(b))
    return np.clip((-dzdx * L[0] - dzdn * L[1] + L[2]) / np.sqrt(dzdx ** 2 + dzdn ** 2 + 1.0), 0.0, 1.0)


def _hillshade_exaggeration(h, res_m):
    """Vertical exaggeration for the preview: gentle land steepened until its steeper slopes read (≤ 50×)."""
    step = max(1, h.shape[0] // 1024)                           # a sample of the slopes is enough
    hh = np.asarray(h[::step, ::step], dtype=np.float64)
    gy, gx = np.gradient(hh, res_m * step)
    p95 = float(np.percentile(np.hypot(gx, gy), 95))
    return float(np.clip(0.35 / max(p95, 1e-6), 1.0, 50.0))


MAX_TILES = 1600        # ≈ an 80 km square at 10 m near the equator (≈ 1 GB of tile data): more is refused
POLE_LIMIT_DEG = 85.0   # squares reaching closer to a pole are refused (tiles narrow toward the poles)


def plan(lat, lon, size_km, res_m, radius_km, pixels=None) -> dict:
    """Samples, tile zoom and the tile rectangle an export needs (ValueError near a pole)."""
    n = int(pixels) if pixels else samples_for(size_km, res_m)
    half = (n - 1) / 2 * res_m
    reach = math.degrees(half * math.sqrt(2) / (radius_km * 1000.0))
    if abs(lat) + reach > POLE_LIMIT_DEG:
        raise ValueError(f"the square reaches within {90 - POLE_LIMIT_DEG:g}° of a pole: move it or make it smaller")
    # the tile zoom whose pixels are at least as fine as res_m (north-south; east-west pixels are finer still)
    z = int(min(T.MAX_Z, max(5, math.ceil(math.log2(math.pi * radius_km * 1000.0 / (T.TILE * res_m))))))
    span = 180.0 / 2 ** z
    edge = np.r_[0, n - 1]                                      # the square's extremes lie on its edges
    la1, lo1 = grid_latlon(lat, lon, n, res_m, radius_km, rows=(0, 1))
    la2, lo2 = grid_latlon(lat, lon, n, res_m, radius_km, rows=(n - 1, n))
    la3, lo3 = (np.concatenate(v) for v in zip(*(grid_latlon(lat, lon, n, res_m, radius_km, rows=(k, k + 1))
                                                  for k in range(0, n, max(1, (n - 1) // 64)))))
    la = np.concatenate([la1.ravel(), la2.ravel(), la3[:, edge].ravel()])
    lo = np.concatenate([lo1.ravel(), lo2.ravel(), lo3[:, edge].ravel()])
    lonu = lon + ((lo - lon + 180.0) % 360.0 - 180.0)
    fx, fy = (lonu + 180.0) / span, (90.0 - la) / span
    x0, x1 = int(math.floor(fx.min())) - 1, int(math.floor(fx.max())) + 1
    y0, y1 = max(0, int(math.floor(fy.min())) - 1), min(2 ** z - 1, int(math.floor(fy.max())) + 1)
    return {"n": n, "z": z, "span": span, "x0": x0, "x1": x1, "y0": y0, "y1": y1,
            "tiles": (x1 - x0 + 1) * (y1 - y0 + 1)}


def _parallel(jobs, fn, threads, put, stop):
    """fn(job) on daemon threads (never delaying the server's exit), put(job, result) in turn; the first error wins."""
    import threading
    it, lock, err = iter(jobs), threading.Lock(), []

    def run():
        while not err:
            with lock:
                job = next(it, None)
            if job is None:
                return
            try:
                r = fn(job)
                with lock:
                    put(job, r)
                stop()
            except BaseException as e:                          # noqa: BLE001 — re-raised below
                err.append(e)
    ts = [threading.Thread(target=run, daemon=True) for _ in range(max(1, threads))]
    for t in ts:
        t.start()
    for t in ts:
        t.join()
    if err:
        raise err[0]


def export(root: Path, res: int, lat: float, lon: float, size_km: float = 40.0, res_m: float = 10.0,
           out_dir: Path | None = None, name: str | None = None, pixels: int | None = None,
           pins_path: Path | None = None, regions_dir: Path | None = None, workers: int = 0, progress=None,
           src=None, max_tiles: int = MAX_TILES, era: str | None = None) -> Path:
    """Write the export folder and return it. progress(done, total) per tile. src: a running server's tile source
    (its world, regions and render workers are used; nothing is loaded again). RuntimeError if the refined regions
    change meanwhile (the export would mix two maps). era: a built era (default: [eras] default, else the base);
    a server's src already is its chosen era's."""
    global _SRC
    from mapgen import config as C
    root = Path(root)
    cfg = C.load(root)[0]
    out_dir = Path(out_dir or (root / "exports"))
    name = name or default_name(lat, lon, size_km, res_m)
    dest = out_dir / name
    if dest.exists():
        raise FileExistsError(f"export {dest} exists: pick another name or remove it")
    own = src is None
    if own:
        import refine
        built = [n for n, _ in refine.world_dirs(root, res, log=lambda m: None)] if (
            root / "out" / f"r{res}" / "cells.npz").exists() else ["base"]
        default = (C.load(root)[1].get("eras") or {}).get("default")
        name_era = era or (default if default in built else built[0])
        if name_era not in built:
            raise SystemExit(f"no built era {name_era!r} at res {res} (have: {built})")
        world = dict(serve.load_worlds(root, res, log=lambda m: None, only=name_era))[name_era]
        src = T.TileSource(world, int(cfg["build"]["seed"]),
                           cache_dir=Path(tempfile.mkdtemp(prefix="export-tiles-")),
                           regions_dir=regions_dir or (root / "out" / f"r{res}" / "regions"))
    try:
        return _export(src, own, root, res, lat, lon, size_km, res_m, out_dir, name, dest, pixels, pins_path,
                       workers, progress, max_tiles, cfg)
    finally:
        if own:
            shutil.rmtree(src.cache_root, ignore_errors=True)


def _export(src, own, root, res, lat, lon, size_km, res_m, out_dir, name, dest, pixels, pins_path, workers,
            progress, max_tiles, cfg):
    global _SRC
    world, R = src.w, src.R
    P = plan(lat, lon, size_km, res_m, R, pixels)
    if P["tiles"] > max_tiles:
        raise ValueError(f"{P['tiles']} tiles is too many (≤ {max_tiles}): a coarser resolution or a smaller square")
    n, z, x0, x1, y0, y1 = P["n"], P["z"], P["x0"], P["x1"], P["y0"], P["y1"]
    half = (n - 1) / 2 * res_m
    rs = src.regions                                            # one region set for the whole export
    changed = lambda: src.regions is not rs
    W, H = (x1 - x0 + 1), (y1 - y0 + 1)
    jobs = [(z, xx % 2 ** (z + 1), yy) for yy in range(y0, y1 + 1) for xx in range(x0, x1 + 1)]
    mosaic = {"z": np.zeros((H * T.TILE, W * T.TILE), np.float32), "rgb": np.zeros((H * T.TILE, W * T.TILE, 3), np.uint8),
              **{k: np.zeros((H * T.TILE, W * T.TILE), np.uint8) for k in ("water", "biome", "ground", "landform")}}
    done = [0]

    def put(job, part):
        zz, xx, yy = job
        r, c = (yy - y0) * T.TILE, ((xx - x0) % 2 ** (z + 1)) * T.TILE
        for k, v in part.items():
            mosaic[k][r:r + T.TILE, c:c + T.TILE] = v
        done[0] += 1
        if progress:
            progress(done[0], len(jobs))

    def check():
        if changed():
            raise RuntimeError("the refined regions changed during the export (a region build finished): export again")
    src.tree()
    src.river_net
    src.raster("elevation")
    src.regions.warm()                                          # region search trees too: forked workers share them
    if not own:                                                 # a server: several render workers at once
        pool = src.pool
        _parallel(jobs, lambda j: src.export_tile(*j), max(1, min(4, pool.alive // 2)) if pool else 1, put, check)
    elif workers and workers > 0:
        _SRC = src
        with multiprocessing.get_context("fork").Pool(workers) as mp:
            for job, part in zip(jobs, mp.imap(_tile, jobs, chunksize=1)):
                put(job, part)
    else:
        for job in jobs:
            put(job, tile_data(src, *job))
    check()
    # sample the square in row blocks (memory: the mosaic and the outputs, not n² float64 grids)
    span = P["span"]
    h = np.empty((n, n), np.float32)
    cls = {k: np.empty((n, n), np.uint8) for k in ("water", "biome", "ground", "landform")}
    rgb = np.empty((n, n, 3), np.uint8)
    for r0 in range(0, n, 256):
        r1 = min(n, r0 + 256)
        la, lo = grid_latlon(lat, lon, n, res_m, R, rows=(r0, r1))
        lonu = lon + ((lo - lon + 180.0) % 360.0 - 180.0)
        pr = ((90.0 - la) / span - y0) * T.TILE - 0.5            # mosaic pixel positions (pixel centres)
        pc = ((lonu + 180.0) / span - x0) * T.TILE - 0.5
        h[r0:r1] = map_coordinates(mosaic["z"], [pr, pc], order=1, mode="nearest")
        ri = np.clip(np.rint(pr).astype(np.int32), 0, H * T.TILE - 1)
        ci = np.clip(np.rint(pc).astype(np.int32), 0, W * T.TILE - 1)
        for k in cls:
            cls[k][r0:r1] = mosaic[k][ri, ci]
        rgb[r0:r1] = mosaic["rgb"][ri, ci]
    del mosaic
    out_dir.mkdir(parents=True, exist_ok=True)
    tmp = Path(tempfile.mkdtemp(prefix=f".{name}-", dir=out_dir))   # same file system: the final rename is atomic
    try:
        h.astype("<f4").tofile(tmp / "height.f32")
        lo_, hi_ = float(h.min()), float(h.max())
        v = np.rint((h - lo_) / max(hi_ - lo_, 1e-9) * 65535).astype(np.uint16)
        Image.fromarray(v).save(tmp / "height.png")
        v.astype("<u2").tofile(tmp / "height.r16")
        del v
        for k, f in (("water", "water.png"), ("biome", "biome.png"), ("ground", "ground.png"), ("landform", "landform.png")):
            Image.fromarray(cls[k]).save(tmp / f)
        exag = _hillshade_exaggeration(h, res_m)
        hs = _hillshade(h, res_m, exag)
        land = cls["water"] != 1
        for r0 in range(0, n, 512):                             # shade the preview in blocks (float32)
            r1 = min(n, r0 + 512)
            sh = np.where(land[r0:r1], 0.55 + 0.45 * hs[r0:r1], 0.85 + 0.15 * hs[r0:r1]).astype(np.float32)
            rgb[r0:r1] = np.clip(rgb[r0:r1] * sh[..., None], 0, 255).astype(np.uint8)
        del hs
        Image.fromarray(rgb).save(tmp / "preview.png")
        (tmp / "rivers.json").write_text(json.dumps(_rivers(src, rs, lat, lon, half, res_m)))
        pins = []
        pp = Path(pins_path) if pins_path else root / "places" / "pins.json"
        for p in (json.loads(pp.read_text()).get("pins", []) if pp.exists() else []):
            px, py = to_local(lat, lon, float(p["lat"]), float(p["lon"]), R)
            if abs(px) <= half and abs(py) <= half:
                pins.append({"name": p.get("name"), "lore": p.get("lore"), "epoch": p.get("epoch"),
                             "x": round(float(px), 1), "y": round(float(py), 1), "lat": p["lat"], "lon": p["lon"]})
        (tmp / "pins.json").write_text(json.dumps(pins, ensure_ascii=False, indent=1))
        leg = world.legends
        sea = cls["water"] == 1
        i0 = world.index_of(lat, lon)
        zones = {k: round(float(world.arrays[k][i0]), 4) for k in
                 ("gravity_g", "o2_fraction", "po2_bar", "pressure_bar", "fire_reactivity") if k in world.arrays}
        vents = []
        for v in serve.vents_of(src)["vents"]:
            vx, vy = to_local(lat, lon, v["lat"], v["lon"], R)
            if abs(vx) <= half and abs(vy) <= half:
                vents.append({"x": round(float(vx), 1), "y": round(float(vy), 1),
                              **{k: v[k] for k in ("type", "temp_c", "flow", "mineral")}})
        (tmp / "legend.json").write_text(json.dumps({"water": WATER, "biome": leg["holdridge"], "ground": leg["ground"],
                                                     "landform": leg["landform"]}, ensure_ascii=False, indent=1))
        meta = {"format": "worldmap-terrain-export", "version": FORMAT_VERSION, "name": name,
                "center": {"lat": lat, "lon": lon}, "samples": n, "res_m": res_m, "size_m": (n - 1) * res_m,
                "projection": {"kind": "azimuthal equidistant", "sphere_radius_km": R, "x": "east", "y": "north",
                               "origin": "the centre sample", "row_0": "north edge", "grid": "samples on the edges"},
                "height": {"file": "height.f32 (float32 LE, m)", "min": lo_, "max": hi_,
                           "u16": "h = min + v / 65535 * (max - min)  (height.png, height.r16 LE)"},
                "sea_level_m": 0, "center_height_m": round(float(h[n // 2, n // 2]), 3), "zoom": z,
                "era": world.era,
                "water": {"sea_fraction": round(float(sea.mean()), 4),
                          "max_depth_m": round(float(-h[sea].min()), 1) if sea.any() else 0.0,
                          "center_depth_m": round(float(max(0.0, -h[n // 2, n // 2])), 1)
                          if sea[n // 2, n // 2] else 0.0},
                "vents": vents, "zones": zones,
                "preview": {"file": "preview.png", "exaggeration": round(exag, 2),
                            "note": "relief colours, hillshade from the north-west with this vertical exaggeration"},
                "tile_px_m": round(math.pi * R * 1000.0 / (2 ** z * T.TILE), 3),
                "sources": {"world_build": src.fingerprint, "regions": "" if rs.empty else rs.fingerprint,
                            "tile_version": T.VERSION, "region_version": T.REGION_VERSION, "refine_model": refine.MODEL,
                            "seed": int(cfg["build"]["seed"]), "world_res": res},
                "created": date.today().isoformat(),
                "notes": ["Generated terrain: build data ≥ ≈ 34 km, refined regions ≥ ≈ 5 km, procedural detail "
                          "below — plausible, not surveyed.",
                          "Heights: ground, lake and river surfaces; the sea floor below 0 m."]}
        (tmp / "meta.json").write_text(json.dumps(meta, ensure_ascii=False, indent=1))
        (tmp / "README.txt").write_text(README.format(**{**meta, "size_km": meta["size_m"] / 1000, "lo": lo_, "hi": hi_}))
        check()
        os.replace(tmp, dest)                                   # complete or not there at all
    except BaseException:
        shutil.rmtree(tmp, ignore_errors=True)
        raise
    return dest


def default_name(lat, lon, size_km, res_m) -> str:
    """e.g. s13.400-w30.400-40km-10m"""
    return f"{'n' if lat >= 0 else 's'}{abs(lat):.3f}-{'e' if lon >= 0 else 'w'}{abs(lon):.3f}-{size_km:g}km-{res_m:g}m"


README = """Terrain export "{name}" (format {version})

Square of {size_km:g} km around {center[lat]}, {center[lon]}: {samples} x {samples} samples every {res_m:g} m.
Row 0 is the north edge, samples lie on the edges (vertex grid); x east, y north from the centre sample
(azimuthal equidistant on a sphere of {projection[sphere_radius_km]:g} km).

height.f32   float32 little-endian metres (ground, lake and river surfaces; sea floor below 0 m; sea level 0 m)
height.png   16-bit grayscale, height.r16 16-bit little-endian raw: h = {lo:.2f} + v / 65535 * ({hi:.2f} - {lo:.2f}) m
water.png    0 land, 1 sea, 2 lake, 3 river channel
biome.png, ground.png, landform.png   class indices, names in legend.json
preview.png  the map's relief colours with hillshade
rivers.json  river centrelines [{{points: [[x, y], ...] m, width_m, discharge_km3_yr, refined}}]
pins.json    map pins in the square (x, y in m)
meta.json    all of the above, sources and versions; the era, water depth, vents in the square (x, y m; types in the
             viewer legend) and the centre's gravity / O₂ / air pressure / fire reactivity

Import hints: Unity: Terrain > Import Raw, height.r16, {samples} x {samples}, 16 bit, byte order Windows (little),
terrain size {size_m:g} x {size_m:g} m, height = max - min ({hi:.2f} - {lo:.2f}), and the terrain's Y position = min
({lo:.2f}) so heights come out in metres (sea level at Y = 0). Godot (Terrain3D / HTerrain): height.png or height.f32
(heights in metres). Unreal: height.png (resample to a landscape size such as 4033 or 8129 if asked).

Own engines: height.f32 is a plain row-major grid (x east, row 0 north); subtract center_height_m (meta.json) to put
the centre at 0 m.

Generated terrain.
"""