import { getPath } from "./common.js";
import { loadGltf } from "./index2.js";
import { loadLight } from "./light.js";

if (!navigator.gpu) {
  throw new Error("WebGPU not supported on this browser.");
}

const adapter = await navigator.gpu.requestAdapter();
if (!adapter) {
  throw new Error("No WebGPU adapter");
}

const device = await adapter.requestDevice();

const canvas = document.querySelector("canvas");
canvas.width = canvas.clientWidth
canvas.height = canvas.clientHeight

const context = canvas.getContext("webgpu");
const canvasFormat = navigator.gpu.getPreferredCanvasFormat();
context.configure({
  device: device,
  format: canvasFormat,
});

const encoder = device.createCommandEncoder();

const renderPass = encoder.beginRenderPass({
  colorAttachments: [{
    view: context.getCurrentTexture().createView(),
    loadOp: "clear",
    clearValue: { r: 0.9, g: 0.5, b: 0.2, a: 1 },
    storeOp: "store",
  }]
});
renderPass.end();

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

const cubeResponse = await fetch("obj/cube.bin");
const cubeObj = await cubeResponse.arrayBuffer();

const cubeView = new DataView(cubeObj);
const indexCount = cubeView.getUint32(0, true);
const indexBufferOffset = cubeView.getUint32(4, true);
const indexBufferSize = cubeView.getUint32(8, true);
const vertexBufferOffset = cubeView.getUint32(12, true);
const vertexBufferSize = cubeView.getUint32(16, true);
//console.log(`indexBufferOffset ${indexBufferOffset} indexBufferSize ${indexBufferSize}`);
//console.log(`vertexBufferOffset ${vertexBufferOffset} vertexBufferSize ${vertexBufferSize}`);
const cubeIndexCount = indexBufferSize / 2;

const cubeVerticesStride = 36;

