import { getPath } from "./common.js";
import { parseGltfChunks, vertexBufferLayouts, primitiveAccessors, splitGltf } from "./gltf.js";


function loadGltfNode(module, handle, gltf, nodeIndex)
{
  const node = gltf.json.nodes[nodeIndex];
  var translation = [0, 0, 0];
  if ("translation" in node) {
    translation = node.translation;
  }
  console.log("init", nodeIndex, translation);
  module.instance.exports.node_init_translation(handle, nodeIndex, translation[0], translation[1], translation[2]);

  if ("children" in node) {
    for (let childIndex in node.children) {
      loadGltfNode(module, handle, gltf, childIndex);
    }
  }
}

function loadGltfNodes(module, gltf)
{
  const nodesHandle = module.instance.exports.node_alloc(gltf.json.nodes.length);
  for (let nodeIndex of gltf.json.scenes[0].nodes) {
    loadGltfNode(module, nodesHandle, gltf, nodeIndex);
  }
  return nodesHandle;
}

function printMatrix(m, o)
{
  console.log(m[o+0], m[o+1], m[o+2], m[o+3]);
  console.log(m[o+4], m[o+5], m[o+6], m[o+7]);
  console.log(m[o+8], m[o+9], m[o+10], m[o+11]);
  console.log(m[o+12], m[o+13], m[o+14], m[o+15]);
}

class GltfRenderer {
  constructor(device, canvasFormat, shaderModule, gltf, texture)
  {
    this.gltf = gltf;

    //////////////////////////////////////////////////////////////////////
    // textures
    //////////////////////////////////////////////////////////////////////

    this.sampler = device.createSampler();
    this.texture = texture;

    //////////////////////////////////////////////////////////////////////
    // buffer
    //////////////////////////////////////////////////////////////////////

    this.buffers = [];

    //console.assert(gltf.json.buffers.length === gltf.bin.length);
    for (let i = 0; i < gltf.bin.length; i++) {
      //console.assert(gltf.json.buffers[i].byteLength == gltf.bin[i].length);
      const buffer = device.createBuffer({
        label: `gltf vertex buffer {i}`,
        size: gltf.bin[i].length,
        usage: GPUBufferUsage.INDEX | GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
      });
      device.queue.writeBuffer(buffer, 0,
                               gltf.bin[i].buffer, gltf.bin[i].offset, gltf.bin[i].length);
      this.buffers.push(buffer);
    }

    //////////////////////////////////////////////////////////////////////
    // pipeline
    //////////////////////////////////////////////////////////////////////

    const bindGroupLayouts = [
      device.createBindGroupLayout({
        label: "gltf bind group layout 0",
        entries: [{
          binding: 0,
          visibility: GPUShaderStage.VERTEX,
          buffer: { type: "uniform" }
        }, {
          binding: 1,
          visibility: GPUShaderStage.VERTEX,
          buffer: { type: "read-only-storage" }
        }, {
          binding: 2,
          visibility: GPUShaderStage.FRAGMENT,
          sampler: {}
        }, {
          binding: 3,
          visibility: GPUShaderStage.FRAGMENT,
          texture: {},
        }]
      }),
    ];

    const pipelineLayout = device.createPipelineLayout({
      label: "gltf pipeline layout",
      bindGroupLayouts: bindGroupLayouts,
    });

    // by vertex buffer layout
    const renderPipelineKeys = {};
    this.renderPipelines = {};

    this.meshes = [];
    for (let i = 0; i < gltf.json.meshes.length; i++) {
      const mesh = gltf.json.meshes[i];
      const meshPrimitives = [];

      for (let j = 0; j < mesh.primitives.length; j++) {
        const accessors = primitiveAccessors(mesh.primitives[j]);
        const key = accessors.join(";");
        meshPrimitives.push([key, accessors]);
        renderPipelineKeys[key] = accessors;
      }
      this.meshes.push(meshPrimitives);
    }

    for (let [key, accessors] of Object.entries(renderPipelineKeys)) {
      const bufferLayouts = vertexBufferLayouts(gltf, accessors);

      const renderPipeline = device.createRenderPipeline({
        label: "gltf pipeline",
        layout: pipelineLayout,
        vertex: {
          module: shaderModule,
          entryPoint: "vertexMain",
          buffers: bufferLayouts,
        },
        fragment: {
          module: shaderModule,
          entryPoint: "fragmentMain",
          targets: [{
            format: canvasFormat,
          }]
        },
        primitive: {
          topology: 'triangle-list',
        },
        depthStencil: {
          depthWriteEnabled: true,
          depthCompare: 'less',
          format: 'depth24plus',
        },
      });

      this.renderPipelines[key] = renderPipeline;
    }

    //////////////////////////////////////////////////////////////////////
    // bind group
    //////////////////////////////////////////////////////////////////////

    const matrixSize = 4 * 4 * 4;
    this.uniformBufferSize = matrixSize * 2;
    this.uniformBuffer = device.createBuffer({
      label: "gltf uniform 0",
      size: this.uniformBufferSize,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });

    this.nodeBufferSize = matrixSize * this.gltf.json.nodes.length;
    this.nodeBuffer = device.createBuffer({
      label: "gltf node buffer",
      size: this.nodeBufferSize,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
    });

    this.bindGroups = [
      device.createBindGroup({
        label: "cell bind group 0",
        layout: bindGroupLayouts[0],
        entries: [{
          binding: 0,
          resource: { buffer: this.uniformBuffer },
        }, {
          binding: 1,
          resource: { buffer: this.nodeBuffer },
        }, {
          binding: 2,
          resource: this.sampler,
        }, {
          binding: 3,
          resource: this.texture,
        }]
      }),
    ];
  }

