import { add, mul, sub, dot, cross, unit, point } from "./math.js";

export function triangleHit(mesh, f, origin, direction) {
  const ids = mesh.indices.subarray(f * 3, f * 3 + 3);
  const a = point(mesh.positions, ids[0]);
  const e1 = sub(point(mesh.positions, ids[1]), a),
    e2 = sub(point(mesh.positions, ids[2]), a);
  const h = cross(direction, e2),
    det = dot(e1, h);
  if (Math.abs(det) < 1e-9) return null;
  const s = sub(origin, a),
    u = dot(s, h) / det;
  if (u < 0 || u > 1) return null;
  const q = cross(s, e1),
    v = dot(direction, q) / det;
  if (v < 0 || u + v > 1) return null;
  const t = dot(e2, q) / det;
  return t > 0
    ? {
        t,
        position: add(origin, mul(direction, t)),
        normal: unit(cross(e1, e2)),
        face: f,
      }
    : null;
}

export class BVH {
  constructor(mesh) {
    this.mesh = mesh;
    this.leaves = new Array(mesh.indices.length / 3);
    const centers = Array.from({ length: this.leaves.length }, (_, f) =>
      [0, 1, 2].map(
        (k) =>
          (mesh.positions[mesh.indices[f * 3] * 3 + k] +
            mesh.positions[mesh.indices[f * 3 + 1] * 3 + k] +
            mesh.positions[mesh.indices[f * 3 + 2] * 3 + k]) /
          3,
      ),
    );
    const build = (faces, parent = null) => {
      const node = { parent, min: [0, 0, 0], max: [0, 0, 0] };
      if (faces.length <= 8) {
        node.faces = faces;
        for (const f of faces) this.leaves[f] = node;
      } else {
        const spans = [0, 1, 2].map((k) => {
          let lo = Infinity,
            hi = -Infinity;
          for (const f of faces) {
            lo = Math.min(lo, centers[f][k]);
            hi = Math.max(hi, centers[f][k]);
          }
          return hi - lo;
        });
        const axis = spans.indexOf(Math.max(...spans));
        faces.sort((a, b) => centers[a][axis] - centers[b][axis]);
        const mid = faces.length >> 1;
        node.left = build(faces.slice(0, mid), node);
        node.right = build(faces.slice(mid), node);
      }
      this.bounds(node);
      return node;
    };
    this.root = this.leaves.length
      ? build(Array.from({ length: this.leaves.length }, (_, i) => i))
      : null;
    this.marks = new Uint32Array(mesh.positions.length / 3);
    this.generation = 0;
  }
  bounds(node) {
    for (let k = 0; k < 3; k++) {
      let lo = Infinity,
        hi = -Infinity;
      if (node.faces)
        for (const f of node.faces)
          for (let j = 0; j < 3; j++) {
            const v = this.mesh.positions[this.mesh.indices[f * 3 + j] * 3 + k];
            lo = Math.min(lo, v);
            hi = Math.max(hi, v);
          }
      else {
        lo = Math.min(node.left.min[k], node.right.min[k]);
        hi = Math.max(node.left.max[k], node.right.max[k]);
      }
      node.min[k] = lo;
      node.max[k] = hi;
    }
  }
  refit(faces) {
    if (!this.root) return;
    const dirty = new Set(faces.map((f) => this.leaves[f]));
    for (const leaf of dirty) this.bounds(leaf);
    let parents = new Set([...dirty].map((n) => n.parent).filter(Boolean));
    while (parents.size) {
      const next = new Set();
      for (const n of parents) {
        this.bounds(n);
        if (n.parent) next.add(n.parent);
      }
      parents = next;
    }
  }
  raycast(origin, direction) {
    if (!this.root) return null;
    let best = null;
    const entry = (n) => {
      let near = 0,
        far = best?.t ?? Infinity;
      for (let k = 0; k < 3; k++) {
        if (Math.abs(direction[k]) < 1e-12) {
          if (origin[k] < n.min[k] || origin[k] > n.max[k]) return Infinity;
          continue;
        }
        const a = (n.min[k] - origin[k]) / direction[k],
          b = (n.max[k] - origin[k]) / direction[k];
        near = Math.max(near, Math.min(a, b));
        far = Math.min(far, Math.max(a, b));
        if (near > far) return Infinity;
      }
      return near;
    };
    const visit = (n) => {
      if (entry(n) === Infinity) return;
      if (n.faces) {
        for (const f of n.faces) {
          const hit = triangleHit(this.mesh, f, origin, direction);
          if (hit && (!best || hit.t < best.t)) best = hit;
        }
      } else if (entry(n.left) < entry(n.right)) {
        visit(n.left);
        visit(n.right);
      } else {
        visit(n.right);
        visit(n.left);
      }
    };
    visit(this.root);
    return best;
  }
  within(center, radius) {
    if (!this.root) return [];
    const result = [],
      gen = ++this.generation,
      r2 = radius * radius;
    const visit = (n) => {
      let d = 0;
      for (let k = 0; k < 3; k++) {
        const e = Math.max(n.min[k] - center[k], 0, center[k] - n.max[k]);
        d += e * e;
      }
      if (d > r2) return;
      if (n.faces)
        for (const f of n.faces)
          for (let j = 0; j < 3; j++) {
            const v = this.mesh.indices[f * 3 + j];
            if (this.marks[v] === gen) continue;
            this.marks[v] = gen;
            const p = this.mesh.positions;
            const d2 =
              (p[v * 3] - center[0]) ** 2 +
              (p[v * 3 + 1] - center[1]) ** 2 +
              (p[v * 3 + 2] - center[2]) ** 2;
            if (d2 < r2) result.push(v);
          }
      else {
        visit(n.left);
        visit(n.right);
      }
    };
    visit(this.root);
    return result;
  }
}