import assert from "node:assert/strict";
import { spawnSync } from "node:child_process";
import { access, mkdtemp, readFile, rm, writeFile } from "node:fs/promises";
import os from "node:os";
import path from "node:path";
import { fileURLToPath } from "node:url";
import { deflateSync } from "node:zlib";
import test from "node:test";
import { Document, NodeIO } from "@gltf-transform/core";
import { ALL_EXTENSIONS } from "@gltf-transform/extensions";

const scriptDirectory = path.dirname(fileURLToPath(import.meta.url));
const projectRoot = path.resolve(scriptDirectory, "../..");
const wrapper = path.join(scriptDirectory, "compress-environment-textures.mjs");
const ktxBinDirectory = path.join(projectRoot, ".tools/ktx/bin");
const ktxExecutable = path.join(ktxBinDirectory, process.platform === "win32" ? "toktx.exe" : "toktx");
const hasToktx = await access(ktxExecutable).then(() => true).catch(() => false);

function crc32(bytes) {
  let crc = 0xffffffff;
  for (const byte of bytes) {
    crc ^= byte;
    for (let bit = 0; bit < 8; bit += 1) {
      crc = (crc >>> 1) ^ ((crc & 1) ? 0xedb88320 : 0);
    }
  }
  return (crc ^ 0xffffffff) >>> 0;
}

function pngChunk(type, data) {
  const typeBytes = Buffer.from(type, "ascii");
  const contents = Buffer.concat([typeBytes, data]);
  const header = Buffer.alloc(4);
  header.writeUInt32BE(data.length);
  const checksum = Buffer.alloc(4);
  checksum.writeUInt32BE(crc32(contents));
  return Buffer.concat([header, contents, checksum]);
}

function makePng(width, height, seed) {
  const header = Buffer.alloc(13);
  header.writeUInt32BE(width, 0);
  header.writeUInt32BE(height, 4);
  header[8] = 8;
  header[9] = 6;
  const rows = [];
  for (let y = 0; y < height; y += 1) {
    const row = Buffer.alloc(1 + width * 4);
    for (let x = 0; x < width; x += 1) {
      const offset = 1 + x * 4;
      row[offset] = (x * 17 + y * 3 + seed) % 256;
      row[offset + 1] = (x * 5 + y * 11 + seed * 2) % 256;
      row[offset + 2] = (x * 7 + y * 13 + seed * 3) % 256;
      row[offset + 3] = (x + y) % 5 === 0 ? 127 : 255;
    }
    rows.push(row);
  }
  return Buffer.concat([
    Buffer.from([137, 80, 78, 71, 13, 10, 26, 10]),
    pngChunk("IHDR", header),
    pngChunk("IDAT", deflateSync(Buffer.concat(rows))),
    pngChunk("IEND", Buffer.alloc(0)),
  ]);
}

function patchGlb(buffer, mutateJson) {
  const chunks = [];
  let json;
  let jsonChunkIndex = -1;
  let offset = 12;
  while (offset + 8 <= buffer.length) {
    const length = buffer.readUInt32LE(offset);
    const type = buffer.readUInt32LE(offset + 4);
    const start = offset + 8;
    const end = start + length;
    const data = buffer.subarray(start, end);
    chunks.push({ type, data });
    if (type === 0x4e4f534a) {
      jsonChunkIndex = chunks.length - 1;
      json = JSON.parse(data.toString("utf8").trim());
    }
    offset = end;
  }
  assert.ok(json);
  mutateJson(json);
  const encoded = Buffer.from(JSON.stringify(json));
  const padded = Buffer.alloc(Math.ceil(encoded.length / 4) * 4, 0x20);
  encoded.copy(padded);
  const outputChunks = chunks.map((chunk, index) => (index === jsonChunkIndex ? { ...chunk, data: padded } : chunk));
  const totalLength = 12 + outputChunks.reduce((sum, chunk) => sum + 8 + chunk.data.length, 0);
  const output = Buffer.alloc(totalLength);
  buffer.copy(output, 0, 0, 12);
  output.writeUInt32LE(totalLength, 8);
  offset = 12;
  for (const chunk of outputChunks) {
    output.writeUInt32LE(chunk.data.length, offset);
    output.writeUInt32LE(chunk.type, offset + 4);
    chunk.data.copy(output, offset + 8);
    offset += 8 + chunk.data.length;
  }
  return output;
}

