planet/frontend/public/earth/js/terrain.js

import * as THREE from "three";

import { CONFIG, TERRAIN_CONFIG } from "./constants.js";
import { vector3ToLatLon } from "./utils.js";

const EARTH_RADIUS_METERS = 6371000;
const TERRAIN_COLOR_STOPS = [
  { height: 0, color: new THREE.Color(0x5f7f5b) },
  { height: 800, color: new THREE.Color(0x7f9564) },
  { height: 1800, color: new THREE.Color(0x9e956c) },
  { height: 3200, color: new THREE.Color(0x9f866a) },
  { height: 5200, color: new THREE.Color(0xc6c0b1) },
  { height: 7800, color: new THREE.Color(0xe8e5de) },
];

let terrainMesh = null;
let terrainLoadPromise = null;
let terrainReady = false;
let terrainFailed = false;
let terrainTileCache = new Map();
let resolvedTileCache = new Map();
let terrainVertexSamples = null;
let terrainOpacity = TERRAIN_CONFIG.opacity;

function clampLatitude(lat) {
  return THREE.MathUtils.clamp(lat, -85.05112878, 85.05112878);
}

function buildTileUrl(z, x, y) {
  return TERRAIN_CONFIG.urlTemplate
    .replace("{z}", String(z))
    .replace("{x}", String(x))
    .replace("{y}", String(y));
}

function buildTerrainTileKey(z, x, y) {
  return `${z}/${x}/${y}`;
}

function getTerrainCanvas(size) {
  if (typeof OffscreenCanvas !== "undefined") {
    return new OffscreenCanvas(size, size);
  }

  const canvas = document.createElement("canvas");
  canvas.width = size;
  canvas.height = size;
  return canvas;
}

async function decodeTerrainBlob(blob, cacheKey) {
  const bitmap = await createImageBitmap(blob);
  const canvas = getTerrainCanvas(TERRAIN_CONFIG.tileSize);
  const ctx = canvas.getContext("2d", { willReadFrequently: true });
  ctx.drawImage(bitmap, 0, 0, TERRAIN_CONFIG.tileSize, TERRAIN_CONFIG.tileSize);
  bitmap.close?.();
  const { data, width, height } = ctx.getImageData(
    0,
    0,
    TERRAIN_CONFIG.tileSize,
    TERRAIN_CONFIG.tileSize,
  );
  const tileData = { data, width, height };
  resolvedTileCache.set(cacheKey, tileData);
  return tileData;
}

function base64ToBlob(base64, contentType = "image/png") {
  const binary = atob(base64);
  const bytes = new Uint8Array(binary.length);
  for (let i = 0; i < binary.length; i++) {
    bytes[i] = binary.charCodeAt(i);
  }
  return new Blob([bytes], { type: contentType });
}

async function decodeTerrainTile(z, x, y) {
  const cacheKey = buildTerrainTileKey(z, x, y);
  if (terrainTileCache.has(cacheKey)) {
    return terrainTileCache.get(cacheKey);
  }

  const tilePromise = (async () => {
    const response = await fetch(buildTileUrl(z, x, y), { mode: "cors" });
    if (!response.ok) {
      throw new Error(`HTTP ${response.status} for terrain tile ${cacheKey}`);
    }

    return decodeTerrainBlob(await response.blob(), cacheKey);
  })();

  terrainTileCache.set(cacheKey, tilePromise);
  return tilePromise;
}

async function fetchTerrainTileBatch(keys) {
  const uncachedKeys = keys.filter((key) => !terrainTileCache.has(key));
  if (uncachedKeys.length === 0) {
    return;
  }

  if (!TERRAIN_CONFIG.batchUrl || typeof atob !== "function") {
    uncachedKeys.forEach((key) => {
      const [z, x, y] = key.split("/").map(Number);
      terrainTileCache.set(key, decodeTerrainTile(z, x, y));
    });
    return;
  }

  const batchPromise = (async () => {
    const response = await fetch(TERRAIN_CONFIG.batchUrl, {
      method: "POST",
      headers: {
        "Content-Type": "application/json",
      },
      body: JSON.stringify({
        tiles: uncachedKeys.map((key) => {
          const [z, x, y] = key.split("/").map(Number);
          return { z, x, y };
        }),
      }),
    });
    if (!response.ok) {
      throw new Error(`HTTP ${response.status} for terrain tile batch`);
    }

    const payload = await response.json();
    const tiles = Array.isArray(payload?.tiles) ? payload.tiles : [];
    const returnedKeys = new Set();

    await Promise.all(
      tiles.map(async (tile) => {
        const z = Number(tile?.z);
        const x = Number(tile?.x);
        const y = Number(tile?.y);
        const data = typeof tile?.data === "string" ? tile.data : "";
        if (!Number.isFinite(z) || !Number.isFinite(x) || !Number.isFinite(y) || !data) {
          return;
        }
        const key = buildTerrainTileKey(z, x, y);
        returnedKeys.add(key);
        const blob = base64ToBlob(data, tile.content_type || "image/png");
        terrainTileCache.set(key, decodeTerrainBlob(blob, key));
      }),
    );

    uncachedKeys.forEach((key) => {
      if (returnedKeys.has(key)) return;
      const [z, x, y] = key.split("/").map(Number);
      terrainTileCache.delete(key);
      terrainTileCache.set(key, decodeTerrainTile(z, x, y));
    });
  })().catch((error) => {
    console.warn("批量地形瓦片加载失败，回退到单瓦片请求:", error);
    uncachedKeys.forEach((key) => {
      terrainTileCache.delete(key);
      const [z, x, y] = key.split("/").map(Number);
      terrainTileCache.set(key, decodeTerrainTile(z, x, y));
    });
  });

  uncachedKeys.forEach((key) => {
    terrainTileCache.set(
      key,
      batchPromise.then(() => terrainTileCache.get(key)),
    );
  });

  await batchPromise;
}

