import {
  Document,
  NodeIO,
} from "@gltf-transform/core";
import {
  ALL_EXTENSIONS,
  KHRMaterialsVariants,
  KHRTextureTransform,
} from "@gltf-transform/extensions";

const REPEAT = 10497;
const CLAMP_TO_EDGE = 33071;

function addAccessor(document, buffer, name, type, array) {
  return document.createAccessor(name, buffer).setType(type).setArray(array);
}

function createTerrainMaterial(document, textureTransformExtension, name, texture, tileScale, extras = null) {
  const material = document.createMaterial(name)
    .setBaseColorTexture(texture)
    .setBaseColorFactor([1, 1, 1, 1])
    .setMetallicFactor(0)
    .setRoughnessFactor(1)
    .setAlphaMode("OPAQUE");
  const textureInfo = material.getBaseColorTextureInfo();
  textureInfo.setWrapS(REPEAT).setWrapT(REPEAT);
  textureInfo.setExtension(KHRTextureTransform.EXTENSION_NAME, textureTransformExtension.createTransform()
    .setScale([Number(tileScale), -Number(tileScale)]));
  if (extras) material.setExtras(extras);
  return material;
}

function addVariantMapping(document, variantExtension, mappingList, material, name) {
  const variant = variantExtension.createVariant(name);
  mappingList.addMapping(variantExtension.createMapping().setMaterial(material).addVariant(variant));
  return variant;
}

function setTextureProvenance(texture, sourceDDS, extra = {}) {
  texture.setExtras({
    ...texture.getExtras(),
    runewakerSourceDDS: sourceDDS,
    ...extra,
  });
}

