import { PropertyType } from "@gltf-transform/core";
import { prune } from "@gltf-transform/functions";

export const CREATURE_MESH_CONSOLIDATION_POLICY = "wow-creature-material-join-v1";

function stableJson(value) {
  return JSON.stringify(value, Object.keys(value ?? {}).sort());
}

function isEmptyObject(value) {
  return !value || Object.keys(value).length === 0;
}

function propertyIndexMap(properties) {
  return new Map(properties.map((property, index) => [property, index]));
}

function parentOf(root, node) {
  const parentNode = node.getParentNode();
  if (parentNode) return parentNode;
  return root.listScenes().find((scene) => scene.listChildren().includes(node)) ?? null;
}

function vectorKey(values) {
  return values.map((value) => Object.is(value, -0) ? 0 : value).join(",");
}

function hasCustomExtensions(property) {
  return typeof property.listExtensions === "function" && property.listExtensions().length > 0;
}

function indexArrayType(maxIndex) {
  if (maxIndex <= 0xff) return Uint8Array;
  if (maxIndex <= 0xffff) return Uint16Array;
  return Uint32Array;
}

function primitiveCount(root) {
  return root.listMeshes().reduce((total, mesh) => total + mesh.listPrimitives().length, 0);
}

function triangleCount(root) {
  return root.listMeshes().reduce((meshTotal, mesh) => (
    meshTotal + mesh.listPrimitives().reduce((primitiveTotal, primitive) => {
      const indices = primitive.getIndices();
      const vertices = primitive.getAttribute("POSITION");
      const count = indices?.getCount() ?? vertices?.getCount() ?? 0;
      return primitiveTotal + (primitive.getMode() === 4 ? Math.floor(count / 3) : 0);
    }, 0)
  ), 0);
}

/**
 * Consolidates sibling skinned render nodes that differ only by material/index
 * range. WoW character exports commonly emit one node per geoset while all
 * geosets share the same vertex attributes and skin. Concatenating those index
 * ranges is lossless and changes N draw calls into one draw call per material.
 *
 * The transform intentionally skips animated nodes, joints, morph targets,
 * nodes with children/custom metadata, and primitives with extensions/custom
 * metadata. Skeleton and animation node names are therefore never rewritten.
 */
