#!/usr/bin/env python3
"""
gpx_3d_viz.py — turn any GPX track into a self-contained, interactive 3D
HTML visualisation (Three.js), with a real terrain surface built from
public elevation data (Open Topo Data), an optional satellite/street-map
texture draped over that terrain, and an animated "follow the route"
playback control.

Requires Pillow for the basemap texture (pip install pillow). Everything
else is stdlib only. Pass --basemap none to skip imagery entirely (no
Pillow needed in that case).

Basic usage:
    python3 gpx_3d_viz.py track.gpx
    python3 gpx_3d_viz.py track.gpx -o out/my-hike-3d.html --title "My Hike"
    python3 gpx_3d_viz.py track.gpx --no-terrain          # skip the API call
    python3 gpx_3d_viz.py track.gpx --basemap satellite   # satellite only, no toggle
    python3 gpx_3d_viz.py track.gpx --basemap none         # elevation colours only
    python3 gpx_3d_viz.py track.gpx --export-track-json track.json \\
                                     --export-terrain-json terrain.json

Embed the result in a page with:
    <iframe src="/path/to/output.html" style="width:100%; height:520px; border:0;"></iframe>
"""

import argparse
import base64
import html
import json
import math
import re
import sys
import time
import urllib.request
import xml.etree.ElementTree as ET
from io import BytesIO
from pathlib import Path

try:
    from PIL import Image
except ImportError:
    Image = None

R_EARTH = 6371000.0
TILE_SIZE = 256

OTD_DATASETS = {
    "srtm30m": "https://api.opentopodata.org/v1/srtm30m",
    "srtm90m": "https://api.opentopodata.org/v1/srtm90m",
    "aster30m": "https://api.opentopodata.org/v1/aster30m",
    "mapzen": "https://api.opentopodata.org/v1/mapzen",
}

# Slippy-map (Web Mercator, 256px tiles) basemap providers usable for
# draping a real-world map/photo over the terrain mesh. Both are free and
# require no API key; the OpenStreetMap tile usage policy requires a
# descriptive User-Agent and discourages bulk/automated fetching, which is
# respected here by fetching only the handful of tiles needed for one
# static build, with a short delay between requests.
BASEMAP_PROVIDERS = {
    "satellite": {
        "label": "Satellite",
        "tile_url": lambda z, x, y: (
            f"https://server.arcgisonline.com/ArcGIS/rest/services/World_Imagery/MapServer/tile/{z}/{y}/{x}"
        ),
        "headers": {"User-Agent": "gpx-3d-viz/1.0 (+https://crisly.nl)"},
    },
    "map": {
        "label": "Map",
        "tile_url": lambda z, x, y: f"https://{'abc'[(x + y) % 3]}.tile.openstreetmap.org/{z}/{x}/{y}.png",
        "headers": {
            "User-Agent": "gpx-3d-viz/1.0 (+https://crisly.nl; contact c.lanting@gmail.com) "
                          "- one-off static site build, very low volume"
        },
    },
}


# --------------------------------------------------------------------------
# GPX parsing
# --------------------------------------------------------------------------

def parse_gpx(path):
    """Return (points, name) where points is a list of dicts:
    {lat, lon, ele, time} (ele/time may be None), in track order.
    Concatenates all tracks/segments in the file."""
    tree = ET.parse(path)
    root = tree.getroot()
    tag = root.tag
    nsuri = tag[tag.find("{") + 1: tag.find("}")] if "{" in tag else ""
    ns = {"g": nsuri}

    name_el = root.find(".//g:trk/g:name", ns)
    name = name_el.text.strip() if name_el is not None and name_el.text else None

    points = []
    for trkpt in root.findall(".//g:trkpt", ns):
        lat = float(trkpt.get("lat"))
        lon = float(trkpt.get("lon"))
        ele_el = trkpt.find("g:ele", ns)
        ele = float(ele_el.text) if ele_el is not None and ele_el.text else None
        time_el = trkpt.find("g:time", ns)
        t = time_el.text if time_el is not None else None
        points.append({"lat": lat, "lon": lon, "ele": ele, "time": t})

    if not points:
        raise ValueError(f"No <trkpt> elements found in {path}")
    if any(p["ele"] is None for p in points):
        raise ValueError("GPX file is missing elevation data on one or more points")

    return points, name


def smooth(values, window):
    if window <= 1:
        return list(values)
    out = []
    half = window // 2
    n = len(values)
    for i in range(n):
        lo = max(0, i - half)
        hi = min(n, i + half + 1)
        out.append(sum(values[lo:hi]) / (hi - lo))
    return out


def haversine(lat1, lon1, lat2, lon2):
    phi1, phi2 = math.radians(lat1), math.radians(lat2)
    dphi = math.radians(lat2 - lat1)
    dlmb = math.radians(lon2 - lon1)
    a = math.sin(dphi / 2) ** 2 + math.cos(phi1) * math.cos(phi2) * math.sin(dlmb / 2) ** 2
    return 2 * R_EARTH * math.asin(math.sqrt(a))


def track_stats(points, smoothed_eles):
    dist = sum(
        haversine(points[i - 1]["lat"], points[i - 1]["lon"], points[i]["lat"], points[i]["lon"])
        for i in range(1, len(points))
    )
    gain = loss = 0.0
    for i in range(1, len(smoothed_eles)):
        d = smoothed_eles[i] - smoothed_eles[i - 1]
        if d > 0:
            gain += d
        else:
            loss += -d

    date_label = None
    first_time = points[0]["time"]
    if first_time:
        date_label = first_time[:10]  # YYYY-MM-DD

    return {
        "distance_km": round(dist / 1000, 1),
        "gain_m": round(gain),
        "loss_m": round(loss),
        "min_ele": round(min(smoothed_eles)),
        "max_ele": round(max(smoothed_eles)),
        "date": date_label,
    }