  renderPrimitive(renderPass, accessors, indices, nodeIndex)
  {
    for (let i = 0; i < accessors.length; i++) {
      const accessor = accessors[i];
      const bufferViewIndex = this.gltf.json.accessors[accessor].bufferView;
      const bufferView = this.gltf.json.bufferViews[bufferViewIndex];

      renderPass.setVertexBuffer(i, this.buffers[bufferView.buffer], bufferView.byteOffset, bufferView.byteLength);
    }

    const indicesAccessor = this.gltf.json.accessors[indices];
    const indicesBufferView = this.gltf.json.bufferViews[indicesAccessor.bufferView];
    renderPass.setIndexBuffer(this.buffers[indicesBufferView.buffer],
                              "uint16",
                              indicesBufferView.byteOffset,
                              indicesBufferView.byteLength);

    renderPass.setBindGroup(0, this.bindGroups[0]);
    //console.log("drawIndexed", nodeIndex);
    renderPass.drawIndexed(indicesAccessor.count, 1, 0, 0, nodeIndex);
    //renderPass.drawIndexed(indicesAccessor.count);
  }

  renderNode(renderPass, lastKey, nodeIndex)
  {
    const node = this.gltf.json.nodes[nodeIndex];

    if ("mesh" in node) {
      const mesh = this.gltf.json.meshes[node.mesh];
      for (let j = 0; j < mesh.primitives.length; j++) {
        const [key, accessors] = this.meshes[node.mesh][j];
        if (key !== lastKey)
          renderPass.setPipeline(this.renderPipelines[key]);
        lastKey = key;
        const primitive = mesh.primitives[j];
        this.renderPrimitive(renderPass, accessors, primitive.indices, nodeIndex);
      }
    }

    if ("children" in node) {
      for (let childIndex of node.children) {
        lastKey = this.renderNode(renderPass, lastKey, childIndex);
      }
    }
    return lastKey;
  }