export async function consolidateCreatureSkinnedMeshes(document) {
  const root = document.getRoot();
  const nodes = root.listNodes();
  const accessors = root.listAccessors();
  const materials = root.listMaterials();
  const skins = root.listSkins();
  const scenes = root.listScenes();
  const nodeIndices = propertyIndexMap(nodes);
  const accessorIndices = propertyIndexMap(accessors);
  const materialIndices = propertyIndexMap(materials);
  const skinIndices = propertyIndexMap(skins);
  const parentIndices = propertyIndexMap([...scenes, ...nodes]);
  const animatedNodes = new Set(root.listAnimations().flatMap((animation) => (
    animation.listChannels().map((channel) => channel.getTargetNode()).filter(Boolean)
  )));
  const jointNodes = new Set(skins.flatMap((skin) => [skin.getSkeleton(), ...skin.listJoints()]).filter(Boolean));

  const before = {
    drawCalls: primitiveCount(root),
    triangles: triangleCount(root),
    meshNodes: nodes.filter((node) => node.getMesh()).length,
  };
  const candidateGroups = new Map();
  let skippedNodes = 0;

  for (const node of nodes) {
    const mesh = node.getMesh();
    const skin = node.getSkin();
    const parent = parentOf(root, node);
    const primitives = mesh?.listPrimitives() ?? [];
    if (
      !mesh
      || !skin
      || !parent
      || primitives.length !== 1
      || node.listChildren().length > 0
      || animatedNodes.has(node)
      || jointNodes.has(node)
      || !isEmptyObject(node.getExtras())
      || !isEmptyObject(mesh.getExtras())
      || hasCustomExtensions(node)
      || hasCustomExtensions(mesh)
    ) {
      if (mesh) skippedNodes += 1;
      continue;
    }

    const primitive = primitives[0];
    const indices = primitive.getIndices();
    const position = primitive.getAttribute("POSITION");
    if (
      !indices
      || !indices.getArray()
      || !position
      || primitive.listTargets().length > 0
      || !isEmptyObject(primitive.getExtras())
      || hasCustomExtensions(primitive)
      || hasCustomExtensions(indices)
    ) {
      skippedNodes += 1;
      continue;
    }

    const semantics = primitive.listSemantics().slice().sort();
    const compatibilityKey = [
      parentIndices.get(parent),
      skinIndices.get(skin),
      vectorKey(node.getTranslation()),
      vectorKey(node.getRotation()),
      vectorKey(node.getScale()),
      vectorKey(node.getWeights()),
      vectorKey(mesh.getWeights()),
      primitive.getMode(),
      materialIndices.get(primitive.getMaterial()) ?? -1,
      ...semantics.map((semantic) => `${semantic}:${accessorIndices.get(primitive.getAttribute(semantic))}`),
      stableJson(primitive.getExtras()),
    ].join("|");
    const group = candidateGroups.get(compatibilityKey) ?? [];
    group.push({ node, mesh, skin, parent, primitive, indices, semantics });
    candidateGroups.set(compatibilityKey, group);
  }

  let groups = 0;
  let sourceDrawCalls = 0;
  let runtimeDrawCalls = 0;
  let mergedNodes = 0;
  let mergedIndexCount = 0;
  for (const group of candidateGroups.values()) {
    if (group.length < 2) continue;
    const representative = group[0];
    const arrays = group.map(({ indices }) => indices.getArray());
    const total = arrays.reduce((sum, array) => sum + array.length, 0);
    let maximum = 0;
    for (const array of arrays) {
      for (let index = 0; index < array.length; index += 1) maximum = Math.max(maximum, array[index]);
    }
    const IndexArray = indexArrayType(maximum);
    const combined = new IndexArray(total);
    let offset = 0;
    for (const array of arrays) {
      combined.set(array, offset);
      offset += array.length;
    }

    const suffix = String(groups + 1).padStart(2, "0");
    const parentName = representative.parent.getName() || "creature";
    const mergedPrimitive = document.createPrimitive()
      .setMode(representative.primitive.getMode())
      .setMaterial(representative.primitive.getMaterial())
      .setIndices(document.createAccessor(
        `${parentName}_merged_indices_${suffix}`,
        representative.indices.getBuffer(),
      ).setType("SCALAR").setArray(combined));
    for (const semantic of representative.semantics) {
      mergedPrimitive.setAttribute(semantic, representative.primitive.getAttribute(semantic));
    }
    const mergedMesh = document.createMesh(`${parentName}_MergedMaterial_${suffix}`)
      .setWeights(representative.mesh.getWeights())
      .addPrimitive(mergedPrimitive);
    const sourceNames = group.map(({ node }) => node.getName()).filter(Boolean);
    const mergedNode = document.createNode(`${parentName}_MergedSkinned_${suffix}`)
      .setTranslation(representative.node.getTranslation())
      .setRotation(representative.node.getRotation())
      .setScale(representative.node.getScale())
      .setWeights(representative.node.getWeights())
      .setMesh(mergedMesh)
      .setSkin(representative.skin)
      .setExtras({
        healerMan: {
          meshConsolidationPolicy: CREATURE_MESH_CONSOLIDATION_POLICY,
          sourceNodes: sourceNames,
        },
      });
    representative.parent.addChild(mergedNode);

    for (const { node } of group) node.dispose();
    groups += 1;
    sourceDrawCalls += group.length;
    runtimeDrawCalls += 1;
    mergedNodes += group.length;
    mergedIndexCount += total;
  }

  if (groups) {
    await document.transform(prune({
      propertyTypes: [
        PropertyType.ANIMATION_CHANNEL,
        PropertyType.ANIMATION_SAMPLER,
        PropertyType.ACCESSOR,
        PropertyType.MESH,
        PropertyType.PRIMITIVE,
      ],
      keepAttributes: true,
      keepExtras: true,
      keepLeaves: true,
      keepSolidTextures: true,
    }));
  }

  const after = {
    drawCalls: primitiveCount(root),
    triangles: triangleCount(root),
    meshNodes: root.listNodes().filter((node) => node.getMesh()).length,
  };
  if (before.triangles !== after.triangles) {
    throw new Error(
      `Creature mesh consolidation changed triangle count (${before.triangles} -> ${after.triangles}).`,
    );
  }

  return {
    policy: CREATURE_MESH_CONSOLIDATION_POLICY,
    groups,
    mergedNodes,
    mergedIndexCount,
    skippedNodes,
    sourceDrawCalls,
    runtimeDrawCalls,
    savedDrawCalls: sourceDrawCalls - runtimeDrawCalls,
    before,
    after,
  };
}