const indexVertexBuffer = device.createBuffer({
  label: "cell vertices",
  size: indexBufferSize + vertexBufferSize,
  usage: GPUBufferUsage.INDEX | GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(indexVertexBuffer, 0, cubeObj, indexBufferOffset, indexBufferSize);
device.queue.writeBuffer(indexVertexBuffer, indexBufferSize, cubeObj, vertexBufferOffset, vertexBufferSize);

//////////////////////////////////////////////////////////////////////
// instance buffer
//////////////////////////////////////////////////////////////////////

async function loadVoxModel(path)
{
  const voxResponse = await fetch(path);
  const vox = await voxResponse.arrayBuffer();
  const voxView = new DataView(vox);
  const sizeX = voxView.getUint32(0, true);
  const sizeY = voxView.getUint32(4, true);
  const sizeZ = voxView.getUint32(8, true);
  const voxelCount = voxView.getUint32(12, true);

  const instanceBuffer = device.createBuffer({
    label: `instanceBuffer %{}`,
    size: voxelCount * 4,
    usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
  });
  device.queue.writeBuffer(instanceBuffer, 0, vox, 16, voxelCount * 4);

  const modelConfigurationBuffer = device.createBuffer({
    label: "modelConfigurationBuffer",
    size: 4 * 4 + 256 * 4 * 4,
    usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
  });
  const paletteOffset = 16 + voxelCount * 4;

  const sizeArray = new Float32Array([sizeX, sizeY, sizeZ, 0]);
  device.queue.writeBuffer(modelConfigurationBuffer, 0, sizeArray);
  device.queue.writeBuffer(modelConfigurationBuffer, 4 * 4, vox, paletteOffset, 256 * 4 * 4);

  return {
    instanceBuffer: instanceBuffer,
    modelConfigurationBuffer: modelConfigurationBuffer,
    voxelCount: voxelCount,
    path: path,
  }
}

const modelPaths = [
  "vox/monu1.bin",
  "vox/monu7.bin",
  "vox/monu9.bin",
  "vox/monu16.bin",
  "vox/chr_fox.bin",
  "vox/ff3.bin",
  "vox/dragon.bin",
];

const models = {}
for (const path of modelPaths) {
  models[path] = loadVoxModel(path);
}
for (const path of modelPaths) {
  models[path] = await models[path];
}
var currentModel = modelPaths[0];

//////////////////////////////////////////////////////////////////////
// vertex layout
//////////////////////////////////////////////////////////////////////

const vertexBufferLayouts = [
  {
    arrayStride: cubeVerticesStride,
    attributes: [{
      format: "float32x3",
      offset: 0,
      shaderLocation: 0,
    }, {
      format: "float32x2",
      offset: 12,
      shaderLocation: 1,
    }, {
      format: "float16x4",
      offset: 20,
      shaderLocation: 2,
    }, {
      format: "float16x4",
      offset: 28,
      shaderLocation: 3,
    }],
    stepMode: "vertex",
  },
  {
    arrayStride: 4,
    attributes: [{
      format: "uint8x4",
      offset: 0,
      shaderLocation: 4,
    }],
    stepMode: "instance",
  },
];

const indexWgsl = getPath("index.wgsl");
const computeWgsl = getPath("compute.wgsl");

const cellShaderModule = device.createShaderModule({
  label: "cell shader",
  code: await indexWgsl,
});

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

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

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

const computeShaderModule = device.createShaderModule({
  label: "compute shader",
  code: await computeWgsl,
});

const computePipelineLayout = device.createPipelineLayout({
  label: "compute pipeline layout",
  bindGroupLayouts: [ bindGroupLayouts[0] ],
});

const computePipeline = device.createComputePipeline({
  label: "compute pipeline",
  layout: computePipelineLayout,
  compute: {
    module: computeShaderModule,
    entryPoint: "computeMain",
  },
});

const gridSize = 32;
const workGroupSize = 16;

//const uniformArray2 = new Float32Array([gridSize, gridSize, 1, 1]);
const memory = new WebAssembly.Memory({ initial: 1 });
const uniformArray = new Float32Array(memory.buffer);
uniformArray[0] = 1;
uniformArray[1] = 0;
uniformArray[2] = 0;
uniformArray[3] = 0;
const uniformArraySize = 4 * 4 + (4 * 4 * 4) * 2;
const rotate = await WebAssembly.instantiateStreaming(fetch("rotate.wasm"), {
  env: { memory: memory }
});

const memory2 = new WebAssembly.Memory({ initial: 258 * 2 });
const module = await WebAssembly.instantiateStreaming(fetch("src/module.wasm"), {
  env: {
    memory: memory2,
    log: console.log,
  }
});
//document.module = module;

const uniformBuffer = device.createBuffer({
  label: "grid uniform",
  size: uniformArraySize,
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});
device.queue.writeBuffer(uniformBuffer, 0, uniformArray, 0, uniformArraySize / 4);

const cellStateArray = new Uint32Array(gridSize * gridSize);

const cellStateStorage = [
  device.createBuffer({
    label: "cell state 0",
    size: cellStateArray.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  }),
  device.createBuffer({
    label: "cell state 1",
    size: cellStateArray.byteLength,
    usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
  }),
];

for (let i = 0; i < cellStateArray.length; i += 1) {
  cellStateArray[i] = Math.random() > 0.6 ? 1 : 0;
}
device.queue.writeBuffer(cellStateStorage[0], 0, cellStateArray);
device.queue.writeBuffer(cellStateStorage[1], 0, cellStateArray);

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

const modelBindGroups = {};
for (const path of modelPaths) {
  const bindGroup = device.createBindGroup({
    label: "palette bind group",
    layout: bindGroupLayouts[1],
    entries: [{
      binding: 0,
      resource: { buffer: models[path].modelConfigurationBuffer },
    }]
  });
  modelBindGroups[path] = bindGroup;
}

let tick = 0;
let which = 0;

function createDepthTexture() {
  return device.createTexture({
    size: [canvas.width, canvas.height],
    format: 'depth24plus',
    usage: GPUTextureUsage.RENDER_ATTACHMENT,
  });
}

var depthTexture = undefined;

const modelSelect = document.getElementById("model-select");
for (const path of modelPaths) {
  const option = document.createElement("option");
  option.value = path;
  option.innerHTML = path;
  modelSelect.appendChild(option);
}
modelSelect.value = currentModel;

const lightingSelect = document.getElementById("lighting-select");
lightingSelect.value = "unlit";

function recreateDepth()
{
  if (canvas.clientWidth !== canvas.width || canvas.clientWidth !== canvas.height) {
    canvas.width = canvas.clientWidth;
    canvas.height = canvas.clientHeight;

    if (depthTexture !== undefined) {
      depthTexture.destroy();
    }
    depthTexture = createDepthTexture();
  }
}

function render()
{
  recreateDepth();

  const encoder = device.createCommandEncoder();

  const computePass = encoder.beginComputePass();
  computePass.setPipeline(computePipeline);
  computePass.setBindGroup(0, bindGroups[which]);
  const workgroupCount = Math.ceil((gridSize * gridSize) / workGroupSize);
  computePass.dispatchWorkgroups(workgroupCount);
  computePass.end();

  const aspect = canvas.width / canvas.height;
  rotate.instance.exports.rotate(tick * 0.01, aspect, 4 * 4);

  tick += 1;
  if ((tick % 5) == 0) {
    which ^= 1;
  }

  if (lightingSelect.value === "diffuse") {
    uniformArray[0] = 1;
  } else {
    uniformArray[0] = 0;
  }
  device.queue.writeBuffer(uniformBuffer, 0, uniformArray, 0, uniformArraySize / 4);

  const renderPass = encoder.beginRenderPass({
    colorAttachments: [{
      view: context.getCurrentTexture().createView(),
      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',
    },
  });

  currentModel = modelSelect.value;

  renderPass.setPipeline(cellPipeline);
  const vertexOffset = indexBufferSize;
  renderPass.setVertexBuffer(0, indexVertexBuffer, vertexOffset, vertexBufferSize);
  renderPass.setVertexBuffer(1, models[currentModel].instanceBuffer);
  renderPass.setIndexBuffer(indexVertexBuffer, "uint16", 0, indexBufferSize);
  renderPass.setBindGroup(0, bindGroups[which]);
  renderPass.setBindGroup(1, modelBindGroups[currentModel]);
  //renderPass.draw(cubeVerticesCount, gridSize * gridSize);
  //renderPass.drawIndexed(cubeIndexCount, gridSize * gridSize);
  //renderPass.drawIndexed(cubeIndexCount);
  renderPass.drawIndexed(cubeIndexCount, models[currentModel].voxelCount);

  renderPass.end();

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

  requestAnimationFrame(render)
}
//requestAnimationFrame(render);

const matrixSize = 4 * 4 * 4;
const float4Size = 4 * 4;
const float3Size = 4 * 3;
const viewUniformBufferSize = (matrixSize * 2) + (float4Size * 1);

const viewUniformBuffer = device.createBuffer({
  label: "view uniform buffer",
  size: viewUniformBufferSize,
  usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
});

const cameraStateSize = (float3Size * 3) + (4 * 3) + (float3Size * 1);

const viewStateAddress = module.instance.exports.mem_alloc(viewUniformBufferSize);
const cameraStateAddress = module.instance.exports.mem_alloc(cameraStateSize);

module.instance.exports.camera_init(cameraStateAddress);

const gltfRenderer = await loadGltf(device, canvasFormat, viewUniformBuffer, memory2, module);
const lightRenderer = await loadLight(device, canvasFormat, viewUniformBuffer);


const KEY = {
  A: 65,
  B: 66,
  C: 67,
  D: 68,
  E: 69,
  F: 70,
  G: 71,
  H: 72,
  I: 73,
  J: 74,
  K: 75,
  L: 76,
  M: 77,
  N: 78,
  O: 79,
  P: 80,
  Q: 81,
  R: 82,
  S: 83,
  T: 84,
  U: 85,
  V: 86,
  W: 87,
  X: 88,
  Y: 89,
  Z: 90,

  Left: 37,
  Up: 38,
  Right: 39,
  Down: 40,
};

const keyCodes = new Set(Object.values(KEY));
const keyState = {};
for (const code of Object.values(KEY)) {
  keyState[code] = false;
}

function handleKeydown(e)
{
  //console.log(e.keyCode);
  if (keyCodes.has(e.keyCode)) {
    keyState[e.keyCode] = true;
  }
}

function handleKeyup(e)
{
  if (keyCodes.has(e.keyCode)) {
    keyState[e.keyCode] = false;
  }
}

window.addEventListener('keydown', handleKeydown, false);
window.addEventListener('keyup', handleKeyup, false);

function updateView()
{
  const lu = keyState[KEY.W] === true;
  const ll = keyState[KEY.A] === true;
  const ld = keyState[KEY.S] === true;
  const lr = keyState[KEY.D] === true;

  const lal = keyState[KEY.Q] === true;
  const lar = keyState[KEY.E] === true;

  const ru = keyState[KEY.Up] === true;
  const rl = keyState[KEY.Left] === true;
  const rd = keyState[KEY.Down] === true;
  const rr = keyState[KEY.Right] === true;

  module.instance.exports.camera_move(cameraStateAddress,
                                      lu, ll, ld, lr,
                                      lal, lar,
                                      ru, rl, rd, rr);

  const aspect = canvas.width / canvas.height;

  module.instance.exports.camera_view_projection(cameraStateAddress,
                                                 viewStateAddress,
                                                 aspect);

  device.queue.writeBuffer(viewUniformBuffer, 0,
                           memory2.buffer, viewStateAddress,
                           viewUniformBufferSize);
}

function render2()
{
  recreateDepth();
  const colorView = context.getCurrentTexture().createView();

  updateView();

  {
    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',
      },
    });

    gltfRenderer.render(renderPass);
    lightRenderer.render(renderPass);

    renderPass.end();

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

  requestAnimationFrame(render2);
}
requestAnimationFrame(render2);