import { sphere, updateNormals } from "./mesh.js";
import { BVH } from "./bvh.js";
import { point, sub, dot, clamp } from "./math.js";
import {
  ensureDocument,
  snapshot,
  hydrate,
  saveDocument,
  loadDocument,
} from "./document.js";

export const pressureValue = (pressure, gain = 1, gamma = 1) =>
  clamp(gain * clamp(pressure, 0, 1) ** gamma, 0, 1);
export const defaults = {
  brush: "draw",
  radius: 0.22,
  strength: 0.45,
  gain: 1,
  gamma: 1,
  symmetry: true,
  invert: false,
  pressureRadius: false,
};

const persistedJson = (value) =>
  JSON.stringify(value, (_key, child) => {
    if (!child || Array.isArray(child) || typeof child !== "object")
      return child;
    return Object.fromEntries(
      Object.keys(child)
        .sort()
        .filter((key) => child[key] !== undefined)
        .map((key) => [key, child[key]]),
    );
  });
const canonicalEqual = (a, b) => {
  if (
    a.level !== b.level ||
    a.positions.length !== b.positions.length ||
    a.indices.length !== b.indices.length ||
    a.masks.length !== b.masks.length ||
    a.pieces.length !== b.pieces.length
  )
    return false;
  for (const name of ["positions", "indices", "masks"])
    for (let i = 0; i < a[name].length; i++)
      if (a[name][i] !== b[name][i]) return false;
  return persistedJson(a.pieces) === persistedJson(b.pieces);
};
const snapshotBytes = (value) =>
  value.positions.byteLength +
  value.indices.byteLength +
  value.masks.byteLength +
  8 +
  new TextEncoder().encode(JSON.stringify(value.pieces)).byteLength;
const allDirty = (mesh) =>
  new Set(
    Array.from({ length: mesh.positions.length / 3 }, (_, vertex) => vertex),
  );