function createVisualDocument(converted, textureSources, provenance) {
  const slug = provenance.dungeonId ?? "sand-vortex";
  const title = provenance.dungeonTitle ?? "Sand Vortex";
  const document = new Document();
  const root = document.getRoot();
  const buffer = document.createBuffer(`${slug}-native-terrain-buffer`);
  const scene = document.createScene(`${title} native terrain`);
  const partitionChunks = provenance.partitionVisualChunks === true;
  const mesh = partitionChunks ? null : document.createMesh(`${title} terrain tiles`);
  const chunkNodes = [];
  const variants = document.createExtension(KHRMaterialsVariants).setRequired(false);
  const textureTransform = document.createExtension(KHRTextureTransform).setRequired(false);
  const textureBySource = new Map();
  for (const [sourceKey, details] of textureSources.textureBySource) {
    const texture = document.createTexture(details.imageName)
      .setImage(details.image)
      .setMimeType("image/png")
      .setExtras({
        runewakerSourceDDS: details.sourceDDS,
        ...(details.sourceCompressedFile ? { runewakerSourceCompressed: details.sourceCompressedFile } : {}),
        ...(details.channelLayerIds ? { runewakerWeightChannels: details.channelLayerIds } : {}),
        runewakerColorPolicy: details.alphaPolicy ?? details.colorSpace,
      });
    textureBySource.set(sourceKey, texture);
  }
  const positions = new Float32Array(converted.vertexCount * 3);
  const normals = new Float32Array(converted.vertexCount * 3);
  const uvs = new Float32Array(converted.vertexCount * 2);

  for (const chunk of converted.chunks) {
    positions.set(chunk.positions, chunk.vertexOffset * 3);
    normals.set(chunk.normals, chunk.vertexOffset * 3);
    uvs.set(chunk.uvs, chunk.vertexOffset * 2);
  }

  const positionAccessor = partitionChunks ? null : addAccessor(document, buffer, `${slug}-native-terrain-positions`, "VEC3", positions);
  const normalAccessor = partitionChunks ? null : addAccessor(document, buffer, `${slug}-native-terrain-normals`, "VEC3", normals);
  const uvAccessor = partitionChunks ? null : addAccessor(document, buffer, `${slug}-native-terrain-normalized-local-xz`, "VEC2", uvs);

  for (const chunk of converted.chunks) {
    const layers = chunk.layers;
    const weights = chunk.splatMaps;
    if (!Array.isArray(layers) || layers.length < 1 || layers.length > 6) {
      throw new Error(`${chunk.id} must contain 1–6 source terrain layers.`);
    }
    if (!Array.isArray(weights) || weights.length !== Math.ceil(layers.length / 3)) {
      throw new Error(`${chunk.id} weight-map count does not match its terrain-layer count.`);
    }
    const metadata = {
      version: 1,
      layerVariants: chunk.material.variantNames,
      weightVariants: chunk.material.weightVariantNames,
      tileScales: layers.map((layer) => Number(layer.tileScale)),
      ...(provenance.cliffProjection ? { cliffProjection: provenance.cliffProjection } : {}),
    };
    if (metadata.layerVariants.length !== layers.length || metadata.weightVariants.length !== weights.length) {
      throw new Error(`${chunk.id} variant name counts do not match its source terrain data.`);
    }

    const firstLayer = layers[0];
    const firstTexture = textureBySource.get(firstLayer.sourceDDS);
    if (!firstTexture) throw new Error(`${chunk.id} default tile image is missing: ${firstLayer.sourceDDS}`);
    const defaultMaterial = createTerrainMaterial(
      document,
      textureTransform,
      `${chunk.id} default layer ${firstLayer.layerId}`,
      firstTexture,
      firstLayer.tileScale,
      { runewakerTerrain: metadata },
    );
    const mappingList = variants.createMappingList();
    addVariantMapping(document, variants, mappingList, defaultMaterial, metadata.layerVariants[0]);

    for (let layerIndex = 1; layerIndex < layers.length; layerIndex += 1) {
      const layer = layers[layerIndex];
      const texture = textureBySource.get(layer.sourceDDS);
      if (!texture) throw new Error(`${chunk.id} layer image is missing: ${layer.sourceDDS}`);
      const material = createTerrainMaterial(
        document,
        textureTransform,
        `${chunk.id} layer ${layer.layerId}`,
        texture,
        layer.tileScale,
      );
      addVariantMapping(document, variants, mappingList, material, metadata.layerVariants[layerIndex]);
    }

    for (let weightIndex = 0; weightIndex < weights.length; weightIndex += 1) {
      const weight = weights[weightIndex];
      const texture = textureBySource.get(weight.linearPngFile);
      if (!texture) throw new Error(`${chunk.id} linear splat image is missing: ${weight.linearPngFile}`);
      const material = document.createMaterial(`${chunk.id} linear weights ${weight.mapIndex}`)
        .setMetallicFactor(0)
        .setRoughnessFactor(1)
        .setAlphaMode("OPAQUE")
        .setMetallicRoughnessTexture(texture)
        .setOcclusionTexture(texture)
        .setExtras({
          runewakerTerrainWeightMap: {
            mapIndex: weight.mapIndex,
            channelLayerIds: weight.channelLayerIds,
            sourceDDS: weight.splatDDS,
            linearPngFile: weight.linearPngFile,
            channels: "linear RGB; source alpha ignored",
          },
        });
      material.getMetallicRoughnessTextureInfo().setWrapS(CLAMP_TO_EDGE).setWrapT(CLAMP_TO_EDGE);
      material.getOcclusionTextureInfo().setWrapS(CLAMP_TO_EDGE).setWrapT(CLAMP_TO_EDGE);
      addVariantMapping(document, variants, mappingList, material, metadata.weightVariants[weightIndex]);
    }

    const indices = addAccessor(
      document,
      buffer,
      `${chunk.id}-source-triangle-indices-with-mask-slots`,
      "SCALAR",
      partitionChunks ? Uint16Array.from(chunk.visualIndices, index => index - chunk.vertexOffset) : chunk.visualIndices,
    );
    const primitive = document.createPrimitive()
      .setAttribute("POSITION", positionAccessor ?? addAccessor(document, buffer, `${chunk.id}-positions`, "VEC3", chunk.positions))
      .setAttribute("NORMAL", normalAccessor ?? addAccessor(document, buffer, `${chunk.id}-normals`, "VEC3", chunk.normals))
      .setAttribute("TEXCOORD_0", uvAccessor ?? addAccessor(document, buffer, `${chunk.id}-uvs`, "VEC2", chunk.uvs))
      .setIndices(indices)
      .setMaterial(defaultMaterial)
      .setExtension(KHRMaterialsVariants.EXTENSION_NAME, mappingList)
      .setExtras({
        runewakerTerrainChunk: {
          id: chunk.id,
          coordinate: chunk.coordinate,
          sourceVertexFlagsSha256: chunk.vertexFlagsSha256 ?? null,
          sourceCellVisibilityFlagsSha256: chunk.cellVisibilityFlagsSha256 ?? null,
          sourceTriangleVisibilityFlags: Buffer.from(chunk.triangleVisibilityFlags).toString("base64"),
          sourceTriangleVisibilityFlagsEncoding: "base64 encoded Uint8Array; one value per source triangle; 1 visible / 0 masked",
          sourceTriangleVisibilitySha256: chunk.triangleVisibilitySha256 ?? null,
          sourceTriangleSlots: chunk.triangleVisibilityFlags.length,
          visibleSourceTriangles: chunk.visibleSourceTriangles,
          maskedSourceTriangles: chunk.maskedSourceTriangles,
        },
      });
    if (mesh) mesh.addPrimitive(primitive);
    else chunkNodes.push(document.createNode(chunk.id)
      .setMesh(document.createMesh(`${chunk.id} native tile`).addPrimitive(primitive)));
  }

  const node = document.createNode(`${title} native terrain`)
    .setExtras({
      runewakerTerrain: {
        version: 1,
        source: "native WDB terrain chunks; no resampling or layer baking",
        chunkCount: converted.chunks.length,
        partition: partitionChunks ? "native-chunk-meshes" : "shared-grid-mesh",
        sourceManifestSha256: provenance.sourceManifestSha256,
        primaryPlacementGuid: provenance.primaryPlacementGuid,
      },
    });
  if (mesh) node.setMesh(mesh);
  for (const chunkNode of chunkNodes) node.addChild(chunkNode);
  scene.addChild(node);
  root.setExtras({
    runewakerTerrain: {
      version: 1,
      source: "Runewaker native WDB heightfield and texture list",
      sourceManifestSha256: provenance.sourceManifestSha256,
      sourceWdbSha256: provenance.sourceWdbSha256,
      primaryPlacementGuid: provenance.primaryPlacementGuid,
      metersPerSourceUnit: provenance.metersPerSourceUnit,
      coordinateBridge: "inverse primary descriptor TRS, then (x,y,z)=(x,y,-z)*metersPerSourceUnit",
      vertexNormalSource: "decoded native vertexData RGB; native terrain shader normal",
      textureSourceDDs: provenance.textureSourceDDS,
    },
  });
  root.setDefaultScene(scene);
  return document;
}