async function createFixture(file) {
  const document = new Document();
  const buffer = document.createBuffer();
  const textureDefinitions = [
    { imageName: "11111111_lava", materialName: "lava-material", alphaMode: "BLEND", seed: 23 },
    { imageName: "22222222_lymor grave_relief_01", materialName: "relief-material", alphaMode: "MASK", seed: 59 },
    { imageName: "33333333_lava_roughness", materialName: "near-match-material", alphaMode: "OPAQUE", seed: 91 },
    { imageName: "44444444_elves_vine", materialName: "reviewed-vine-mask-material", alphaMode: "MASK", seed: 107 },
  ];
  const primitives = [];
  for (const [index, definition] of textureDefinitions.entries()) {
    const texture = document.createTexture(definition.imageName)
      .setMimeType("image/png")
      .setImage(makePng(128, 64, definition.seed));
    const material = document.createMaterial(definition.materialName)
      .setBaseColorTexture(texture)
      .setAlphaMode(definition.alphaMode)
      .setAlphaCutoff(definition.alphaMode === "MASK" ? 0.37 : 0.5);
    const x = index * 2;
    const primitive = document.createPrimitive()
      .setAttribute("POSITION", document.createAccessor(`position-${index}`)
        .setType("VEC3")
        .setArray(new Float32Array([x, 0, 0, x + 1, 0, 0, x, 1, 0]))
        .setBuffer(buffer))
      .setAttribute("TEXCOORD_0", document.createAccessor(`uv-${index}`)
        .setType("VEC2")
        .setArray(new Float32Array([0.03, 0.17, 0.91, 0.24, 0.08, 0.96]))
        .setBuffer(buffer))
      .setIndices(document.createAccessor(`indices-${index}`)
        .setType("SCALAR")
        .setArray(new Uint16Array([0, 1, 2]))
        .setBuffer(buffer))
      .setMaterial(material);
    primitives.push(primitive);
  }
  const mesh = document.createMesh("texture-fixture-mesh");
  for (const primitive of primitives) mesh.addPrimitive(primitive);
  document.createScene("scene").addChild(document.createNode("texture-fixture").setMesh(mesh));
  document.getRoot().setDefaultScene(document.getRoot().listScenes()[0]);

  const io = new NodeIO().registerExtensions(ALL_EXTENSIONS);
  await io.write(file, document);
  const initial = await readFile(file);
  const rawImageNames = textureDefinitions.map(({ imageName }) => imageName);
  await writeFile(file, patchGlb(initial, (json) => {
    for (const [index, texture] of json.textures.entries()) {
      const imageIndex = texture.source;
      json.images[imageIndex].name = rawImageNames[index];
      delete texture.name;
    }
    // Exercise a real GLTF alias: two base-color texture objects share the
    // same image bytes while retaining independent sampler settings.
    json.samplers ??= [];
    const aliasSampler = json.samplers.length;
    json.samplers.push({ wrapS: 33071, wrapT: 33071 });
    json.textures.push({ ...json.textures[0], name: "lava-clamp-sampler-alias", sampler: aliasSampler });
    json.materials.push({
      name: "lava-repeat-alias-material",
      pbrMetallicRoughness: {
        baseColorTexture: { index: json.textures.length - 1, texCoord: 0 },
        metallicFactor: 0,
        roughnessFactor: 1,
      },
    });
    json.meshes[0].primitives.push({ ...json.meshes[0].primitives[0], material: json.materials.length - 1 });
  }));
  return { io, textureDefinitions };
}

function inspectGlbTextures(buffer) {
  const jsonLength = buffer.readUInt32LE(12);
  const json = JSON.parse(buffer.toString("utf8", 20, 20 + jsonLength));
  const binStart = 20 + jsonLength + 8;
  return json.images.map((image, imageIndex) => {
    const view = json.bufferViews[image.bufferView];
    const start = binStart + (view.byteOffset ?? 0);
    const ktx = buffer.subarray(start, start + view.byteLength);
    assert.equal(ktx.subarray(0, 12).toString("hex"), "ab4b5458203230bb0d0a1a0a", `${image.name} must be KTX2`);
    const scheme = ktx.readUInt32LE(44);
    return {
      imageIndex,
      name: image.name,
      width: ktx.readUInt32LE(20),
      height: ktx.readUInt32LE(24),
      mipLevels: Math.max(1, ktx.readUInt32LE(40)),
      supercompressionScheme: scheme,
      codec: scheme === 1 ? "ETC1S" : "UASTC",
    };
  });
}

