Made of little things.

// Bake the level's lighting into per-vertex colours: a warm sun through the
// window (with shadows), a cool sky term weighted by ambient occlusion, and a
// weak unshadowed fill. The game's shader uses these instead of its dynamic
// lights for baked meshes, the way the original multiplied in lightmaps.
import { BufferAttribute, BufferGeometry, DoubleSide, Ray, Vector3 } from 'three';
import { MeshBVH } from 'three-mesh-bvh';
import type { Piece } from './mesh.ts';
import { rng } from './textures.ts';

/** Direction sunlight travels: from high in the south-west, down into the room. */
const SUN = new Vector3(0.45, 0.55, -0.7).normalize();
const SUN_COLOR = [1.0, 0.86, 0.68];
const SKY_COLOR = [0.42, 0.48, 0.6];
const FILL = new Vector3(-0.5, -0.3, -0.8).normalize();
const FILL_COLOR = [0.14, 0.14, 0.16];
/** Warm light bounced off the floor onto downward faces (the ceiling, bridge deck). */
const BOUNCE_COLOR = [0.42, 0.34, 0.26];
const AO_SAMPLES = 40;
const AO_RANGE = 700;

export function bakeLevel(pieces: Piece[]): Map<string, Uint8Array> {
  // One BVH over every visual triangle: decor casts shadows too.
  const pos: number[] = [];
  const idx: number[] = [];
  for (const p of pieces.filter((x) => x.castsShadow)) {
    for (const pr of p.prims.values()) {
      const base = pos.length / 3;
      pos.push(...pr.positions);
      for (const i of pr.indices) idx.push(base + i);
    }
  }
  const geometry = new BufferGeometry();
  geometry.setAttribute('position', new BufferAttribute(new Float32Array(pos), 3));
  geometry.setIndex(idx);
  const bvh = new MeshBVH(geometry);

  const ray = new Ray();
  const occluded = (o: Vector3, d: Vector3, far: number) => {
    ray.origin.copy(o);
    ray.direction.copy(d);
    const hit = bvh.raycastFirst(ray, DoubleSide);
    return !!hit && hit.distance < far;
  };

  // Cosine-weighted hemisphere directions (tangent space), rotated per vertex.
  const random = rng(99);
  const dirs = Array.from({ length: AO_SAMPLES }, (_, i) => {
    const u = (i + random()) / AO_SAMPLES;
    const a = i * 2.39996323;
    const r = Math.sqrt(u);
    return [Math.cos(a) * r, Math.sin(a) * r, Math.sqrt(1 - u)];
  });

  const out = new Map<string, Uint8Array>();
  const n = new Vector3();
  const o = new Vector3();
  const t = new Vector3();
  const b = new Vector3();
  const d = new Vector3();
  const sunTo = SUN.clone().negate();
  const fillTo = FILL.clone().negate();
  for (const p of pieces) {
    for (const pr of p.prims.values()) {
      const colors = new Uint8Array((pr.positions.length / 3) * 4);
      for (let v = 0; v < pr.positions.length / 3; v++) {
        n.fromArray(pr.normals, v * 3);
        // Start just off the surface so it does not shadow itself.
        o.fromArray(pr.positions, v * 3).addScaledVector(n, 1.5);
        if (Math.abs(n.z) < 0.9) t.set(0, 0, 1);
        else t.set(1, 0, 0);
        t.cross(n).normalize();
        b.crossVectors(n, t);
        const spin = (v * 0.618034) % 1 * Math.PI * 2;
        const cs = Math.cos(spin);
        const sn = Math.sin(spin);
        let open = 0;
        for (const [x0, y0, z] of dirs) {
          const x = x0 * cs - y0 * sn;
          const y = x0 * sn + y0 * cs;
          d.copy(t).multiplyScalar(x).addScaledVector(b, y).addScaledVector(n, z);
          if (!occluded(o, d, AO_RANGE)) open++;
        }
        const ao = open / AO_SAMPLES;
        const sunDot = Math.max(0, n.dot(sunTo));
        const sun = sunDot > 0 && !occluded(o, sunTo, 1e5) ? sunDot : 0;
        const fill = Math.max(0, n.dot(fillTo));
        const sky = 0.35 + 0.65 * ao;
        for (let k = 0; k < 3; k++) {
          const light =
            SKY_COLOR[k] * sky * (0.75 + 0.25 * Math.max(0, n.z)) +
            SUN_COLOR[k] * sun * (0.6 + 0.4 * ao) +
            FILL_COLOR[k] * fill +
            BOUNCE_COLOR[k] * Math.max(0, -n.z) * ao;
          colors[v * 4 + k] = Math.round(Math.min(1, light) * 255);
        }
        colors[v * 4 + 3] = 255;
      }
      out.set(`${p.name}/${pr.material}`, colors);
    }
  }
  return out;
}