function createCollisionDocument(converted, provenance) {
  const slug = provenance.dungeonId ?? "sand-vortex";
  const title = provenance.dungeonTitle ?? "Sand Vortex";
  const document = new Document();
  const root = document.getRoot();
  const buffer = document.createBuffer(`${slug}-native-terrain-collision-buffer`);
  const scene = document.createScene(`${title} native terrain collision`);
  const mesh = document.createMesh(`${title} native terrain collision surfaces`);
  const partitionChunks = provenance.partitionVisualChunks === true;
  const positions = new Float32Array(partitionChunks ? 0 : converted.vertexCount * 3);
  if (!partitionChunks) for (const chunk of converted.chunks) positions.set(chunk.positions, chunk.vertexOffset * 3);
  const positionAccessor = partitionChunks ? null : addAccessor(document, buffer, `${slug}-native-terrain-collision-positions`, "VEC3", positions);

  for (const chunk of converted.chunks) {
    const indices = addAccessor(
      document,
      buffer,
      `${chunk.id}-visible-collision-triangles`,
      "SCALAR",
      partitionChunks ? Uint16Array.from(chunk.collisionIndices, index => index - chunk.vertexOffset) : chunk.collisionIndices,
    );
    mesh.addPrimitive(document.createPrimitive()
      .setAttribute("POSITION", positionAccessor ?? addAccessor(document, buffer, `${chunk.id}-collision-positions`, "VEC3", chunk.positions))
      .setIndices(indices)
      .setExtras({
        runewakerTerrainCollisionChunk: {
          id: chunk.id,
          coordinate: chunk.coordinate,
          sourceTriangleSlots: chunk.triangleVisibilityFlags.length,
          sourceVisibleTriangles: chunk.visibleSourceTriangles,
          includedCollisionTriangles: chunk.collisionIndices.length / 3,
          maskedTrianglesExcluded: chunk.maskedSourceTriangles,
        },
      }));
  }

  scene.addChild(document.createNode(`${title} native terrain collision`).setMesh(mesh));
  root.setExtras({
    runewakerTerrainCollision: {
      version: 1,
      sourceManifestSha256: provenance.sourceManifestSha256,
      sourceTriangleSlots: converted.report.sourceTriangleSlots,
      includedVisibleTriangles: converted.report.collisionTriangles,
      maskedOrDegenerateTrianglesExcluded: converted.report.excludedCollisionTriangles,
    },
  });
  root.setDefaultScene(scene);
  return document;
}

export async function writeNativeTerrainGlbs({ converted, textureSources, provenance, visualFile, collisionFile }) {
  const io = new NodeIO().registerExtensions(ALL_EXTENSIONS);
  await io.write(visualFile, createVisualDocument(converted, textureSources, provenance));
  await io.write(collisionFile, createCollisionDocument(converted, provenance));
}