// import {Glslang} from '@webgpu/glslang/dist/web-devel-onefile/glslang';
import {Glslang} from '@webgpu/glslang/dist/web-devel/glslang.onefile';
import * as utils from './fixture';

export class BufferOp {
  device: GPUDevice;
  queue: GPUQueue;
  glslang: Glslang;
  commandQueue: GPUCommandEncoder[];
  times: [];
  resultMatrixBuffer: GPUBuffer;
  resultMatrixBufferSize: number;
  // TODO; This workaround memory usage. But not work.
  // gpuReadBuffer: GPUBuffer;
  shape: Uint32Array;
  computePipeline: any;
  bindGroup: any;
  bufferID: number;
  freeBuffers: Map<number, GPUBuffer[]> = new Map();
  // enableTimeStamp: boolean;
  constructor(device: GPUDevice, glslang: Glslang) {
    this.device = device;
    this.queue = device.defaultQueue;
    this.glslang = glslang;
    this.commandQueue = [];
    this.bufferID = 0;
    // this.enableTimeStamp = false;
  }

  /*
    var x = new Int32Array(1);
    x[0] = 17;
    console.log(x[0]);
    console.log(x.length);
    this.uploadToGPU(x, 4, GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST);
  */
  /*
  private uploadToGPU(values: ArrayBufferView, byteSize: number, usage: any) {
    const buffer = this.device.createBuffer({size: byteSize, usage});

    if (values) {
      buffer.setSubData(0, values as ArrayBufferView);
      values = null;
    }
    return buffer;
  }
  */
  // TIMESTAMP
  /*
  async getQueryTime(dstBuffer: GPUBuffer) {
    const dstStagingBuffer = this.device.createBuffer({
      size: 16,
      usage: GPUBufferUsage.MAP_READ | GPUBufferUsage.COPY_DST,
    });

    const commandEncoder = this.device.createCommandEncoder();
    commandEncoder.copyBufferToBuffer(dstBuffer, 0, dstStagingBuffer, 0, 16);
    this.device.defaultQueue.submit([commandEncoder.finish()]);

    // @ts-ignore
    const arrayBuf = new BigUint64Array(await dstStagingBuffer.mapReadAsync());

    // Time delta is a gpu ticks, we need convert to time using gpu frequency
    const timeDelta = arrayBuf[1] - arrayBuf[0];
    // gpu frequency can be got from console in chromium. For Winodws on
    // UHD630/RX560 = 25000000;
    const frequency = 25000000;
    // 1 ms = 1000 000 ns
    const timeInMS = Number(timeDelta) * 1000 / frequency;
    console.log('timestamp: ' + timeInMS + 'ms');
    dstStagingBuffer.destroy();
    dstBuffer.destroy();
    return timeInMS;
  }
  */

  compile(
      firstMatrix: Float32Array|Uint32Array,
      secondMatrix: Float32Array|Uint32Array, shape: Uint32Array,
      computeShaderCode: any) {
    this.shape = shape;
    /*
    const [gpuBufferFirstMatrix, arrayBufferFirstMatrix] =
        this.device.createBufferMapped({
          size: (firstMatrix as Float32Array).byteLength,
          usage: GPUBufferUsage.STORAGE
        });
    new Float32Array(arrayBufferFirstMatrix).set(firstMatrix);
    gpuBufferFirstMatrix.unmap();
    */
    const gpuBufferFirstMatrix = this.device.createBuffer({
      mappedAtCreation: true,
      size: (firstMatrix as Float32Array).byteLength,
      usage: GPUBufferUsage.STORAGE
    });
    // TODO: turn this into type of firstMatrix.
    new Float32Array(gpuBufferFirstMatrix.getMappedRange()).set(firstMatrix);
    gpuBufferFirstMatrix.unmap();

    /*
    const [gpuBufferSecondMatrix, arrayBufferSecondMatrix] =
        this.device.createBufferMapped({
          size: (secondMatrix as Float32Array).byteLength,
          usage: GPUBufferUsage.STORAGE
        });
    new Float32Array(arrayBufferSecondMatrix).set(secondMatrix);
    gpuBufferSecondMatrix.unmap();
    */
    const gpuBufferSecondMatrix = this.device.createBuffer({
      mappedAtCreation: true,
      size: (secondMatrix as Float32Array).byteLength,
      usage: GPUBufferUsage.STORAGE
    });
    // TODO: turn this into type of secondMatrix.
    new Float32Array(gpuBufferSecondMatrix.getMappedRange()).set(secondMatrix);
    gpuBufferSecondMatrix.unmap();

    // Result Matrix
    this.resultMatrixBufferSize =
        Float32Array.BYTES_PER_ELEMENT * (shape[4] * shape[5]);
    this.resultMatrixBuffer = this.device.createBuffer({
      size: this.resultMatrixBufferSize,
      usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC
    });
    // console.log(this.resultMatrixBufferSize);

    // This works.
    /*
    const [shapeBuffer, shapeMapping] = this.device.createBufferMapped({
      size: shape.byteLength,
      usage: GPUBufferUsage.UNIFORM,
    });
    new Uint32Array(shapeMapping).set(shape);
    shapeBuffer.unmap();
    */
    const shapeBuffer = this.device.createBuffer({
      mappedAtCreation: true,
      size: shape.byteLength,
      usage: GPUBufferUsage.UNIFORM
    });
    // TODO: turn this into type of shape.
    new Uint32Array(shapeBuffer.getMappedRange()).set(shape);
    shapeBuffer.unmap();

    // This works too.
    /*
     const shapeBuffer = this.uploadToGPU(
         shape, shape.byteLength,
         GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_SRC |
             GPUBufferUsage.COPY_DST);
     */
    this.createLayout(
        gpuBufferFirstMatrix, gpuBufferSecondMatrix, shapeBuffer,
        computeShaderCode);
    this.releaseBuffer(gpuBufferFirstMatrix);
    this.releaseBuffer(gpuBufferSecondMatrix);
    this.releaseBuffer(this.resultMatrixBuffer);
    this.releaseBuffer(shapeBuffer);
    return;
  }

