"use strict";

function updateGetModelProgress(name, loaded, total) {
  const progressElement = document.getElementById('model-progress');
  if (total === 0) {
    progressElement.innerHTML = `Model ${name} is already in local cache`;
    return;
  }
  const percent = ((loaded / total) * 100).toFixed(2);
  let progress;
  if (loaded >= total) {
    progress = "Downloaded";
  } else {
    progress = "Downloading";
  }
  progressElement.innerHTML = `${progress} model ${name}: Total ${total} bytes, loaded ${loaded} bytes, ${percent}%`;
}

// Get model via Origin Private File System
async function getModelOPFS(name, url, updateModel) {
  const root = await navigator.storage.getDirectory();
  let fileHandle;

  async function updateFile() {
    const response = await fetch(url);
    const buffer = await readResponse(name, response, updateGetModelProgress);
    fileHandle = await root.getFileHandle(name, { create: true });
    const writable = await fileHandle.createWritable();
    await writable.write(buffer);
    await writable.close();
    return buffer;
  }

  if (updateModel) {
    return await updateFile();
  }

  try {
    fileHandle = await root.getFileHandle(name);
    const blob = await fileHandle.getFile();
    updateGetModelProgress(name, 0, 0);
    return await blob.arrayBuffer();
  } catch (e) {
    return await updateFile();
  }
}

async function readResponse(name, response, progressCallback) {
  const contentLength = response.headers.get('Content-Length');
  let total = parseInt(contentLength ?? '0');
  let buffer = new Uint8Array(total);
  let loaded = 0;

  const reader = response.body.getReader();
  async function read() {
    const { done, value } = await reader.read();
    if (done) return;

    let newLoaded = loaded + value.length;
    if (newLoaded > total) {
      total = newLoaded;
      let newBuffer = new Uint8Array(total);
      newBuffer.set(buffer);
      buffer = newBuffer;
    }
    buffer.set(value, loaded);
    loaded = newLoaded;

    if (progressCallback) {
      progressCallback(name, loaded, total);
    }

    return read();
  }

  await read();
  return buffer;
}

function isInternal() {
  if (window.location.href.includes('intel')) {
    return true;
  } else {
    return false;
  }
}

function getModelsPath() {
  if (isInternal()) {
    return "https://wp-27.sh.intel.com/workspace/project/models/";
  } else {
    return "https://huggingface.co/webai-community/models/resolve/main/";
  }
}

function getParam(name, type, _default) {
  name = name.replace(/[\[\]]/g, "\\$&");
  let regex = new RegExp("[?&]" + name + "(=([^&#]*)|&|#|$)", "i");
  let results = regex.exec(window.location.href);
  if (!results || !results[2]) return _default;

  const result = decodeURIComponent(results[2].replace(/\+/g, " "));
  if (type === "Boolean") {
    if (result === "true") {
      return true;
    } else if (result === "false") {
      return false;
    }
  } else if (type === "Number") {
    return parseInt(result);
  } else {
    return result;
  }
}

async function loadOrt() {
  const ortUrl = getParam("ortUrl", "String", "default");
  if (ortUrl !== "default") {
    await loadScript(ortUrl);
  } else if (isInternal()) {
    await loadScript("https://wp-27.sh.intel.com/workspace/project/onnxruntime/js/web/dist/ort.webgpu.min.js");
  } else {
    await loadScript("https://cdn.jsdelivr.net/npm/onnxruntime-web@dev/dist/ort.webgpu.min.js");
  }
}

async function loadScript(url) {
  return new Promise((resolve, reject) => {
    const script = document.createElement('script');
    script.onload = resolve;
    script.onerror = reject;
    script.src = url;
    if (url.startsWith('http')) {
      script.crossOrigin = 'anonymous';
    }
    document.body.append(script);
  })
}

/*
const webgpuCanvas = new OffscreenCanvas(canvasWidth, canvasHeight);
const webglCanvas = new OffscreenCanvas(canvasWidth, canvasHeight);
const gl = webglCanvas.getContext('webgl2');
const webglTexture = gl.createTexture();
*/
function readGPUBufferSync(buffer, byteOffset, webgpuCanvas, gl, glTexture, device) {
  const bufferSize = buffer.size;
  const dataBytes = bufferSize - byteOffset;

  const canvasWidth = 8192;
  const canvasHeight = dataBytes / 4 / canvasWidth + 1;

  // Copy WebGPU buffer to WebGPU texture
  const webgpuContext = webgpuCanvas.getContext('webgpu');
  webgpuContext.configure({
    device: device,
    format: 'rgb8unorm',
    usage: GPUTextureUsage.COPY_DST,
    alphaMode: 'premultiplied',
  });
  const webgpuTexture = webgpuContext.getCurrentTexture();
  const bytesPerRow = canvasWidth * 4;
  const commandEncoder = device.createCommandEncoder();
  commandEncoder.copyBufferToTexture(
    {
      buffer,
      bytesPerRow,
      byteOffset,
    },
    {
      webgpuTexture,
    },
    {
      canvasWidth,
      canvasHeight,
    });
  device.queue.submit([commandEncoder.finish()]);
  commandEncoder = null;

  // Read WebGPU texture via WebGL
  gl.bindTexture(gl.TEXTURE_2D, webglTexture);
  gl.pixelStorei(gl.UNPACK_PREMULTIPLY_ALPHA_WEBGL, true);
  gl.texImage2D(gl.TEXTURE_2D, 0, gl.RGBA8, gl.RGBA, gl.UNSIGNED_BYTE, webgpuCanvas);
  gl.bindFramebuffer(gl.FRAMEBUFFER, gl.createFramebuffer());
  gl.framebufferTexture2D(gl.FRAMEBUFFER, gl.COLOR_ATTACHMENT0, gl.TEXTURE_2D, webglTexture, 0);
  const pixels = new Uint8Array(dataBytes);
  gl.readPixels(0, 0, width, height, gl.RGBA, gl.UNSIGNED_BYTE, pixels);

  return pixels;
}