export class Sculpt {
  constructor(mesh = sphere(), historyCap = 16 * 1024 * 1024) {
    this.mesh = ensureDocument(mesh);
    this.bvh = new BVH(mesh);
    this.undoStack = [];
    this.redoStack = [];
    this.historyCap = historyCap;
    this.revision = 0;
    this.state = 0;
    this.nextState = 1;
    this.activeId = this.mesh.pieces[0]?.id ?? null;
    this.dirty = new Set();
    this.markAll();
    this.stroke = null;
  }
  markAll() {
    for (let v = 0; v < this.mesh.positions.length / 3; v++) this.dirty.add(v);
    this.revision++;
  }
  get historyBytes() {
    return [...this.undoStack, ...this.redoStack].reduce(
      (sum, e) => sum + (e.bytes ?? 0),
      0,
    );
  }
  commit(entry) {
    entry.beforeState = this.state;
    entry.afterState = this.nextState++;
    this.state = entry.afterState;
    this.redoStack = [];
    this.undoStack.push(entry);
    while (this.undoStack.length > 50 || this.historyBytes > this.historyCap)
      this.undoStack.shift();
  }
  begin(settings) {
    if (this.stroke) throw new Error("A stroke is already active.");
    const requested = settings?.activeId;
    if (requested && this.mesh.pieces.some((p) => p.id === requested))
      this.activeId = requested;
    const piece = this.mesh.pieces.find((p) => p.id === this.activeId);
    this.stroke = {
      settings: { ...defaults, ...settings },
      piece,
      before: this.mesh.positions.slice(),
      beforeMasks: this.mesh.masks.slice(),
      touched: new Set(),
      grab: null,
    };
  }
  changed(ids) {
    if (!ids.length) return;
    const { faces, vertices } = updateNormals(this.mesh, ids);
    this.bvh.refit(faces);
    for (const v of vertices) this.dirty.add(v);
    this.revision++;
  }
  dab(hit, rawPressure = 1, drag = [0, 0, 0]) {
    if (!this.stroke || !hit) return;
    const { settings: s, piece } = this.stroke;
    if (
      !piece ||
      (hit.face != null &&
        (hit.face * 3 < piece.indexStart ||
          hit.face * 3 >= piece.indexStart + piece.indexCount))
    )
      return;
    const pressure = pressureValue(rawPressure, s.gain, s.gamma);
    if (pressure === 0) return;
    const radius = s.radius * (s.pressureRadius ? 0.25 + 0.75 * pressure : 1);
    const p = this.mesh.positions,
      n = this.mesh.normals;
    // Merge reflected influence instead of applying two sequential dabs at the seam.
    let affected = this.stroke.grab;
    if (!affected) {
      affected = new Map();
      for (const mirror of s.symmetry ? [1, -1] : [1]) {
        const center = [
          piece.symmetryX + (hit.position[0] - piece.symmetryX) * mirror,
          hit.position[1],
          hit.position[2],
        ];
        const normal = [hit.normal[0] * mirror, hit.normal[1], hit.normal[2]];
        for (const v of this.bvh.within(center, radius)) {
          if (
            v < piece.vertexStart ||
            v >= piece.vertexStart + piece.vertexCount
          )
            continue;
          if (dot(point(n, v), normal) < 0.05) continue;
          const delta = sub(point(p, v), center),
            weight = (1 - dot(delta, delta) / (radius * radius)) ** 2;
          if (!affected.has(v)) affected.set(v, []);
          affected.get(v).push({ weight, center, normal, mirror });
        }
      }
      if (s.brush === "grab") this.stroke.grab = affected;
    }
    const moved = [],
      updates = [];
    for (const [v, contributions] of affected) {
      const position = point(p, v),
        displacement = [0, 0, 0];
      let total = 0,
        envelope = 0;
      for (const { weight, center, normal, mirror } of contributions) {
        total += weight;
        envelope = Math.max(envelope, weight);
        let delta;
        if (s.brush === "smooth") {
          const neighbors = this.mesh.neighbors[v];
          delta = [0, 1, 2].map((k) => {
            let sum = 0;
            for (const neighbor of neighbors) sum += p[neighbor * 3 + k];
            return sum / neighbors.length - position[k];
          });
        } else if (s.brush === "flatten") {
          const distance = dot(sub(position, center), normal);
          delta = normal.map((x) => -x * distance);
        } else if (s.brush === "grab") {
          delta = drag.map((x, k) => x * (k === 0 ? mirror : 1));
        } else {
          const direction = s.brush === "inflate" ? point(n, v) : normal;
          delta = direction.map((x) => x * radius * 0.12 * (s.invert ? -1 : 1));
        }
        for (let k = 0; k < 3; k++) displacement[k] += delta[k] * weight;
      }
      if (s.brush === "mask") {
        const old = this.mesh.masks[v];
        const direction = s.invert ? -1 : 1;
        const next = clamp(
          old + direction * s.strength * pressure * envelope,
          0,
          1,
        );
        if (Math.fround(next) !== old) {
          this.mesh.masks[v] = next;
          this.stroke.touched.add(v);
          this.dirty.add(v);
        }
        continue;
      }
      const base = s.brush === "grab" ? point(this.stroke.before, v) : position;
      let next = base.map(
        (x, k) =>
          x +
          (displacement[k] *
            s.strength *
            pressure *
            envelope *
            (1 - this.mesh.masks[v])) /
            total,
      );
      if (s.symmetry && this.stroke.before[v * 3] === piece.symmetryX)
        next[0] = piece.symmetryX;
      // Bounded document coordinates keep camera and project validation meaningful.
      next = next.map((x) => clamp(x, -8, 8));
      if (next.some((x, k) => Math.fround(x) !== position[k])) {
        moved.push(v);
        updates.push(...next);
        this.stroke.touched.add(v);
      }
    }
    // Jacobi update: smoothing reads only the pre-dab mesh, independent of iteration order.
    for (let i = 0; i < moved.length; i++)
      p.set(updates.slice(i * 3, i * 3 + 3), moved[i] * 3);
    if (s.brush === "mask") {
      if (this.stroke.touched.size) this.revision++;
    } else this.changed(moved);
  }
  end(cancel = false) {
    if (!this.stroke) return;
    const { before, beforeMasks, touched, settings } = this.stroke;
    this.stroke = null;
    const maskStroke = settings.brush === "mask";
    const ids = Uint32Array.from(
      [...touched].filter((v) =>
        maskStroke
          ? beforeMasks[v] !== this.mesh.masks[v]
          : [0, 1, 2].some(
              (k) => before[v * 3 + k] !== this.mesh.positions[v * 3 + k],
            ),
      ),
    );
    if (cancel) {
      if (maskStroke) {
        for (const v of ids) {
          this.mesh.masks[v] = beforeMasks[v];
          this.dirty.add(v);
        }
        if (ids.length) this.revision++;
      } else {
        for (const v of ids)
          this.mesh.positions.set(before.subarray(v * 3, v * 3 + 3), v * 3);
        this.changed(Array.from(ids));
      }
      return;
    }
    if (!ids.length) return;
    const width = maskStroke ? 1 : 3;
    const entry = {
      kind: maskStroke ? "mask" : "positions",
      ids,
      before: new Float32Array(ids.length * width),
      after: new Float32Array(ids.length * width),
    };
    ids.forEach((v, i) => {
      if (maskStroke) {
        entry.before[i] = beforeMasks[v];
        entry.after[i] = this.mesh.masks[v];
      } else {
        entry.before.set(before.subarray(v * 3, v * 3 + 3), i * 3);
        entry.after.set(this.mesh.positions.subarray(v * 3, v * 3 + 3), i * 3);
      }
    });
    entry.bytes =
      entry.ids.byteLength + entry.before.byteLength + entry.after.byteLength;
    this.commit(entry);
  }
  restore(redo = false) {
    if (this.stroke) return false;
    const source = redo ? this.redoStack : this.undoStack,
      target = redo ? this.undoStack : this.redoStack,
      entry = source[source.length - 1];
    if (!entry) return false;
    const expectedState = redo ? entry.beforeState : entry.afterState;
    const nextState = redo ? entry.afterState : entry.beforeState;
    if (
      !Number.isInteger(expectedState) ||
      !Number.isInteger(nextState) ||
      this.state !== expectedState
    )
      throw new Error("Invalid history entry.");
    let prepared = null;
    if (entry.kind === "structural") {
      const value = redo ? entry.after : entry.before;
      const mesh = hydrate(value?.mesh);
      if (
        value.activeId !== null &&
        !mesh.pieces.some((piece) => piece.id === value.activeId)
      )
        throw new Error("Invalid history active piece.");
      prepared = {
        mesh,
        activeId: value.activeId,
        bvh: new BVH(mesh),
        dirty: allDirty(mesh),
      };
    } else if (entry.kind === "color") {
      const piece = this.mesh.pieces.find((p) => p.id === entry.id);
      const color = redo ? entry.after : entry.before;
      if (
        !piece ||
        !Array.isArray(color) ||
        color.length !== 3 ||
        color.some((x) => !Number.isFinite(x) || x < 0 || x > 1)
      )
        throw new Error("Invalid history entry.");
      piece.color = [...color];
      for (
        let v = piece.vertexStart;
        v < piece.vertexStart + piece.vertexCount;
        v++
      )
        this.dirty.add(v);
      this.revision++;
    } else {
      const values = redo ? entry.after : entry.before;
      const width =
        entry.kind === "mask" ? 1 : entry.kind === "positions" ? 3 : 0;
      if (
        !(entry.ids instanceof Uint32Array) ||
        !(values instanceof Float32Array) ||
        !width ||
        values.length !== entry.ids.length * width ||
        [...entry.ids].some((v) => v >= this.mesh.positions.length / 3) ||
        [...values].some((x) => !Number.isFinite(x)) ||
        (entry.kind === "mask" && [...values].some((x) => x < 0 || x > 1))
      )
        throw new Error("Invalid history entry.");
      if (entry.kind === "mask") {
        entry.ids.forEach((v, i) => {
          this.mesh.masks[v] = values[i];
          this.dirty.add(v);
        });
        this.revision++;
      } else {
        entry.ids.forEach((v, i) =>
          this.mesh.positions.set(values.subarray(i * 3, i * 3 + 3), v * 3),
        );
        this.changed(Array.from(entry.ids));
      }
    }
    if (prepared) {
      this.mesh = prepared.mesh;
      this.activeId = prepared.activeId;
      this.bvh = prepared.bvh;
      this.dirty = prepared.dirty;
      this.revision++;
    }
    source.pop();
    target.push(entry);
    this.state = nextState;
    return true;
  }
  install(candidateMesh, selectedId = this.activeId) {
    if (this.stroke)
      throw new Error("Cannot install topology during a stroke.");
    const next = snapshot(candidateMesh),
      current = snapshot(this.mesh);
    if (canonicalEqual(current, next)) return false;
    const before = { mesh: current, activeId: this.activeId };
    const activeId = next.pieces.some((p) => p.id === selectedId)
      ? selectedId
      : (next.pieces[0]?.id ?? null);
    const mesh = hydrate(next),
      bvh = new BVH(mesh),
      dirty = allDirty(mesh);
    const entry = {
      kind: "structural",
      before,
      after: { mesh: next, activeId },
      bytes: snapshotBytes(before.mesh) + snapshotBytes(next),
    };
    this.mesh = mesh;
    this.activeId = activeId;
    this.bvh = bvh;
    this.dirty = dirty;
    this.revision++;
    this.commit(entry);
    return true;
  }
  setColor(color) {
    const piece = this.mesh.pieces.find((p) => p.id === this.activeId);
    if (
      !piece ||
      !Array.isArray(color) ||
      color.length !== 3 ||
      color.some((x) => !Number.isFinite(x) || x < 0 || x > 1)
    )
      throw new Error("Invalid piece color.");
    if (color.every((x, i) => x === piece.color[i])) return false;
    const before = [...piece.color];
    piece.color = [...color];
    for (
      let v = piece.vertexStart;
      v < piece.vertexStart + piece.vertexCount;
      v++
    )
      this.dirty.add(v);
    this.revision++;
    this.commit({
      kind: "color",
      id: piece.id,
      before,
      after: [...color],
      bytes: 24,
    });
    return true;
  }
  editMask(invert = false) {
    const piece = this.mesh.pieces.find((p) => p.id === this.activeId);
    if (!piece) return false;
    const ids = [],
      before = [],
      after = [];
    for (
      let v = piece.vertexStart;
      v < piece.vertexStart + piece.vertexCount;
      v++
    ) {
      const n = invert ? 1 - this.mesh.masks[v] : 0;
      if (n !== this.mesh.masks[v]) {
        ids.push(v);
        before.push(this.mesh.masks[v]);
        after.push(n);
        this.mesh.masks[v] = n;
        this.dirty.add(v);
      }
    }
    if (!ids.length) return false;
    const entry = {
      kind: "mask",
      ids: Uint32Array.from(ids),
      before: Float32Array.from(before),
      after: Float32Array.from(after),
    };
    entry.bytes =
      entry.ids.byteLength + entry.before.byteLength + entry.after.byteLength;
    this.revision++;
    this.commit(entry);
    return true;
  }
}