  private createLayout(
      gpuBufferFirstMatrix: GPUBuffer, gpuBufferSecondMatrix: GPUBuffer,
      shapeBuffer: GPUBuffer, computeShaderCode: any) {
    // Use old layout:
    // Bind group layout and bind group
    const bindGroupLayout = this.device.createBindGroupLayout({
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.COMPUTE,
          type: 'uniform-buffer'
        },
        {
          binding: 1,
          visibility: GPUShaderStage.COMPUTE,
          type: 'readonly-storage-buffer'
        },
        {
          binding: 2,
          visibility: GPUShaderStage.COMPUTE,
          type: 'readonly-storage-buffer'
        },
        {binding: 3, visibility: GPUShaderStage.COMPUTE, type: 'storage-buffer'}
      ]
    });

    this.bindGroup = this.device.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        {binding: 0, resource: {buffer: shapeBuffer}},
        {binding: 1, resource: {buffer: gpuBufferFirstMatrix}},
        {binding: 2, resource: {buffer: gpuBufferSecondMatrix}},
        {binding: 3, resource: {buffer: this.resultMatrixBuffer}}
      ]
    });
    // Old layout end.

    // Pipeline setup
    const result =
        this.glslang.compileGLSLZeroCopy(computeShaderCode, 'compute', false);
    if (result.data.length === 0) {
      throw new Error('Shader compilation failed');
    }
    this.computePipeline = this.device.createComputePipeline({
      // For new layout, remove this line.
      layout: this.device.createPipelineLayout(
          {bindGroupLayouts: [bindGroupLayout]}),
      computeStage: {
        module: this.device.createShaderModule({code: result.data}),
        entryPoint: 'main'
      }
    });

    /* New layout
    const bindGroup = this.device.createBindGroup({
      layout: computePipeline.getBindGroupLayout(0),
      entries: [
        {binding: 0, resource: {buffer: shapeBuffer}},
        {binding: 1, resource: {buffer: gpuBufferFirstMatrix}},
        {binding: 2, resource: {buffer: gpuBufferSecondMatrix}},
        {binding: 3, resource: {buffer: this.resultMatrixBuffer}}
      ]
    });
        */
  }

  compileAndRunSync(workGroupSize: [number, number, number], workPerThread: [
    number, number, number
  ] = [1, 1, 1]) {
    // TODO: figure out how to return non const two values.
    // if (mode == 0) {
    return this.dispatchAndSubmit(
        this.computePipeline, this.bindGroup, this.shape[0], this.shape[1],
        workGroupSize, workPerThread);
  }

  private dispatchAndSubmit(
      computePipeline: any, bindGroup: any, dispatchX: number,
      dispatchY: number, workGroupSize: [number, number, number],
      workPerThread: [number, number, number] = [1, 1, 1]) {
    // TIMESTAMP
    // TODO: necessary to destroy querySet?
    /*
    let querySet: any;
    let dstBuffer: GPUBuffer;
    if (this.enableTimeStamp) {
      querySet = this.device.createQuerySet({
        type: 'timestamp',
        count: 2,
      });
      dstBuffer = this.device.createBuffer({
        size: 16,
        usage: GPUBufferUsage.COPY_SRC | GPUBufferUsage.COPY_DST,
      });
    }
    */
    // Commands submission
    const commandEncoder = this.device.createCommandEncoder();

    const passEncoder = commandEncoder.beginComputePass();
    /*
    if (this.enableTimeStamp) {
      passEncoder.writeTimestamp(querySet, 0);
    }
    */
    passEncoder.setPipeline(computePipeline);
    passEncoder.setBindGroup(0, bindGroup);

    // passEncoder.dispatch(dispatchX, dispatchY);
    if (workGroupSize[1] == 1 && workGroupSize[2] == 1) {
      passEncoder.dispatch(
          Math.ceil(
              dispatchX * dispatchY / workGroupSize[0] / workPerThread[0] /
              workPerThread[1]),
          1);
    } else {
      passEncoder.dispatch(
          Math.ceil(dispatchX / workGroupSize[0] / workPerThread[0]),
          Math.ceil(dispatchY / workGroupSize[1] / workPerThread[1]));
      /*
  console.log(
      'dispatchX / workGroupSize[0] / workPerThread[0] =' +
      dispatchX / workGroupSize[0] / workPerThread[0]);
  console.log(
      'dispatchY / workGroupSize[1] / workPerThread[1] =' +
      dispatchY / workGroupSize[1] / workPerThread[1]);
      */
    }
    /*
    if (this.enableTimeStamp) {
      passEncoder.writeTimestamp(querySet, 1);
    }
    */
    passEncoder.endPass();
    /*
    if (this.enableTimeStamp) {
      commandEncoder.resolveQuerySet(querySet, 0, 2, dstBuffer, 0);
    }
    */
    // Submit GPU commands.

    const gpuCommands = commandEncoder.finish();
    this.device.defaultQueue.submit([gpuCommands]);
    /*
    if (this.enableTimeStamp) {
      await this.getQueryTime(dstBuffer);
    }
    */

    // return (this.now() - start);
    return true;
  }

  async data() {
    const arrayBuffer = await this.getBufferData();
    return new Float32Array(arrayBuffer);
  }

  async getBufferData() {
    // TODO: this is only used in readback.
    const gpuReadBuffer = this.device.createBuffer({
      size: this.resultMatrixBufferSize,
      usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ
    });

    // Commands submission
    const commandEncoder = this.device.createCommandEncoder();
    // Encode commands for copying buffer to buffer.
    commandEncoder.copyBufferToBuffer(
        this.resultMatrixBuffer /* source buffer */, 0 /* source offset */,
        gpuReadBuffer /* destination buffer */, 0 /* destination offset */,
        this.resultMatrixBufferSize /* size */
    );

    // Submit GPU commands.
    const gpuCommands = commandEncoder.finish();
    this.device.defaultQueue.submit([gpuCommands]);
    // Read buffer.
    // const mapped = await gpuReadBuffer.mapReadAsync();
    // const arrayBuffer = mapped.slice(0);
    await gpuReadBuffer.mapAsync(GPUMapMode.READ);
    const arrayBuffer = gpuReadBuffer.getMappedRange().slice(0);
    gpuReadBuffer.unmap();
    gpuReadBuffer.destroy();
    return arrayBuffer;
  }

  // ---------Below code is not used!-------------------
  // Do not USE!
  private async dispatchAndSubmitWithFence(
      computePipeline: any, bindGroup: any, dispatchX: number,
      dispatchY: number, workGroupSize: [number, number, number]) {
    const start = utils.now();
    this.dispatchAndSubmit(
        this.computePipeline, this.bindGroup, this.shape[0], this.shape[1],
        workGroupSize);
    const fence = this.queue.createFence();
    this.queue.signal(fence, 1);
    await fence.onCompletion(1);
    console.log((utils.now() - start).toFixed(2));
  }

  // TODO: Float32Array is bad. And buffer is bad.
  async compileAndRun(workGroupSize: [number, number, number]) {
    // TODO: figure out how to return non const two values.
    // if (mode == 0) {
    return await this.dispatchAndSubmitWithFence(
        this.computePipeline, this.bindGroup, this.shape[0], this.shape[1],
        workGroupSize);
  }

  getBufferKey() {
    return this.bufferID++;
  }

  // Below is not used currently.
  releaseBuffer(buffer: GPUBuffer) {
    if (this.freeBuffers == null) {
      return;
    }

    const key = this.getBufferKey();
    if (!this.freeBuffers.has(key)) {
      this.freeBuffers.set(key, []);
    }

    this.freeBuffers.get(key).push(buffer);
  }
  // Call this after execute.
  disposeReadBackBuffer() {
    if (this.freeBuffers == null) {
      return;
    }

    this.freeBuffers.forEach((buffers, key) => {
      buffers.forEach(buff => {
        // console.log(' freeBuffers destroy key = ' + key);
        // buff.unmap();
        buff.destroy();
      });
    });
  }

  dispose() {
    this.disposeReadBackBuffer();
  }
}