import assert from "node:assert/strict";
import fs from "node:fs";
import factory from "../../web/app/src/vendor/blender/web_engine.js";

const root = new URL("../../", import.meta.url);
const wasmBinary = fs.readFileSync(new URL("web/app/src/vendor/blender/web_engine.wasm", root));
const colorFixture = fs.readFileSync(new URL("tests/files/web/attribute_scene.blend", root));
const weightFixture = fs.readFileSync(new URL("tests/files/web/rigged_shape_scene.blend", root));

function open(engine, handle, bytes) {
  const pointer = engine._malloc(bytes.byteLength);
  try {
    engine.HEAPU8.set(bytes, pointer);
    assert.equal(engine._web_engine_open_blend(handle, pointer, bytes.byteLength), 0,
      engine.UTF8ToString(engine._web_engine_last_error_message()));
  }
  finally { engine._free(pointer); }
}

function output(engine, handle, fn, owned = false) {
  const dataOut = engine._malloc(4);
  const lengthOut = engine._malloc(4);
  try {
    assert.equal(fn(handle, dataOut, lengthOut), 0, engine.UTF8ToString(engine._web_engine_last_error_message()));
    const pointer = engine.HEAPU32[dataOut >>> 2];
    const length = engine.HEAPU32[lengthOut >>> 2];
    const bytes = engine.HEAPU8.slice(pointer, pointer + length);
    if (owned) engine._web_engine_free_buffer(pointer);
    return bytes;
  }
  finally {
    engine._free(dataOut);
    engine._free(lengthOut);
  }
}

function snapshot(engine, handle) {
  return JSON.parse(new TextDecoder().decode(output(engine, handle, engine._web_engine_get_scene_snapshot)));
}

function apply(engine, handle, payload) {
  const bytes = new TextEncoder().encode(JSON.stringify(payload));
  const pointer = engine._malloc(bytes.byteLength);
  try {
    engine.HEAPU8.set(bytes, pointer);
    assert.equal(engine._web_engine_apply_command(handle, pointer, bytes.byteLength), 0,
      `${payload.type}: ${engine.UTF8ToString(engine._web_engine_last_error_message())}`);
  }
  finally { engine._free(pointer); }
}

function reject(engine, handle, payload, code) {
  const bytes = new TextEncoder().encode(JSON.stringify(payload));
  const pointer = engine._malloc(bytes.byteLength);
  try {
    engine.HEAPU8.set(bytes, pointer);
    assert.notEqual(engine._web_engine_apply_command(handle, pointer, bytes.byteLength), 0,
      `${payload.type} unexpectedly succeeded`);
    assert.match(engine.UTF8ToString(engine._web_engine_last_error_message()), new RegExp(code));
  }
  finally { engine._free(pointer); }
}

function close(actual, expected, label) {
  assert.ok(Math.abs(actual - expected) <= 1e-6, `${label}: ${actual} != ${expected}`);
}

function assertUniformColor(mesh, expected) {
  assert.equal(mesh.colors.length, mesh.cornerCount * 4);
  for (let offset = 0; offset < mesh.colors.length; offset += 4) {
    for (let channel = 0; channel < 4; channel++) {
      close(mesh.colors[offset + channel], expected[channel], `color[${offset / 4}][${channel}]`);
    }
  }
}

function vertexWeight(mesh, vertex, groupName) {
  const group = mesh.skinWeights?.boneNames.indexOf(groupName) ?? -1;
  if (group < 0) return 0;
  for (let slot = 0; slot < 4; slot++) {
    const offset = vertex * 4 + slot;
    if (mesh.skinWeights.indices[offset] === group) return mesh.skinWeights.weights[offset];
  }
  return 0;
}

function assertVertexGroup(mesh, groupName) {
  assert.ok(mesh.vertexGroups?.some((group) => group.name === groupName), `${groupName} missing from vertexGroups`);
  assert.ok(mesh.skinWeights?.boneNames.includes(groupName), `${groupName} missing from skinWeights.boneNames`);
}

const engine = await factory({ wasmBinary: wasmBinary.slice() });

const colorHandle = engine._web_engine_create();
open(engine, colorHandle, colorFixture);
const originalColorMesh = snapshot(engine, colorHandle).meshes.find((mesh) => mesh.id === "mesh:AttributeMesh");
assert.equal(originalColorMesh.cornerCount, 7);
const paintColor = [0.125, 0.375, 0.625, 0.875];
apply(engine, colorHandle, {
  type: "setVertexColors",
  meshId: originalColorMesh.id,
  attributeName: "WebPaintColor",
  domain: "CORNER",
  indices: Array.from({ length: originalColorMesh.cornerCount }, (_, index) => index),
  colors: Array.from({ length: originalColorMesh.cornerCount }, () => paintColor).flat(),
});
assertUniformColor(snapshot(engine, colorHandle).meshes.find((mesh) => mesh.id === originalColorMesh.id), paintColor);
assert.equal(engine._web_engine_undo(colorHandle), 0, engine.UTF8ToString(engine._web_engine_last_error_message()));
assert.notDeepEqual(snapshot(engine, colorHandle).meshes.find((mesh) => mesh.id === originalColorMesh.id).colors,
  Array.from({ length: originalColorMesh.cornerCount }, () => paintColor).flat());