# --------------------------------------------------------------------------
# Local ENU projection (equirectangular, fine for areas up to a few tens of km)
# --------------------------------------------------------------------------

def make_projector(lat0, lon0):
    def to_local(lat, lon):
        x = math.radians(lon - lon0) * R_EARTH * math.cos(math.radians(lat0))
        north = math.radians(lat - lat0) * R_EARTH
        return x, -north  # z = -north, matching the convention used throughout this tool
    return to_local


def downsample(seq, target):
    if len(seq) <= target:
        return list(seq)
    step = max(1, len(seq) // target)
    out = seq[::step]
    if out[-1] is not seq[-1]:
        out.append(seq[-1])
    return out


def slugify(value):
    slug = re.sub(r"[^a-z0-9]+", "-", value.lower())
    return slug.strip("-") or "track"


def build_embed_snippet(src, title, height=520):
    escaped_src = html.escape(src, quote=True)
    escaped_title = html.escape(title, quote=True)
    return f"""<div markdown="0">
  <iframe src="{escaped_src}" style="width:100%; height:{height}px; border:0; border-radius:10px; display:block;" loading="lazy" title="3D track of {escaped_title}"></iframe>
  <p style="font-size:0.8em; color:var(--text-muted, #666); margin-top:0.4rem;">The recorded GPX track set against the surrounding terrain (SRTM elevation data) &mdash; drag to rotate, scroll to zoom, or press play to follow the route. Elevation is exaggerated &times;2.4 for clarity.</p>
</div>"""


# --------------------------------------------------------------------------
# Terrain: fetch a grid of elevations around the track's bounding box
# --------------------------------------------------------------------------

def fetch_terrain_grid(lat_min, lat_max, lon_min, lon_max, nx, ny, dataset,
                        batch_size=90, delay=1.1, retries=3, timeout=20, quiet=False):
    if dataset not in OTD_DATASETS:
        raise ValueError(f"Unknown terrain dataset '{dataset}', choose from {list(OTD_DATASETS)}")
    base_url = OTD_DATASETS[dataset]

    lats = [lat_min + (lat_max - lat_min) * i / (ny - 1) for i in range(ny)]
    lons = [lon_min + (lon_max - lon_min) * j / (nx - 1) for j in range(nx)]
    all_pts = [(i, j, lats[i], lons[j]) for i in range(ny) for j in range(nx)]

    grid = {}
    for b in range(0, len(all_pts), batch_size):
        chunk = all_pts[b:b + batch_size]
        loc_str = "|".join(f"{lat:.6f},{lon:.6f}" for (_, _, lat, lon) in chunk)
        url = f"{base_url}?locations={loc_str}"
        for attempt in range(retries):
            try:
                with urllib.request.urlopen(url, timeout=timeout) as resp:
                    data = json.loads(resp.read())
                break
            except Exception as e:
                if not quiet:
                    print(f"  terrain fetch retry {attempt}: {e}", file=sys.stderr)
                time.sleep(2)
        else:
            raise RuntimeError("Terrain API request failed after retries")
        for (i, j, _, _), r in zip(chunk, data["results"]):
            grid[(i, j)] = r["elevation"] if r["elevation"] is not None else 0.0
        if not quiet:
            print(f"  terrain: {min(b + len(chunk), len(all_pts))}/{len(all_pts)} points", file=sys.stderr)
        time.sleep(delay)

    elevations = [[grid[(i, j)] for j in range(nx)] for i in range(ny)]
    return {
        "lat_min": lat_min, "lat_max": lat_max,
        "lon_min": lon_min, "lon_max": lon_max,
        "nx": nx, "ny": ny,
        "elevations": elevations,
    }


def bilinear_sample(terrain, lat, lon):
    """Sample an elevation from a lat/lon-gridded terrain dict via bilinear
    interpolation. Used to drape the GPX track onto the terrain surface so
    the two never disagree (no floating/clipping "Z-offset")."""
    lat_min, lat_max = terrain["lat_min"], terrain["lat_max"]
    lon_min, lon_max = terrain["lon_min"], terrain["lon_max"]
    nx, ny = terrain["nx"], terrain["ny"]
    elev = terrain["elevations"]

    fi = (lat - lat_min) / (lat_max - lat_min) * (ny - 1)
    fj = (lon - lon_min) / (lon_max - lon_min) * (nx - 1)
    fi = min(max(fi, 0.0), ny - 1)
    fj = min(max(fj, 0.0), nx - 1)

    i0 = int(math.floor(fi))
    i1 = min(i0 + 1, ny - 1)
    j0 = int(math.floor(fj))
    j1 = min(j0 + 1, nx - 1)
    ti = fi - i0
    tj = fj - j0

    e00, e01 = elev[i0][j0], elev[i0][j1]
    e10, e11 = elev[i1][j0], elev[i1][j1]
    e0 = e00 + (e01 - e00) * tj
    e1 = e10 + (e11 - e10) * tj
    return e0 + (e1 - e0) * ti


def upsample_terrain_grid(terrain, factor):
    """Bilinearly upsample a terrain grid by an integer factor for a smoother,
    less blocky mesh — no extra API calls required."""
    if factor <= 1:
        return terrain
    nx, ny = terrain["nx"], terrain["ny"]
    new_nx = (nx - 1) * factor + 1
    new_ny = (ny - 1) * factor + 1
    new_elev = []
    for i in range(new_ny):
        lat = terrain["lat_min"] + (terrain["lat_max"] - terrain["lat_min"]) * i / (new_ny - 1)
        row = []
        for j in range(new_nx):
            lon = terrain["lon_min"] + (terrain["lon_max"] - terrain["lon_min"]) * j / (new_nx - 1)
            row.append(bilinear_sample(terrain, lat, lon))
        new_elev.append(row)
    return {
        "lat_min": terrain["lat_min"], "lat_max": terrain["lat_max"],
        "lon_min": terrain["lon_min"], "lon_max": terrain["lon_max"],
        "nx": new_nx, "ny": new_ny,
        "elevations": new_elev,
    }


# --------------------------------------------------------------------------
# Basemap imagery: fetch satellite/street-map tiles and drape them over the
# terrain mesh as a texture (in addition to, or instead of, the elevation
# colour gradient above).
# --------------------------------------------------------------------------

def lonlat_to_pixel(lat, lon, zoom):
    """Standard Web Mercator slippy-map global pixel coordinates at the
    given zoom (256px tiles) — the same scheme OSM/Esri/Google tile servers
    use, so tile index = pixel // 256."""
    lat = max(min(lat, 85.0511), -85.0511)
    n = 2 ** zoom
    x = (lon + 180.0) / 360.0 * n * TILE_SIZE
    lat_rad = math.radians(lat)
    y = (1.0 - math.log(math.tan(lat_rad) + 1 / math.cos(lat_rad)) / math.pi) / 2.0 * n * TILE_SIZE
    return x, y


def choose_zoom(lat_min, lat_max, lon_min, lon_max, target_px=900, min_zoom=10, max_zoom=17):
    """Pick the highest zoom whose bounding box still maps to roughly
    target_px pixels on its longer side — enough resolution to look sharp
    without fetching more tiles than necessary."""
    best = min_zoom
    for zoom in range(min_zoom, max_zoom + 1):
        x0, y0 = lonlat_to_pixel(lat_max, lon_min, zoom)
        x1, y1 = lonlat_to_pixel(lat_min, lon_max, zoom)
        span = max(abs(x1 - x0), abs(y1 - y0))
        if span > target_px * 1.6:
            break
        best = zoom
    return best


def latlon_to_uv(lat, lon, crop_bounds):
    """Map a lat/lon to (u, v) inside a basemap image previously cropped
    via fetch_basemap_image, using the exact same pixel projection so the
    texture lines up with the terrain mesh with no manual offset."""
    px0, py0, px1, py1, zoom = crop_bounds
    x, y = lonlat_to_pixel(lat, lon, zoom)
    u = (x - px0) / (px1 - px0) if px1 != px0 else 0.5
    v = (y - py0) / (py1 - py0) if py1 != py0 else 0.5
    return u, v


def fetch_basemap_image(lat_min, lat_max, lon_min, lon_max, provider, zoom,
                         max_dim=900, delay=0.15, retries=3, timeout=15, quiet=False):
    """Fetch the slippy-map tiles covering the bbox at `zoom` from the given
    provider (a BASEMAP_PROVIDERS entry), stitch them into one composite
    image, then crop to the exact bbox. Returns (PIL.Image, crop_bounds);
    crop_bounds is reused later to compute matching per-vertex UVs."""
    if Image is None:
        raise RuntimeError(
            "Basemap texture draping requires Pillow. Install it with "
            "'pip install pillow' or pass --basemap none to skip imagery."
        )

    px0, py0 = lonlat_to_pixel(lat_max, lon_min, zoom)  # top-left: max lat, min lon
    px1, py1 = lonlat_to_pixel(lat_min, lon_max, zoom)  # bottom-right: min lat, max lon
    px0, px1 = sorted((px0, px1))
    py0, py1 = sorted((py0, py1))

    tx0, tx1 = int(px0 // TILE_SIZE), int(px1 // TILE_SIZE)
    ty0, ty1 = int(py0 // TILE_SIZE), int(py1 // TILE_SIZE)
    n_tiles = (tx1 - tx0 + 1) * (ty1 - ty0 + 1)
    if n_tiles > 64:
        raise RuntimeError(
            f"Basemap fetch would need {n_tiles} tiles at zoom {zoom} for {provider['label']} — "
            f"aborting to stay polite to the tile server. Try --texture-zoom with a lower value."
        )

    mosaic = Image.new("RGB", ((tx1 - tx0 + 1) * TILE_SIZE, (ty1 - ty0 + 1) * TILE_SIZE), (190, 190, 190))
    for ty in range(ty0, ty1 + 1):
        for tx in range(tx0, tx1 + 1):
            url = provider["tile_url"](zoom, tx, ty)
            tile_img = None
            for attempt in range(retries):
                try:
                    req = urllib.request.Request(url, headers=provider["headers"])
                    with urllib.request.urlopen(req, timeout=timeout) as resp:
                        tile_img = Image.open(BytesIO(resp.read())).convert("RGB")
                    break
                except Exception as e:
                    if not quiet:
                        print(f"  {provider['label']} tile fetch retry {attempt}: {e}", file=sys.stderr)
                    time.sleep(1)
            if tile_img is None:
                tile_img = Image.new("RGB", (TILE_SIZE, TILE_SIZE), (190, 190, 190))
            mosaic.paste(tile_img, ((tx - tx0) * TILE_SIZE, (ty - ty0) * TILE_SIZE))
            time.sleep(delay)

    crop_box = (
        int(round(px0 - tx0 * TILE_SIZE)),
        int(round(py0 - ty0 * TILE_SIZE)),
        int(round(px1 - tx0 * TILE_SIZE)),
        int(round(py1 - ty0 * TILE_SIZE)),
    )
    img = mosaic.crop(crop_box)

    if max(img.size) > max_dim:
        scale = max_dim / max(img.size)
        new_size = (max(1, int(img.width * scale)), max(1, int(img.height * scale)))
        img = img.resize(new_size, Image.LANCZOS)

    return img, (px0, py0, px1, py1, zoom)


def image_to_data_uri(img, quality=83):
    buf = BytesIO()
    img.convert("RGB").save(buf, format="JPEG", quality=quality, optimize=True)
    b64 = base64.b64encode(buf.getvalue()).decode("ascii")
    return f"data:image/jpeg;base64,{b64}"


def terrain_color(t, stops):
    for k in range(len(stops) - 1):
        t0, c0 = stops[k]
        t1, c1 = stops[k + 1]
        if t0 <= t <= t1:
            f = 0 if t1 == t0 else (t - t0) / (t1 - t0)
            return tuple(c0[i] + (c1[i] - c0[i]) * f for i in range(3))
    return stops[-1][1]


DEFAULT_TERRAIN_STOPS = [
    (0.00, (0.20, 0.36, 0.18)),  # vegetated low ground
    (0.35, (0.36, 0.45, 0.22)),  # olive
    (0.65, (0.55, 0.48, 0.38)),  # tan / bare rock
    (1.00, (0.62, 0.58, 0.56)),  # grey rock
]


def build_terrain_mesh(terrain, to_local, min_elev, max_elev, exaggeration, color_stops, uv_bounds=None):
    nx, ny = terrain["nx"], terrain["ny"]
    lats = [terrain["lat_min"] + (terrain["lat_max"] - terrain["lat_min"]) * i / (ny - 1) for i in range(ny)]
    lons = [terrain["lon_min"] + (terrain["lon_max"] - terrain["lon_min"]) * j / (nx - 1) for j in range(nx)]
    lat0_ref = to_local(lats[0], lons[0])  # unused, just keeps signature simple
    elev = terrain["elevations"]

    xs = [to_local(lats[0], lon)[0] for lon in lons]
    zs = [to_local(lat, lons[0])[1] for lat in lats]

    positions = []
    colors = []
    uvs = []
    span = (max_elev - min_elev) or 1.0
    xspan = (max(xs) - min(xs)) or 1.0
    zspan = (max(zs) - min(zs)) or 1.0
    for i in range(ny):
        for j in range(nx):
            e = elev[i][j]
            y = (e - min_elev) * exaggeration
            positions.append((round(xs[j], 1), round(y, 1), round(zs[i], 1)))
            t = (e - min_elev) / span
            colors.append(tuple(round(c, 3) for c in terrain_color(t, color_stops)))
            if uv_bounds:
                u, v = latlon_to_uv(lats[i], lons[j], uv_bounds)
            else:
                # Fallback: normalise local x/z directly (no real basemap to align to).
                u = (xs[j] - min(xs)) / xspan
                v = (zs[i] - min(zs)) / zspan
            uvs.append((round(u, 4), round(v, 4)))

    indices = []
    for i in range(ny - 1):
        for j in range(nx - 1):
            a = i * nx + j
            b = a + 1
            c = a + nx
            d = c + 1
            indices.extend([a, b, c, b, d, c])

    return {
        "positions": [v for p in positions for v in p],
        "colors": [v for p in colors for v in p],
        "uvs": [v for p in uvs for v in p],
        "indices": indices,
    }


# --------------------------------------------------------------------------
# HTML template
# --------------------------------------------------------------------------

HTML_TEMPLATE = """<!DOCTYPE html>
<html lang="en">
<head>
<meta charset="utf-8">
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>{title} &mdash; 3D track</title>
<style>
  html, body {{ margin: 0; padding: 0; overflow: hidden; background: linear-gradient(180deg, #dce9f0 0%, #f4f1ea 100%); font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Helvetica, Arial, sans-serif; }}
  #app {{ position: relative; width: 100%; height: 100vh; }}
  canvas {{ display: block; }}
  #stats {{
    position: absolute; top: 12px; left: 12px;
    background: rgba(255,255,255,0.85); backdrop-filter: blur(4px);
    border-radius: 10px; padding: 10px 14px; font-size: 12.5px; color: #2d2d30;
    box-shadow: 0 2px 10px rgba(0,0,0,0.12); line-height: 1.5; max-width: 220px;
  }}
  #stats b {{ display:block; font-size: 13.5px; margin-bottom: 4px; }}
  #controls {{
    position: absolute; bottom: 12px; left: 12px; right: 12px;
    display: flex; align-items: center; gap: 10px;
    background: rgba(255,255,255,0.85); backdrop-filter: blur(4px);
    border-radius: 10px; padding: 8px 14px; box-shadow: 0 2px 10px rgba(0,0,0,0.12);
  }}
  #controls button {{
    border: none; background: #2d2d30; color: #fff; border-radius: 6px;
    width: 30px; height: 30px; font-size: 14px; cursor: pointer; flex: 0 0 auto;
  }}
  #controls button:hover {{ background: #45454a; }}
  #progress {{ flex: 1 1 auto; }}
  #hint {{
    position: absolute; top: 12px; right: 12px;
    background: rgba(255,255,255,0.7); border-radius: 8px; padding: 6px 10px;
    font-size: 11px; color: #555;
  }}
  #basemapToggle {{
    position: absolute; top: 48px; right: 12px; display: none;
    border: none; background: rgba(255,255,255,0.85); backdrop-filter: blur(4px);
    color: #2d2d30; border-radius: 8px; padding: 6px 12px; font-size: 11px;
    cursor: pointer; box-shadow: 0 2px 10px rgba(0,0,0,0.12);
  }}
  #basemapToggle:hover {{ background: #fff; }}
  @media (max-width: 480px) {{
    #stats {{ font-size: 11px; padding: 8px 10px; max-width: 150px; }}
    #hint {{ display: none; }}
  }}
</style>
</head>
<body>
<div id="app">
  <div id="stats">
    <b>{title}</b>
    {stats_line}
  </div>
  <div id="hint">drag to rotate &middot; scroll to zoom</div>
  <button id="basemapToggle">Terrain</button>
  <div id="controls">
    <button id="playBtn">&#9658;</button>
    <input type="range" id="progress" min="0" max="1000" value="0">
  </div>
</div>

<script type="importmap">
{{
  "imports": {{
    "three": "https://unpkg.com/three@0.160.0/build/three.module.js"
  }}
}}
</script>
<script type="module">
import * as THREE from 'three';
import {{ OrbitControls }} from 'https://unpkg.com/three@0.160.0/examples/jsm/controls/OrbitControls.js';

const track = {track_json}; // [x, z, elevation] in metres, local ENU — z already matches the terrain mesh's z directly
const EXAGGERATION = {exaggeration};
const MIN_ELEV = {min_elev};   // shared baseline with the terrain mesh below (if any)
const LIFT = {lift};          // metres: raise the recorded track slightly above the terrain surface for visibility

const points = track.map(p => new THREE.Vector3(p[0], (p[2] + LIFT - MIN_ELEV) * EXAGGERATION, p[1]));

const app = document.getElementById('app');
const scene = new THREE.Scene();

const camera = new THREE.PerspectiveCamera(45, app.clientWidth / app.clientHeight, 1, 20000);
const renderer = new THREE.WebGLRenderer({{ antialias: true, alpha: true }});
renderer.setPixelRatio(Math.min(window.devicePixelRatio, 2));
renderer.setSize(app.clientWidth, app.clientHeight);
app.appendChild(renderer.domElement);

// Lighting — lower ambient / stronger directional sun gives the terrain more
// relief (closer to a hillshaded map look) than flat, evenly-lit shading.
scene.add(new THREE.AmbientLight(0xffffff, 0.45));
const sun = new THREE.DirectionalLight(0xfff4e0, 1.5);
sun.position.set(600, 900, 300);
scene.add(sun);
const fill = new THREE.DirectionalLight(0xb9d4ff, 0.3);
fill.position.set(-500, 300, -400);
scene.add(fill);

{terrain_block}

// Build the track as a tube, coloured by elevation (low = teal, high = warm orange)
const curve = new THREE.CatmullRomCurve3(points);
const tubeGeo = new THREE.TubeGeometry(curve, Math.max(200, points.length * 2), {tube_radius}, 8, false);

const lowColor = new THREE.Color({track_low_color});
const highColor = new THREE.Color({track_high_color});
const posAttr = tubeGeo.attributes.position;
const colors = new Float32Array(posAttr.count * 3);
const maxY = Math.max(...points.map(p => p.y)) || 1;
for (let i = 0; i < posAttr.count; i++) {{
  const y = posAttr.getY(i);
  const t = Math.min(1, Math.max(0, y / maxY));
  const c = lowColor.clone().lerp(highColor, t);
  colors[i*3] = c.r; colors[i*3+1] = c.g; colors[i*3+2] = c.b;
}}
tubeGeo.setAttribute('color', new THREE.BufferAttribute(colors, 3));

const tubeMat = new THREE.MeshStandardMaterial({{ vertexColors: true, roughness: 0.55, metalness: 0.05 }});
const tubeMesh = new THREE.Mesh(tubeGeo, tubeMat);
scene.add(tubeMesh);

// Trailhead marker
const startMarker = new THREE.Mesh(
  new THREE.SphereGeometry(9, 24, 24),
  new THREE.MeshStandardMaterial({{ color: {track_low_color} }})
);
startMarker.position.copy(points[0]);
scene.add(startMarker);

// Highest point reached
let highestIdx = 0;
points.forEach((p, i) => {{ if (p.y > points[highestIdx].y) highestIdx = i; }});
const summitMarker = new THREE.Mesh(
  new THREE.SphereGeometry(9, 24, 24),
  new THREE.MeshStandardMaterial({{ color: {track_high_color} }})
);
summitMarker.position.copy(points[highestIdx]);
scene.add(summitMarker);

// "You are here" marker, animated along the path
const hiker = new THREE.Mesh(
  new THREE.SphereGeometry({hiker_radius}, 24, 24),
  new THREE.MeshStandardMaterial({{ color: 0xffffff, emissive: 0x8899ff, emissiveIntensity: 0.4 }})
);
hiker.position.copy(points[0]);
scene.add(hiker);
const hikerGlow = new THREE.PointLight(0xffffff, 1.2, {hiker_glow_distance});
hiker.add(hikerGlow);

// Camera framing
const box = new THREE.Box3().setFromObject(tubeMesh);{box_union}
const center = box.getCenter(new THREE.Vector3());
const size = box.getSize(new THREE.Vector3());
const radius = Math.max(size.x, size.z, size.y * 0.6);
camera.position.set(center.x + radius * 0.9, center.y + radius * 0.6, center.z + radius * 0.9);

const controls = new OrbitControls(camera, renderer.domElement);
controls.target.copy(center);
controls.enableDamping = true;
controls.dampingFactor = 0.08;
controls.minDistance = radius * 0.3;
controls.maxDistance = radius * 3;
controls.update();

window.addEventListener('resize', () => {{
  camera.aspect = app.clientWidth / app.clientHeight;
  camera.updateProjectionMatrix();
  renderer.setSize(app.clientWidth, app.clientHeight);
}});

// Playback along the curve
const progressEl = document.getElementById('progress');
const playBtn = document.getElementById('playBtn');
let playing = false;
let t = 0;

function setHikerAt(tt) {{
  const p = curve.getPointAt(Math.min(1, Math.max(0, tt)));
  hiker.position.copy(p);
}}

playBtn.addEventListener('click', () => {{
  playing = !playing;
  playBtn.innerHTML = playing ? '&#10074;&#10074;' : '&#9658;';
}});
progressEl.addEventListener('input', () => {{
  t = parseFloat(progressEl.value) / 1000;
  playing = false;
  playBtn.innerHTML = '&#9658;';
  setHikerAt(t);
}});

function animate() {{
  requestAnimationFrame(animate);
  if (playing) {{
    t += 0.0016;
    if (t > 1) t = 0;
    progressEl.value = String(Math.round(t * 1000));
    setHikerAt(t);
  }}
  controls.update();
  renderer.render(scene, camera);
}}
animate();
</script>
</body>
</html>
"""

TERRAIN_BLOCK_TEMPLATE = """// Terrain surface built from a {nx}x{ny} grid of {dataset} elevation samples
// around the track, projected into the same local ENU coordinate frame.
// flatShading gives each triangle its own face normal (via screen-space
// derivatives — no special non-indexed geometry needed), which reads as a
// rugged, faceted rock surface instead of a smooth, rounded blob.
const terrainGeo = new THREE.BufferGeometry();
const terrainPositions = new Float32Array({positions_json});
const terrainColorsArr = new Float32Array({colors_json});
const terrainUVs = new Float32Array({uvs_json});
terrainGeo.setAttribute('position', new THREE.BufferAttribute(terrainPositions, 3));
terrainGeo.setAttribute('color', new THREE.BufferAttribute(terrainColorsArr, 3));
terrainGeo.setAttribute('uv', new THREE.BufferAttribute(terrainUVs, 2));
terrainGeo.setIndex({indices_json});
terrainGeo.computeVertexNormals();

// Terrain materials: elevation-tint colours always available; satellite
// and/or street-map textures added below if they were fetched at build time.
const terrainMaterials = [
  new THREE.MeshStandardMaterial({{ vertexColors: true, roughness: 0.95, metalness: 0.0, side: THREE.DoubleSide, flatShading: true }})
];
const terrainModeNames = ['Terrain'];
{texture_materials_block}
const terrainMesh = new THREE.Mesh(terrainGeo, terrainMaterials[{default_mode_idx}]);
scene.add(terrainMesh);

const basemapBtn = document.getElementById('basemapToggle');
if (terrainMaterials.length > 1) {{
  let basemapIdx = {default_mode_idx};
  basemapBtn.textContent = terrainModeNames[basemapIdx];
  basemapBtn.style.display = 'block';
  basemapBtn.addEventListener('click', () => {{
    basemapIdx = (basemapIdx + 1) % terrainMaterials.length;
    terrainMesh.material = terrainMaterials[basemapIdx];
    basemapBtn.textContent = terrainModeNames[basemapIdx];
  }});
}}"""

TEXTURE_MATERIAL_SNIPPET = """{{
  const tex = new THREE.TextureLoader().load('{data_uri}');
  tex.flipY = false;
  tex.colorSpace = THREE.SRGBColorSpace;
  terrainMaterials.push(new THREE.MeshStandardMaterial({{ map: tex, roughness: 0.92, metalness: 0.0, side: THREE.DoubleSide, flatShading: true }}));
  terrainModeNames.push('{label}');
}}"""

NO_TERRAIN_BLOCK = """// No terrain surface requested (--no-terrain) — a plain reference grid instead.
const grid = new THREE.GridHelper(2400, 24, 0xb9b2a0, 0xd8d2c2);
grid.position.y = -2;
scene.add(grid);"""


def build_stats_line(stats):
    parts = [f"{stats['distance_km']} km", f"{stats['min_ele']}&ndash;{stats['max_ele']} m"]
    line1 = " &middot; ".join(parts)
    line2_parts = [f"~{stats['gain_m']} m climbed"]
    if stats["date"]:
        line2_parts.append(stats["date"])
    line2 = " &middot; ".join(line2_parts)
    return f"{line1}<br>\n    {line2}"


# --------------------------------------------------------------------------
# Main build
# --------------------------------------------------------------------------

def build(args, gpx_arg=None, output_arg=None):
    gpx_path = Path(gpx_arg if gpx_arg is not None else args.gpx)
    points, gpx_name = parse_gpx(gpx_path)

    lats = [p["lat"] for p in points]
    lons = [p["lon"] for p in points]
    eles = [p["ele"] for p in points]
    smoothed = smooth(eles, args.smooth_window)
    # Stats always describe the actual recorded hike (GPS elevation), regardless
    # of how the track is rendered below.
    stats = track_stats(points, smoothed)

    lat0 = sum(lats) / len(lats)
    lon0 = sum(lons) / len(lons)
    to_local = make_projector(lat0, lon0)

    title = args.title or gpx_name or gpx_path.stem.replace("_", " ").replace("-", " ").title()

    terrain_mesh = None
    terrain_raw = None
    drape = not args.no_drape and not args.no_terrain
    if not args.no_terrain:
        margin = args.terrain_margin
        lat_min, lat_max = min(lats) - margin, max(lats) + margin
        lon_min, lon_max = min(lons) - margin, max(lons) + margin

        if args.terrain_expand:
            # Pad the bounding box further by a fraction of its own span on
            # each side, e.g. 0.10 = +/-10% per side = +20% extent per axis.
            lat_pad = (lat_max - lat_min) * args.terrain_expand
            lon_pad = (lon_max - lon_min) * args.terrain_expand
            lat_min -= lat_pad
            lat_max += lat_pad
            lon_min -= lon_pad
            lon_max += lon_pad

        nx, ny = args.terrain_grid
        print(f"Fetching {nx}x{ny} terrain grid from {args.terrain_dataset}...", file=sys.stderr)
        terrain_raw = fetch_terrain_grid(
            lat_min, lat_max, lon_min, lon_max, nx, ny, args.terrain_dataset,
            batch_size=args.api_batch_size, delay=args.api_delay,
        )

    # Choose the elevation used to *render* the track. Draping it onto the
    # terrain (instead of using the GPX's own recorded/GPS elevation) makes
    # the line sit exactly on the terrain surface with no visible offset —
    # GPS altitude and DEM elevation rarely agree to better than 20-40m, and
    # that gap is what shows up as a track floating above or sinking into
    # the ground.
    if drape:
        render_eles = [bilinear_sample(terrain_raw, p["lat"], p["lon"]) for p in points]
    else:
        render_eles = smoothed

    coords = []
    for i, p in enumerate(points):
        x, z = to_local(p["lat"], p["lon"])
        coords.append((x, z, render_eles[i]))
    sampled = downsample(coords, args.max_track_points)

    track_min_render = min(c[2] for c in sampled)
    track_max_render = max(c[2] for c in sampled)

    # Basemap imagery (satellite/street-map) draped over the terrain mesh —
    # only meaningful if a terrain mesh is actually being built.
    basemap_modes = []
    if args.basemap == "both":
        basemap_modes = ["satellite", "map"]
    elif args.basemap in ("satellite", "map"):
        basemap_modes = [args.basemap]
    if not terrain_raw:
        basemap_modes = []

    uv_bounds = None
    texture_uris = {}
    if basemap_modes:
        zoom = args.texture_zoom or choose_zoom(lat_min, lat_max, lon_min, lon_max)
        print(f"Fetching basemap imagery ({', '.join(basemap_modes)}) at zoom {zoom}...", file=sys.stderr)
        for mode in basemap_modes:
            provider = BASEMAP_PROVIDERS[mode]
            img, crop_bounds = fetch_basemap_image(
                lat_min, lat_max, lon_min, lon_max, provider, zoom,
                max_dim=args.texture_max_dim, delay=args.texture_delay,
            )
            uv_bounds = crop_bounds  # identical for every provider at this bbox/zoom
            texture_uris[mode] = image_to_data_uri(img, quality=args.texture_quality)

    if terrain_raw:
        terrain_for_mesh = upsample_terrain_grid(terrain_raw, args.terrain_upsample)
        flat = [e for row in terrain_for_mesh["elevations"] for e in row]
        min_elev = min(min(flat), track_min_render)
        max_elev = max(max(flat), track_max_render)
        terrain_mesh = build_terrain_mesh(terrain_for_mesh, to_local, min_elev, max_elev,
                                           args.exaggeration, DEFAULT_TERRAIN_STOPS, uv_bounds=uv_bounds)
    else:
        min_elev = track_min_render

    track_json = json.dumps([[round(c[0], 1), round(c[1], 1), round(c[2], 1)] for c in sampled])

    if terrain_mesh:
        texture_materials_block = "".join(
            TEXTURE_MATERIAL_SNIPPET.format(data_uri=texture_uris[mode], label=BASEMAP_PROVIDERS[mode]["label"]) + "\n"
            for mode in basemap_modes
        )
        # Default to the first fetched texture (satellite, if present) rather
        # than the plain elevation tint, since that's the more striking view.
        default_mode_idx = 1 if basemap_modes else 0

        terrain_block = TERRAIN_BLOCK_TEMPLATE.format(
            nx=terrain_for_mesh["nx"], ny=terrain_for_mesh["ny"], dataset=args.terrain_dataset,
            positions_json=json.dumps([round(v, 1) for v in terrain_mesh["positions"]]),
            colors_json=json.dumps([round(v, 3) for v in terrain_mesh["colors"]]),
            uvs_json=json.dumps(terrain_mesh["uvs"]),
            indices_json=json.dumps(terrain_mesh["indices"]),
            texture_materials_block=texture_materials_block,
            default_mode_idx=default_mode_idx,
        )
        box_union = "\nbox.union(new THREE.Box3().setFromObject(terrainMesh));"
    else:
        terrain_block = NO_TERRAIN_BLOCK
        box_union = ""

    html_doc = HTML_TEMPLATE.format(
        title=html.escape(title),
        stats_line=build_stats_line(stats),
        track_json=track_json,
        exaggeration=args.exaggeration,
        min_elev=round(min_elev, 1),
        lift=args.lift,
        terrain_block=terrain_block,
        box_union=box_union,
        tube_radius=args.tube_radius,
        track_low_color=args.track_low_color,
        track_high_color=args.track_high_color,
        hiker_radius=args.hiker_radius,
        hiker_glow_distance=round(args.hiker_radius * 8, 1),
    )

    if output_arg is not None:
        out_path = Path(output_arg)
    elif args.output:
        out_path = Path(args.output)
    elif args.output_dir:
        out_path = Path(args.output_dir) / f"{slugify(title)}-3d.html"
    else:
        out_path = gpx_path.with_name(gpx_path.stem + "-3d.html")
    out_path.parent.mkdir(parents=True, exist_ok=True)
    out_path.write_text(html_doc, encoding="utf-8")
    print(f"Wrote {out_path} ({len(html_doc)} bytes)", file=sys.stderr)

    if args.export_track_json:
        track_export = {
            "title": title,
            "stats": stats,
            "lat0": lat0,
            "lon0": lon0,
            "points": [
                {"x": round(c[0], 2), "z": round(c[1], 2), "ele": round(c[2], 2)}
                for c in coords
            ],
        }
        Path(args.export_track_json).write_text(json.dumps(track_export, indent=2), encoding="utf-8")
        print(f"Wrote {args.export_track_json}", file=sys.stderr)

    if args.export_terrain_json and terrain_raw:
        Path(args.export_terrain_json).write_text(json.dumps(terrain_raw, indent=2), encoding="utf-8")
        print(f"Wrote {args.export_terrain_json}", file=sys.stderr)
    elif args.export_terrain_json:
        print("No terrain data to export (--no-terrain was set)", file=sys.stderr)

    return stats, title, out_path


def parse_grid(value):
    if "x" not in value:
        raise argparse.ArgumentTypeError("expected NXxNY, e.g. 14x32")
    nx, ny = value.split("x", 1)
    return int(nx), int(ny)


def main():
    p = argparse.ArgumentParser(
        description="Build a self-contained interactive 3D HTML visualisation from a GPX track.",
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=__doc__,
    )
    p.add_argument("gpx", nargs="+", help="Path(s) to the input .gpx file(s)")
    p.add_argument("-o", "--output", help="Output HTML path (single GPX only; default: <gpx-name>-3d.html)")
    p.add_argument("--output-dir",
                   help="Directory for generated HTML files. Useful for batch builds, e.g. ../assets/visualizations")
    p.add_argument("--print-embed", action="store_true",
                   help="Print Jekyll/Markdown iframe snippets for the generated visualisations")
    p.add_argument("--embed-prefix", default="{{ site.baseurl }}/assets/visualizations",
                   help="URL prefix used in --print-embed snippets (default: {{ site.baseurl }}/assets/visualizations)")
    p.add_argument("--embed-height", type=int, default=520,
                   help="Iframe height in --print-embed snippets (default: 520)")
    p.add_argument("--title", help="Title shown on the visualisation (default: GPX <name> or filename)")

    p.add_argument("--exaggeration", type=float, default=2.4, help="Vertical exaggeration factor (default: 2.4)")
    p.add_argument("--lift", type=float, default=6.0,
                    help="Metres to raise the track above the terrain surface, for visibility (default: 6). "
                         "Only needs to clear the tube radius — the track is draped onto the terrain by default, "
                         "so there's no real elevation gap to cover any more.")
    p.add_argument("--no-drape", action="store_true",
                    help="Render the track at its own recorded/smoothed GPS elevation instead of draping it onto "
                         "the terrain DEM. Off by default because GPS altitude and DEM elevation routinely differ "
                         "by 20-40m, which shows up as the track floating above or sinking into the terrain.")
    p.add_argument("--smooth-window", type=int, default=15,
                    help="Rolling-average window (in track points) for elevation smoothing (default: 15)")
    p.add_argument("--max-track-points", type=int, default=400,
                    help="Downsample the track to roughly this many points (default: 400)")
    p.add_argument("--tube-radius", type=float, default=4.5, help="Radius of the rendered track tube (default: 4.5)")
    p.add_argument("--track-low-color", default="0x2f7d6b", help="Hex colour (e.g. 0x2f7d6b) for the track's low end")
    p.add_argument("--track-high-color", default="0xe8703a", help="Hex colour for the track's high end")
    p.add_argument("--hiker-radius", type=float, default=10.0,
                    help="Radius of the animated 'you are here' hiker marker sphere (default: 10.0)")

    p.add_argument("--no-terrain", action="store_true", help="Skip fetching terrain data; use a flat reference grid")
    p.add_argument("--terrain-dataset", default="srtm30m", choices=sorted(OTD_DATASETS),
                    help="Open Topo Data dataset to use (default: srtm30m)")
    p.add_argument("--terrain-margin", type=float, default=0.006,
                    help="Degrees of padding around the track's bounding box for the terrain grid (default: 0.006)")
    p.add_argument("--terrain-expand", type=float, default=0.0,
                    help="Pad the terrain bounding box further by this fraction of its own span on each side, "
                         "e.g. 0.10 = +/-10%% per side = +20%% extent per axis (default: 0.0, no extra padding)")
    p.add_argument("--terrain-grid", type=parse_grid, default=(14, 32),
                    help="Terrain sample grid as NXxNY, e.g. 20x20 (default: 14x32). "
                         "Keep nx*ny modest — each point is one API query.")
    p.add_argument("--terrain-upsample", type=int, default=3,
                    help="Bilinearly upsample the fetched terrain grid by this integer factor before building the "
                         "mesh, for a smoother, less blocky surface with no extra API calls (default: 3, set to 1 "
                         "to disable)")
    p.add_argument("--api-batch-size", type=int, default=90, help="Points per terrain API request (default: 90)")
    p.add_argument("--api-delay", type=float, default=1.1, help="Seconds to wait between API batches (default: 1.1)")

    p.add_argument("--basemap", choices=["both", "satellite", "map", "none"], default="both",
                    help="Drape a real-world basemap texture over the terrain mesh, with an on-screen toggle "
                         "button to switch views (default: both — satellite imagery + OpenStreetMap, cycling "
                         "through Satellite/Map/Terrain-colours). Requires Pillow (pip install pillow); has no "
                         "effect with --no-terrain. Needs an internet connection at build time, not at view time "
                         "— imagery is embedded directly in the output HTML.")
    p.add_argument("--texture-zoom", type=int, default=None,
                    help="Slippy-map zoom level for basemap tiles (default: auto, picked so the area maps to "
                         "roughly 900px on its longer side)")
    p.add_argument("--texture-max-dim", type=int, default=900,
                    help="Maximum width/height in pixels for the embedded basemap texture (default: 900)")
    p.add_argument("--texture-quality", type=int, default=83,
                    help="JPEG quality (1-95) for the embedded basemap texture (default: 83)")
    p.add_argument("--texture-delay", type=float, default=0.15,
                    help="Seconds to wait between individual tile requests (default: 0.15)")

    p.add_argument("--export-track-json", help="Also write the processed track (local ENU coords + stats) as JSON")
    p.add_argument("--export-terrain-json", help="Also write the raw fetched terrain grid as JSON")

    args = p.parse_args()

    if args.output and len(args.gpx) > 1:
        p.error("--output can only be used with a single GPX file; use --output-dir for batch builds")
    if args.title and len(args.gpx) > 1:
        p.error("--title can only be used with a single GPX file")

    built = []
    try:
        for gpx_arg in args.gpx:
            stats, title, out_path = build(args, gpx_arg=gpx_arg)
            built.append((stats, title, out_path))
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        sys.exit(1)

    for stats, title, out_path in built:
        print(
            f"{title}: {stats['distance_km']} km, {stats['min_ele']}-{stats['max_ele']} m, "
            f"~{stats['gain_m']} m climbed" + (f", {stats['date']}" if stats["date"] else "")
        )

    if args.print_embed:
        print()
        for _, title, out_path in built:
            src = f"{args.embed_prefix.rstrip('/')}/{out_path.name}"
            print(build_embed_snippet(src, title, height=args.embed_height))
            print()


if __name__ == "__main__":
    main()