/**
 * LANE-RENDER — belts, instanced, one InstancedMesh per (belt def x topology shape).
 *
 * The chevron scroll is injected into a MeshStandardMaterial rather than written as a
 * raw ShaderMaterial, so belts keep real lighting, fog and brownout dimming for free.
 * The shape's flow field is baked into the shader source at material-build time — no
 * runtime branch, and each shape is its own draw call anyway.
 *
 * Straight chevrons run along local -z; corners sweep the quarter-arc from the feeding
 * edge to the output edge; merges add inlet arrows flowing in from the sides. The
 * instance yaw carries all of it to world orientation.
 */
import * as THREE from 'three';
import type { MachineDef, SimSnapshot } from '../contracts';
import { AssetRegistry } from './registry';
import { yawFor } from './coords';
import { ACCENT_BY_KIND, BODY_BY_KIND, HEIGHT_BY_KIND } from './palette';
import { ARC_LEN, BELT_SHAPES, BeltTopology, type BeltShape } from './topology';

const CHEVRONS_PER_TILE = 3;

interface BeltMesh {
  mesh: THREE.InstancedMesh;
  capacity: number;
  assetVersion: number;
}

/** The flow field for one shape, in local tile space. Baked, not branched. */
function chevronGLSL(shape: BeltShape): string {
  const N = CHEVRONS_PER_TILE.toFixed(1);
  if (shape === 'cornerA' || shape === 'cornerB') {
    const side = shape === 'cornerA' ? '1.0' : '-1.0';
    return `
      vec2 rel = vec2(vLocal.x, vLocal.z) - vec2(${side} * 0.5, -0.5);
      float r = length(rel);
      float u = ${side} * (atan(rel.y, rel.x) - 1.5707963) / 1.5707963;
      float band = fract((u * ${ARC_LEN.toFixed(4)} + abs(r - 0.5) * 0.55) * ${N}
                         - uTime * uSpeed * ${N});
      float arrow = (smoothstep(0.40, 0.50, band) - smoothstep(0.50, 0.60, band))
                  * (1.0 - smoothstep(0.85, 1.15, r));
    `;
  }
  if (shape === 'merge') {
    return `
      float band = fract((vLocal.z + abs(vLocal.x) * 0.55) * ${N} + uTime * uSpeed * ${N});
      float arrow = smoothstep(0.40, 0.50, band) - smoothstep(0.50, 0.60, band);
      float sb = fract(abs(vLocal.x) * ${N} + uTime * uSpeed * ${N});
      arrow += (smoothstep(0.40, 0.50, sb) - smoothstep(0.50, 0.60, sb))
             * step(0.26, abs(vLocal.x)) * 0.8;
    `;
  }
  return `
    float band = fract((vLocal.z + abs(vLocal.x) * 0.55) * ${N} + uTime * uSpeed * ${N});
    float arrow = smoothstep(0.40, 0.50, band) - smoothstep(0.50, 0.60, band);
  `;
}

export class BeltLayer {
  readonly group = new THREE.Group();
  private meshes = new Map<string, BeltMesh>();
  private uTime = { value: 0 };
  private dummy = new THREE.Object3D();
  private builtVersion = -1;

  constructor(
    private registry: AssetRegistry,
    private defs: Map<string, MachineDef>,
    private topo: BeltTopology,
  ) {
    this.group.name = 'belts';
  }

