import * as THREE from "three";

import type { PositionFrame } from "./ws_client";

/**
 * Renders synapses (Stage 3) as additive line segments connecting two
 * engrams. Endpoints are looked up from the engram position buffer on every
 * position frame, so the lines track engram motion automatically.
 *
 * Uses a fixed-capacity vertex buffer; one segment per synapse → 2 vertices
 * per synapse → 6 floats of position per synapse.
 */
const VERT_SHADER = /* glsl */ `
  attribute float aAlpha;
  varying vec3 vColor;
  varying float vAlpha;
  void main() {
    vColor = color;
    vAlpha = aAlpha;
    gl_Position = projectionMatrix * modelViewMatrix * vec4(position, 1.0);
  }
`;

const FRAG_SHADER = /* glsl */ `
  varying vec3 vColor;
  varying float vAlpha;
  void main() {
    gl_FragColor = vec4(vColor, vAlpha);
  }
`;

type SynapseEntry = {
  /** Slot in the line-segments geometry (0..capacity-1). */
  slot: number;
  /** instance_idx of the two engrams this connects. */
  aIdx: number;
  bIdx: number;
  weight: number;
};

type Pending = {
  id: string;
  a: string;
  b: string;
  weight: number;
};

export class SynapseMesh {
  private readonly mesh: THREE.LineSegments;
  private readonly positionAttr: THREE.BufferAttribute;
  private readonly colorAttr: THREE.BufferAttribute;
  private readonly alphaAttr: THREE.BufferAttribute;
  private readonly capacity: number;
  /** synapse_id → entry. */
  private readonly bySynapseId = new Map<string, SynapseEntry>();
  /** Engram UUID → instance_idx, populated as engrams arrive. */
  private readonly engramIdx = new Map<string, number>();
  /** Synapses waiting for one of their endpoints to be registered. */
  private readonly pending: Pending[] = [];
  /** Engram colour cache so we don't recompute on every frame. */
  private readonly engramColor: Float32Array;
  private nextSlot = 0;

  constructor(scene: THREE.Scene, capacity = 8000) {
    this.capacity = capacity;
    this.engramColor = new Float32Array(5000 * 3);

    const vertexCount = capacity * 2;
    const geom = new THREE.BufferGeometry();
    this.positionAttr = new THREE.BufferAttribute(new Float32Array(vertexCount * 3), 3);
    this.colorAttr = new THREE.BufferAttribute(new Float32Array(vertexCount * 3), 3);
    this.alphaAttr = new THREE.BufferAttribute(new Float32Array(vertexCount), 1);
    this.positionAttr.setUsage(THREE.DynamicDrawUsage);
    this.colorAttr.setUsage(THREE.DynamicDrawUsage);
    this.alphaAttr.setUsage(THREE.DynamicDrawUsage);
    geom.setAttribute("position", this.positionAttr);
    geom.setAttribute("color", this.colorAttr);
    geom.setAttribute("aAlpha", this.alphaAttr);
    geom.setDrawRange(0, 0);

    const mat = new THREE.ShaderMaterial({
      vertexShader: VERT_SHADER,
      fragmentShader: FRAG_SHADER,
      vertexColors: true,
      transparent: true,
      depthWrite: false,
      depthTest: true,
      blending: THREE.AdditiveBlending,
    });

    this.mesh = new THREE.LineSegments(geom, mat);
    this.mesh.frustumCulled = false;
    scene.add(this.mesh);
  }

  /** Register an engram so synapses referencing it can be wired up. */
  registerEngram(id: string, instanceIdx: number, r: number, g: number, b: number): void {
    this.engramIdx.set(id, instanceIdx);
    if (instanceIdx * 3 + 2 < this.engramColor.length) {
      this.engramColor[instanceIdx * 3] = r;
      this.engramColor[instanceIdx * 3 + 1] = g;
      this.engramColor[instanceIdx * 3 + 2] = b;
    }
    // Try to materialise any pending synapses now that this engram is known.
    if (this.pending.length > 0) {
      const stillPending: Pending[] = [];
      for (const p of this.pending) {
        if (!this.tryMaterialise(p)) stillPending.push(p);
      }
      this.pending.length = 0;
      this.pending.push(...stillPending);
    }
  }

  /** Add a synapse by engram UUIDs. Defers if either endpoint is unknown. */
  addSynapse(id: string, a: string, b: string, weight: number): void {
    if (this.bySynapseId.has(id)) return;
    const p: Pending = { id, a, b, weight };
    if (!this.tryMaterialise(p)) {
      this.pending.push(p);
    }
  }

  private tryMaterialise(p: Pending): boolean {
    const aIdx = this.engramIdx.get(p.a);
    const bIdx = this.engramIdx.get(p.b);
    if (aIdx === undefined || bIdx === undefined) return false;
    if (this.nextSlot >= this.capacity) {
      console.warn("SynapseMesh capacity reached; ignoring further synapses");
      return true; // treat as resolved so we stop waiting on it
    }
    const slot = this.nextSlot++;
    this.bySynapseId.set(p.id, { slot, aIdx, bIdx, weight: p.weight });
    this.applyEndpointColors(slot, aIdx, bIdx);
    // Alpha tied to weight; a small floor so very weak synapses still register.
    const alpha = Math.max(0.08, Math.min(0.75, p.weight));
    const alphaArr = this.alphaAttr.array as Float32Array;
    alphaArr[slot * 2] = alpha;
    alphaArr[slot * 2 + 1] = alpha;
    this.alphaAttr.needsUpdate = true;
    this.mesh.geometry.setDrawRange(0, this.nextSlot * 2);
    return true;
  }

  /** Apply the latest position frame to all synapse endpoints. */
  applyPositionFrame(frame: PositionFrame): void {
    const posArr = this.positionAttr.array as Float32Array;
    const src = frame.positions;
    const n = frame.n;
    let dirty = false;
    for (const entry of this.bySynapseId.values()) {
      const a = entry.aIdx, b = entry.bIdx;
      if (a >= n || b >= n) continue;
      const slot = entry.slot;
      posArr[slot * 6]     = src[a * 3];
      posArr[slot * 6 + 1] = src[a * 3 + 1];
      posArr[slot * 6 + 2] = src[a * 3 + 2];
      posArr[slot * 6 + 3] = src[b * 3];
      posArr[slot * 6 + 4] = src[b * 3 + 1];
      posArr[slot * 6 + 5] = src[b * 3 + 2];
      dirty = true;
    }
    if (dirty) this.positionAttr.needsUpdate = true;
  }

  /** Number of synapses currently rendered (for HUD). */
  count(): number {
    return this.bySynapseId.size;
  }

  /** Re-paint a synapse's endpoint colours after one of its engrams updates. */
  private applyEndpointColors(slot: number, aIdx: number, bIdx: number): void {
    const colorArr = this.colorAttr.array as Float32Array;
    colorArr[slot * 6]     = this.engramColor[aIdx * 3];
    colorArr[slot * 6 + 1] = this.engramColor[aIdx * 3 + 1];
    colorArr[slot * 6 + 2] = this.engramColor[aIdx * 3 + 2];
    colorArr[slot * 6 + 3] = this.engramColor[bIdx * 3];
    colorArr[slot * 6 + 4] = this.engramColor[bIdx * 3 + 1];
    colorArr[slot * 6 + 5] = this.engramColor[bIdx * 3 + 2];
    this.colorAttr.needsUpdate = true;
  }
}