assert.equal(engine._web_engine_redo(colorHandle), 0, engine.UTF8ToString(engine._web_engine_last_error_message()));
assertUniformColor(snapshot(engine, colorHandle).meshes.find((mesh) => mesh.id === originalColorMesh.id), paintColor);
const savedColors = output(engine, colorHandle, engine._web_engine_save_blend, true);
engine._web_engine_destroy(colorHandle);

const reopenedColors = engine._web_engine_create();
open(engine, reopenedColors, savedColors);
assertUniformColor(snapshot(engine, reopenedColors).meshes.find((mesh) => mesh.id === originalColorMesh.id), paintColor);
engine._web_engine_destroy(reopenedColors);

const weightHandle = engine._web_engine_create();
open(engine, weightHandle, weightFixture);
const meshId = "mesh:RiggedShapeMesh";
const objectId = "object:RiggedShapeObject";
apply(engine, weightHandle, {
  type: "setVertexWeights",
  objectId,
  vertexGroup: "WebPaintGroup",
  indices: [0, 1],
  values: [0.75, 0.25],
  normalize: false,
});
let weightedMesh = snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId);
assertVertexGroup(weightedMesh, "WebPaintGroup");
close(vertexWeight(weightedMesh, 0, "WebPaintGroup"), 0.75, "vertex 0 paint weight");
close(vertexWeight(weightedMesh, 1, "WebPaintGroup"), 0.25, "vertex 1 paint weight");
apply(engine, weightHandle, {
  type: "setVertexWeights",
  objectId,
  vertexGroup: "WebPaintGroup",
  indices: [2],
  values: [0.5],
  normalize: true,
});
weightedMesh = snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId);
const normalizedSum = weightedMesh.skinWeights.weights.slice(8, 12).reduce((sum, weight) => sum + weight, 0);
close(normalizedSum, 1, "normalized vertex 2 weight sum");
assert.equal(engine._web_engine_undo(weightHandle), 0, engine.UTF8ToString(engine._web_engine_last_error_message()));
close(vertexWeight(snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId), 2, "WebPaintGroup"), 0,
  "undone vertex 2 paint weight");
assert.equal(engine._web_engine_redo(weightHandle), 0, engine.UTF8ToString(engine._web_engine_last_error_message()));
close(vertexWeight(snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId), 2, "WebPaintGroup"), 0.5 / 1.75,
  "redone normalized vertex 2 paint weight");
apply(engine, weightHandle, {
  type: "setVertexWeights",
  objectId,
  vertexGroup: "WebPaintGroup",
  indices: [3],
  values: [0.4],
  limit: 2,
  normalize: true,
});
weightedMesh = snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId);
close(vertexWeight(weightedMesh, 3, "WebPaintGroup"), 0.4 / 1.4, "limited normalized vertex 3 weight");
apply(engine, weightHandle, {
  type: "setVertexWeights",
  objectId,
  vertexGroup: "WebPaintGroup",
  indices: [0],
  values: [0.9],
  mirror: true,
  mirrorAxis: 0,
  mirrorTolerance: 1e-4,
});
weightedMesh = snapshot(engine, weightHandle).meshes.find((mesh) => mesh.id === meshId);
close(vertexWeight(weightedMesh, 0, "WebPaintGroup"), 0.9, "mirrored vertex 0 weight");
close(vertexWeight(weightedMesh, 1, "WebPaintGroup"), 0.9, "mirrored vertex 1 weight");
const savedWeights = output(engine, weightHandle, engine._web_engine_save_blend, true);
engine._web_engine_destroy(weightHandle);

const reopenedWeights = engine._web_engine_create();
open(engine, reopenedWeights, savedWeights);
weightedMesh = snapshot(engine, reopenedWeights).meshes.find((mesh) => mesh.id === meshId);
assertVertexGroup(weightedMesh, "WebPaintGroup");
close(vertexWeight(weightedMesh, 0, "WebPaintGroup"), 0.9, "reopened mirrored vertex 0 paint weight");
close(vertexWeight(weightedMesh, 1, "WebPaintGroup"), 0.9, "reopened mirrored vertex 1 paint weight");
close(vertexWeight(weightedMesh, 2, "WebPaintGroup"), 0.5 / 1.75, "reopened normalized vertex 2 paint weight");
engine._web_engine_destroy(reopenedWeights);

process.stdout.write("paint-roundtrip-ok vertex-color=passed vertex-weight-normalize=passed vertex-weight-limit=passed vertex-weight-mirror=passed mirror-gate=passed undo-redo=passed save-reopen=passed\n");