function decodeTerrariumHeight(tile, pixelX, pixelY) {
  const safeX = THREE.MathUtils.clamp(pixelX, 0, tile.width - 1);
  const safeY = THREE.MathUtils.clamp(pixelY, 0, tile.height - 1);
  const index = (safeY * tile.width + safeX) * 4;
  const r = tile.data[index];
  const g = tile.data[index + 1];
  const b = tile.data[index + 2];
  return (r * 256 + g + b / 256) - 32768;
}

function latLonToTileSample(lat, lon, z, tileSize) {
  const n = 2 ** z;
  const clampedLat = clampLatitude(lat);
  const latRad = THREE.MathUtils.degToRad(clampedLat);
  const normalizedX = ((lon + 180) / 360) * n;
  const normalizedY =
    ((1 -
      Math.log(Math.tan(latRad) + 1 / Math.cos(latRad)) / Math.PI) /
      2) *
    n;

  const tileX = THREE.MathUtils.euclideanModulo(
    Math.floor(normalizedX),
    n,
  );
  const tileY = THREE.MathUtils.clamp(Math.floor(normalizedY), 0, n - 1);
  const pixelX = Math.floor((normalizedX - Math.floor(normalizedX)) * tileSize);
  const pixelY = Math.floor((normalizedY - Math.floor(normalizedY)) * tileSize);

  return {
    tileX,
    tileY,
    pixelX,
    pixelY,
  };
}

function buildTerrainVertexSamples(positionAttribute) {
  const z = TERRAIN_CONFIG.baseZoom;
  const tileSize = TERRAIN_CONFIG.tileSize;
  const samples = [];

  for (let i = 0; i < positionAttribute.count; i++) {
    const direction = new THREE.Vector3(
      positionAttribute.getX(i),
      positionAttribute.getY(i),
      positionAttribute.getZ(i),
    ).normalize();
    const { lat, lon } = vector3ToLatLon(direction);
    const sample = latLonToTileSample(lat, lon, z, tileSize);
    samples.push({
      index: i,
      direction,
      ...sample,
    });
  }

  return samples;
}

function sampleTerrainColor(heightMeters) {
  if (heightMeters <= TERRAIN_COLOR_STOPS[0].height) {
    return TERRAIN_COLOR_STOPS[0].color;
  }

  for (let i = 1; i < TERRAIN_COLOR_STOPS.length; i++) {
    const lower = TERRAIN_COLOR_STOPS[i - 1];
    const upper = TERRAIN_COLOR_STOPS[i];
    if (heightMeters <= upper.height) {
      const t =
        (heightMeters - lower.height) / Math.max(upper.height - lower.height, 1);
      return lower.color.clone().lerp(upper.color, t);
    }
  }

  return TERRAIN_COLOR_STOPS[TERRAIN_COLOR_STOPS.length - 1].color;
}

async function runWithConcurrency(items, limit, worker) {
  const queue = [...items];
  const workers = Array.from({ length: Math.min(limit, queue.length) }, async () => {
    while (queue.length > 0) {
      const item = queue.shift();
      await worker(item);
    }
  });
  await Promise.all(workers);
}

async function fetchRequiredTiles(samples) {
  const uniqueKeys = Array.from(
    new Set(
      samples.map((sample) =>
        buildTerrainTileKey(TERRAIN_CONFIG.baseZoom, sample.tileX, sample.tileY),
      ),
    ),
  );
  const resolvedTiles = new Map();
  const batchSize = Math.max(1, Number(TERRAIN_CONFIG.batchRequestSize) || 1);
  const batchChunks = [];
  for (let i = 0; i < uniqueKeys.length; i += batchSize) {
    batchChunks.push(uniqueKeys.slice(i, i + batchSize));
  }

  await runWithConcurrency(
    batchChunks,
    TERRAIN_CONFIG.maxConcurrentRequests,
    async (keys) => {
      await fetchTerrainTileBatch(keys);
    },
  );

  await Promise.all(
    uniqueKeys.map(async (key) => {
      const tilePromise = terrainTileCache.get(key);
      if (!tilePromise) {
        const [z, x, y] = key.split("/").map(Number);
        resolvedTiles.set(key, await decodeTerrainTile(z, x, y));
        return;
      }
      resolvedTiles.set(key, await tilePromise);
    }),
  );

  return resolvedTiles;
}