async function readMeshArrays(io, file) {
  const document = await io.read(file);
  return document.getRoot().listMeshes().map((mesh) => mesh.listPrimitives().map((primitive) => ({
    positions: Array.from(primitive.getAttribute("POSITION")?.getArray() ?? []),
    uvs: Array.from(primitive.getAttribute("TEXCOORD_0")?.getArray() ?? []),
    indices: Array.from(primitive.getIndices()?.getArray() ?? []),
  })));
}

test("selective environment KTX2 policy keeps reviewed base colors UASTC and other base colors ETC1S", {
  skip: hasToktx ? false : "KTX-Software toktx is not available",
}, async (context) => {
  const workingDirectory = await mkdtemp(path.join(os.tmpdir(), "runewaker-mixed-codec-"));
  context.after(async () => rm(workingDirectory, { recursive: true, force: true }));
  const file = path.join(workingDirectory, "environment.glb");
  const report = path.join(workingDirectory, "quality-report.json");
  const { io, textureDefinitions } = await createFixture(file);
  const sourceBuffer = await readFile(file);
  const sourceMeshArrays = await readMeshArrays(io, file);
  const result = spawnSync(process.execPath, [
    wrapper,
    file,
    "--jobs",
    "2",
    "--report",
    report,
  ], {
    cwd: projectRoot,
    encoding: "utf8",
    timeout: 120_000,
    windowsHide: true,
    env: { ...process.env, KTX_BIN_DIR: ktxBinDirectory },
  });
  assert.equal(result.status, 0, `${result.stdout}\n${result.stderr}`);

  const outputBuffer = await readFile(file);
  const textures = inspectGlbTextures(outputBuffer);
  assert.equal(textures.length, textureDefinitions.length);
  const byName = new Map(textures.map((texture) => [texture.name, texture]));
  for (const definition of textureDefinitions) {
    const image = byName.get(definition.imageName);
    assert.ok(image, `missing image ${definition.imageName}`);
    assert.deepEqual([image.width, image.height, image.mipLevels], [128, 64, 8]);
    const expectedCodec = definition.imageName.endsWith("_lava_roughness") ? "ETC1S" : "UASTC";
    assert.equal(image.codec, expectedCodec, `${definition.imageName} codec`);
    assert.equal(image.supercompressionScheme, expectedCodec === "ETC1S" ? 1 : 2);
  }

  const originalJsonLength = sourceBuffer.readUInt32LE(12);
  const originalJson = JSON.parse(sourceBuffer.toString("utf8", 20, 20 + originalJsonLength));
  const outputJsonLength = outputBuffer.readUInt32LE(12);
  const outputJson = JSON.parse(outputBuffer.toString("utf8", 20, 20 + outputJsonLength));
  assert.deepEqual(
    outputJson.materials.map((material) => [material.name, material.alphaMode ?? "OPAQUE", material.alphaCutoff ?? null]),
    originalJson.materials.map((material) => [material.name, material.alphaMode ?? "OPAQUE", material.alphaCutoff ?? null]),
  );
  assert.deepEqual(outputJson.textures.map((texture) => texture.name ?? null), originalJson.textures.map((texture) => texture.name ?? null));
  assert.deepEqual(outputJson.images.map((image) => image.name ?? null), originalJson.images.map((image) => image.name ?? null));
  assert.equal(outputJson.textures.length, originalJson.textures.length);
  assert.deepEqual(outputJson.samplers, originalJson.samplers, "independent alias sampler settings must survive");
  assert.equal(outputJson.textures.at(-1).source, outputJson.textures[0].source, "reviewed aliases keep one shared image");
  assert.equal(outputJson.samplers[outputJson.textures.at(-1).sampler].wrapS, 33071);
  assert.deepEqual(await readMeshArrays(io, file), sourceMeshArrays, "geometry, triangle topology, and UVs must be unchanged");

  const qualityReport = JSON.parse(await readFile(report, "utf8"));
  assert.deepEqual(qualityReport.totals, {
    files: 1,
    textures: 5,
    etc1s: 1,
    uastc: 4,
    baseColorTargets: 4,
    beforeBytes: sourceBuffer.length,
    afterBytes: outputBuffer.length,
  });
  assert.deepEqual(qualityReport.files[0].uastcTargets.map((target) => target.name).sort(), [
    "elves_vine",
    "lava",
    "lava",
    "lymor grave_relief_01",
  ]);
});