// Screen-space arc-length sampler. Call break() when the ray misses the surface.
export class StrokePath {
  constructor(spacing, emit) {
    this.spacing = spacing;
    this.emit = emit;
    this.break();
  }
  break() {
    this.generation = (this.generation ?? 0) + 1;
    this.last = null;
    this.remaining = 0;
  }
  push(sample, budget = Infinity) {
    const generation = this.generation;
    if (!this.last) {
      this.last = sample;
      this.remaining = this.spacing;
      this.emit(sample);
      return true;
    }
    const a = this.last,
      dx = sample.x - a.x,
      dy = sample.y - a.y,
      length = Math.hypot(dx, dy);
    let distance = this.remaining;
    let emitted = 0;
    while (distance <= length && length > 0) {
      const t = distance / length;
      const current = {
        x: a.x + dx * t,
        y: a.y + dy * t,
        pressure: a.pressure + (sample.pressure - a.pressure) * t,
      };
      this.emit(current);
      if (this.generation !== generation) return true;
      distance += this.spacing;
      if (++emitted >= budget && distance <= length) {
        this.last = current;
        this.remaining = this.spacing;
        return false;
      }
    }
    this.remaining = distance - length;
    this.last = sample;
    return true;
  }
}

export const saveProject = saveDocument;
export const loadProject = loadDocument;
export function exportOBJ(mesh) {
  const lines = ["# Yawn Sculpt — local mesh export"];
  for (let i = 0; i < mesh.positions.length; i += 3)
    lines.push(
      `v ${mesh.positions[i]} ${mesh.positions[i + 1]} ${mesh.positions[i + 2]}`,
    );
  const pieces = ensureDocument(mesh).pieces;
  for (const piece of pieces) {
    lines.push(`g ${piece.name.replace(/\s+/g, "_")}`);
    for (
      let i = piece.indexStart;
      i < piece.indexStart + piece.indexCount;
      i += 3
    )
      lines.push(
        `f ${mesh.indices[i] + 1} ${mesh.indices[i + 1] + 1} ${mesh.indices[i + 2] + 1}`,
      );
  }
  return lines.join("\n");
}