function applyTerrainDisplacement(samples, mesh, resolvedTiles) {
  const geometry = mesh.geometry;
  const positionAttribute = geometry.getAttribute("position");
  let colorAttribute = geometry.getAttribute("color");
  if (!colorAttribute || colorAttribute.itemSize !== 4) {
    colorAttribute = new THREE.BufferAttribute(
      new Float32Array(positionAttribute.count * 4),
      4,
    );
    geometry.setAttribute("color", colorAttribute);
  }
  const baseRadius = CONFIG.earthRadius + TERRAIN_CONFIG.baseRadiusOffset;

  samples.forEach((sample) => {
    const tileKey = `${TERRAIN_CONFIG.baseZoom}/${sample.tileX}/${sample.tileY}`;
    const tile = resolvedTiles.get(tileKey);
    if (!tile) return;
    const rawElevationMeters = decodeTerrariumHeight(
      tile,
      sample.pixelX,
      sample.pixelY,
    );
    const elevationMeters = Math.max(0, rawElevationMeters);
    const heightWorld =
      (elevationMeters / EARTH_RADIUS_METERS) *
      CONFIG.earthRadius *
      TERRAIN_CONFIG.exaggeration;
    const radius = baseRadius + heightWorld;
    const tint = sampleTerrainColor(elevationMeters);
    const landAlpha = THREE.MathUtils.clamp(
      elevationMeters / Math.max(TERRAIN_CONFIG.landRevealFadeMeters, 1),
      0,
      1,
    );
    positionAttribute.setXYZ(
      sample.index,
      sample.direction.x * radius,
      sample.direction.y * radius,
      sample.direction.z * radius,
    );
    colorAttribute.setXYZW(sample.index, tint.r, tint.g, tint.b, landAlpha);
  });

  positionAttribute.needsUpdate = true;
  colorAttribute.needsUpdate = true;
  geometry.computeVertexNormals();
  geometry.computeBoundingSphere();
}

export function registerTerrainMesh(mesh) {
  terrainMesh = mesh;
  terrainReady = false;
  terrainFailed = false;
  terrainLoadPromise = null;
  terrainTileCache = new Map();
  resolvedTileCache = new Map();
  terrainOpacity = TERRAIN_CONFIG.opacity;
  if (terrainMesh?.material) {
    terrainMesh.material.opacity = terrainOpacity;
    terrainMesh.material.needsUpdate = true;
  }
  terrainVertexSamples = mesh
    ? buildTerrainVertexSamples(mesh.geometry.getAttribute("position"))
    : null;
}

export function isTerrainReady() {
  return terrainReady;
}

export async function ensureTerrainReady() {
  if (!terrainMesh || !TERRAIN_CONFIG.enabled) {
    return false;
  }
  if (terrainReady) {
    return true;
  }
  if (terrainLoadPromise) {
    return terrainLoadPromise;
  }

  terrainLoadPromise = (async () => {
    try {
      const resolvedTiles = await fetchRequiredTiles(terrainVertexSamples);
      applyTerrainDisplacement(terrainVertexSamples, terrainMesh, resolvedTiles);
      terrainReady = true;
      terrainFailed = false;
      return true;
    } catch (error) {
      terrainFailed = true;
      console.error("加载真实地形失败:", error);
      throw error;
    } finally {
      terrainLoadPromise = null;
    }
  })();

  return terrainLoadPromise;
}

export function clearTerrainData() {
  terrainMesh = null;
  terrainLoadPromise = null;
  terrainReady = false;
  terrainFailed = false;
  terrainVertexSamples = null;
  terrainTileCache = new Map();
  resolvedTileCache = new Map();
  terrainOpacity = TERRAIN_CONFIG.opacity;
}

export function sampleElevationAt(lat, lon) {
  if (!terrainReady) return null;
  const z = TERRAIN_CONFIG.baseZoom;
  const { tileX, tileY, pixelX, pixelY } = latLonToTileSample(lat, lon, z, TERRAIN_CONFIG.tileSize);
  const tile = resolvedTileCache.get(`${z}/${tileX}/${tileY}`);
  if (!tile) return null;
  return Math.max(0, decodeTerrariumHeight(tile, pixelX, pixelY));
}

export function setTerrainOpacity(nextOpacity) {
  terrainOpacity = THREE.MathUtils.clamp(nextOpacity, 0.05, 1);
  if (terrainMesh?.material) {
    terrainMesh.material.opacity = terrainOpacity;
    terrainMesh.material.needsUpdate = true;
  }
  return terrainOpacity;
}

export function getTerrainOpacity() {
  return terrainOpacity;
}