  private material(def: MachineDef, shape: BeltShape): THREE.Material {
    // Chassis colour from data when LANE-DATA has art-directed it (see registry.bodyFor).
    const mat = new THREE.MeshStandardMaterial({
      color: def.color ? new THREE.Color(def.color) : new THREE.Color(BODY_BY_KIND.belt),
      roughness: 0.9,
      metalness: 0.1,
    });
    const uSpeed = { value: def.beltSpeed ?? 2 };
    const uAccent = { value: new THREE.Color(ACCENT_BY_KIND.belt) };
    const topY = (HEIGHT_BY_KIND.belt * 0.92).toFixed(4);
    mat.onBeforeCompile = (shader) => {
      shader.uniforms.uTime = this.uTime;
      shader.uniforms.uSpeed = uSpeed;
      shader.uniforms.uAccent = uAccent;
      shader.vertexShader =
        'varying vec3 vLocal;\n' +
        shader.vertexShader.replace('#include <begin_vertex>', '#include <begin_vertex>\n  vLocal = position;');
      shader.fragmentShader =
        'uniform float uTime;\nuniform float uSpeed;\nuniform vec3 uAccent;\nvarying vec3 vLocal;\n' +
        shader.fragmentShader.replace(
          '#include <emissivemap_fragment>',
          `#include <emissivemap_fragment>
           if (vLocal.y > ${topY}) {
             ${chevronGLSL(shape)}
             totalEmissiveRadiance += uAccent * arrow * 0.85;
           }`,
        );
    };
    // Distinct cache key per shape: same base source, different injected flow field.
    mat.customProgramCacheKey = () => `fktry-belt-${shape}`;
    return mat;
  }

  private ensure(def: MachineDef, shape: BeltShape, needed: number): BeltMesh {
    const key = `${def.id}|${shape}`;
    const entry = this.registry.get(def.asset);
    const version = entry?.version ?? 0;
    let bm = this.meshes.get(key);

    if (bm && bm.assetVersion !== version) {
      this.group.remove(bm.mesh);
      bm.mesh.dispose();
      this.meshes.delete(key);
      bm = undefined; // GLB landed for belts: rebuild from it
    }
    if (bm && needed <= bm.capacity) return bm;

    const capacity = Math.max(64, 1 << Math.ceil(Math.log2(Math.max(1, needed))));
    const inst = entry?.instanceable() ?? null;
    const geometry = inst?.geometry ?? new THREE.BoxGeometry(0.94, HEIGHT_BY_KIND.belt, 0.94);
    if (!inst) geometry.translate(0, HEIGHT_BY_KIND.belt / 2, 0);
    // A real GLB brings its own look; the placeholder plate gets the chevron scroll.
    // (Taking inst.material unconditionally here silently dropped the chevrons — the
    // registry hands back a plain plate material for the placeholder case.)
    const material = entry?.isGLB && inst ? inst.material : this.material(def, shape);

    if (bm) {
      this.group.remove(bm.mesh);
      bm.mesh.dispose(); // grow: drop the old instanced mesh, keep geo/mat
    }
    const mesh = new THREE.InstancedMesh(geometry, material, capacity);
    mesh.instanceMatrix.setUsage(THREE.DynamicDrawUsage);
    mesh.castShadow = mesh.receiveShadow = true;
    mesh.frustumCulled = false;
    mesh.count = 0;
    this.group.add(mesh);
    const next: BeltMesh = { mesh, capacity, assetVersion: version };
    this.meshes.set(key, next);
    return next;
  }

  /** `topo` is rebuilt by index.ts once per frame; we only react when it moved. */
  sync(_snap: SimSnapshot, timeSec: number): void {
    this.uTime.value = timeSec;
    const assetVersion = this.registry.version;
    const stamp = this.topo.version * 1000 + assetVersion;
    if (stamp === this.builtVersion) return; // belts are static once placed
    this.builtVersion = stamp;

    for (const def of this.defs.values()) {
      if (def.kind !== 'belt') continue;
      for (const shape of BELT_SHAPES) {
        const nodes = this.topo.nodesOf(def.id, shape);
        const bm = this.ensure(def, shape, nodes.length);
        nodes.forEach((n, i) => {
          this.dummy.position.set(n.x + 0.5, 0, n.y + 0.5); // 1x1 footprint: tile centre
          this.dummy.rotation.set(0, yawFor(n.dir), 0);
          this.dummy.updateMatrix();
          bm.mesh.setMatrixAt(i, this.dummy.matrix);
        });
        bm.mesh.count = nodes.length;
        bm.mesh.instanceMatrix.needsUpdate = true;
        bm.mesh.computeBoundingSphere();
      }
    }
  }
}