  render(device, colorView, depthTexture)
  {
    const encoder = device.createCommandEncoder();
    const renderPass = encoder.beginRenderPass({
      colorAttachments: [{
        view: colorView,
        loadOp: "clear",
        clearValue: { r: 0.2, g: 0.2, b: 0.4, a: 1 },
        storeOp: "store",
      }],
      depthStencilAttachment: {
        view: depthTexture.createView(),
        depthClearValue: 1.0,
        depthLoadOp: 'clear',
        depthStoreOp: 'store',
      },
    });

    var lastKey = undefined;
    for (let nodeIndex of this.gltf.json.scenes[0].nodes) {
      lastKey = this.renderNode(renderPass, lastKey, nodeIndex);
    }

    renderPass.end();

    const commandBuffer = encoder.finish();
    device.queue.submit([commandBuffer]);
  }
}

function memoryString(memory, offset)
{
  const u8 = new Uint8Array(memory.buffer, offset);

  const s = [];
  var i = 0;
  while (u8[i] != 0 && (i - offset) < 256) {
    s.push(u8[i]);
    i++;
  }
  return new TextDecoder().decode(new Uint8Array(s));
}

async function loadTextures(device, memory, module, gltf)
{
  const tempTable = module.instance.exports.mem_alloc(256 * 4);
  const tempHeader = module.instance.exports.mem_alloc(7 * 4);

  const tempInputSize = 8 * 1024 * 1024;
  const tempInput = module.instance.exports.mem_alloc(tempInputSize);

  const tempOutputSize = 2048 * 2048 * 4;
  const tempOutput = module.instance.exports.mem_alloc(tempOutputSize);

  const tempError = module.instance.exports.mem_alloc(4);

  module.instance.exports.png_build_crc_table(tempTable);

  const tempInputU8 = new Uint8Array(memory.buffer, tempInput);
  const tempOutputU8 = new Uint8Array(memory.buffer, tempOutput);
  const tempErrorU32 = new Uint32Array(memory.buffer, tempError);
  const tempHeaderU32 = new Uint32Array(memory.buffer, tempHeader);

  const textures = [];
  for (let image of gltf.json.images) {
    console.log(image.uri);
    const imageResponse = await fetch(`gltf/${image.uri}`);
    const imageBuffer = await imageResponse.arrayBuffer();
    const imageBufferU8 = new Uint8Array(imageBuffer)

    tempInputU8.set(imageBufferU8);
    const ret = module.instance.exports.png_decode(tempTable, tempInput, imageBuffer.byteLength,
                                                   tempHeader,
                                                   tempOutput, tempOutputSize,
                                                   tempError);
    if (ret < 0) {
      console.log("png ret", ret, memoryString(memory, tempErrorU32[0]));
      throw new Error("png error");
    }

    const width = tempHeaderU32[0];
    const height = tempHeaderU32[1];
    const texture = device.createTexture({
      size: [width, height],
      format: 'rgba8unorm',
      usage: GPUTextureUsage.TEXTURE_BINDING | GPUTextureUsage.COPY_DST,
    });

    device.queue.writeTexture({ texture: texture },
                              tempOutputU8,
                              { bytesPerRow: width * 4 },
                              { width: width, height: height });

    textures.push(texture);
  }
  return textures;
}

async function loadGltf(device, canvasFormat, memory, module)
{
  const gltf = await splitGltf("gltf/lamp.gltf", "gltf/lamp.bin");
  const nodesHandle = loadGltfNodes(module, gltf);

  /*
  const nodesM = new Float32Array(memory.buffer, nodesHandle);
  for (let i = 0; i < gltf.json.nodes.length; i++) {
    console.log("[", i, "]");
    printMatrix(nodesM, i * 16);
  }
  */

  const textures = await loadTextures(device, memory, module, gltf);

  const code = await getPath("gltf.wgsl");
  const shaderModule = device.createShaderModule({
    label: "gltf module",
    code: await code,
  });

  const renderer = new GltfRenderer(device, canvasFormat, shaderModule, gltf, textures[0]);
  console.log(renderer.nodeBuffer);
  device.queue.writeBuffer(renderer.nodeBuffer, 0,
                           memory.buffer, nodesHandle, (4 * 4 * 4) * gltf.json.nodes.length);

  return renderer;
}

export { loadGltf };