// This Source Code Form is subject to the terms of the Mozilla Public
// License, v2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/
/// <reference types="@webgpu/types" />
import { expose } from "comlink";
import { i420ToRgbaToPointcloud } from "./utils";
function depthToPointCloudBuffer(uint16Array, width, height, fx, fy, maxSteroDepth) {
    const numPixels = width * height;
    const points = new Float32Array(numPixels * 3);
    for (let v = 0; v < height; v++) {
        for (let u = 0; u < width; u++) {
            const pixelIndex = v * width + u;
            const depthValue = uint16Array[pixelIndex] ?? 0;
            let x = 0.0;
            let y = 0.0;
            let z = 0.0;
            const depthScale = 1.0;
            const scaledDepth = depthValue * depthScale;
            if (scaledDepth > 0 && fx !== 0 && fy !== 0) {
                const imageCenterX = width * 0.5;
                const imageCenterY = height * 0.5;
                x = (u - imageCenterX) * scaledDepth / fx;
                y = (v - imageCenterY) * scaledDepth / fy;
                z = scaledDepth;
                // Filter out points beyond maxSteroDepth
                if (z > maxSteroDepth) {
                    x = 0.0;
                    y = 0.0;
                    z = 0.0;
                }
            }
            const pointOffset = pixelIndex * 3;
            points[pointOffset] = x;
            points[pointOffset + 1] = y;
            points[pointOffset + 2] = z;
        }
    }
    const buffer = new ArrayBuffer(numPixels * 16);
    const view = new DataView(buffer);
    const pointsU8 = new Uint8Array(points.buffer);
    for (let i = 0; i < numPixels; i++) {
        const xyzOffsetBytes = i * 12;
        const outputOffsetBytes = i * 16;
        for (let b = 0; b < 12; ++b) {
            // biome-ignore lint/style/noNonNullAssertion: <explanation>
            view.setUint8(outputOffsetBytes + b, pointsU8[xyzOffsetBytes + b]);
        }
        view.setUint8(outputOffsetBytes + 12, 0); // R
        view.setUint8(outputOffsetBytes + 13, 0); // G
        view.setUint8(outputOffsetBytes + 14, 0); // B
        view.setUint8(outputOffsetBytes + 15, 255); // A (Opaque)
    }
    return new Uint8Array(buffer);
}
const WORKGROUP_SIZE_X = 128;
const POINT_STRIDE_BYTES = 16;
const shader = `
struct Intrinsics {
  fx: f32,
  fy: f32,
  cx: f32,
  cy: f32,
  depthScale: f32,
  depthWidth: f32,
  depthHeight: f32,
  i420Width: f32,
  i420Height: f32,
  numFrames_f: f32,
  bytesPerI420Frame_f: f32,
  maxSteroDepth: f32,
};

struct OutputPointData {
  x: f32,
  y: f32,
  z: f32,
  color: u32,
};

@group(0) @binding(0) var<storage, read> depthBuffer: array<u32>;
@group(0) @binding(1) var<storage, read> intrinsics: Intrinsics;
@group(0) @binding(2) var<storage, read_write> xyzColorBuffer: array<OutputPointData>;
@group(0) @binding(3) var<storage, read> i420Buffer: array<u32>;

fn getI420Byte(byteIndexWithinFrame: u32, frameOffsetElements: u32, bufferLengthElements: u32) -> u32 {
  let globalByteIndex = (frameOffsetElements * 4u) + byteIndexWithinFrame;
  let elementIndex = globalByteIndex / 4u;
  let offsetInElement = globalByteIndex % 4u;
  if (elementIndex >= bufferLengthElements) {
    return 0u;
  }
  let packedValue = i420Buffer[elementIndex];
  return (packedValue >> (offsetInElement * 8u)) & 0xFFu;
}

@compute @workgroup_size(${WORKGROUP_SIZE_X}, 1, 1)
fn main(@builtin(global_invocation_id) global_id: vec3<u32>) {
  let pixel_idx_in_frame = global_id.x;
  let frame_idx = global_id.y;
  let depthWidth = u32(intrinsics.depthWidth);
  let depthHeight = u32(intrinsics.depthHeight);
  let i420WidthU = u32(intrinsics.i420Width);
  let i420HeightU = u32(intrinsics.i420Height);
  let numFrames = u32(intrinsics.numFrames_f);
  let bytesPerI420Frame = u32(intrinsics.bytesPerI420Frame_f);
  let numPixelsPerFrame = depthWidth * depthHeight;
  if (pixel_idx_in_frame >= numPixelsPerFrame || frame_idx >= numFrames || depthWidth == 0u || depthHeight == 0u) {
    return;
  }
  let global_output_idx = (frame_idx * numPixelsPerFrame) + pixel_idx_in_frame;
  let depth_elements_per_frame = (numPixelsPerFrame + 1u) / 2u;
  let depth_frame_offset = frame_idx * depth_elements_per_frame;
  let i420AlignedStrideBytes = (bytesPerI420Frame + 3u) / 4u * 4u;
  let i420_elements_per_frame_aligned = i420AlignedStrideBytes / 4u;
  let i420_frame_offset = frame_idx * i420_elements_per_frame_aligned;
  let depthIndex32_within_frame = pixel_idx_in_frame / 2u;
  let packedDepthU32_idx = depth_frame_offset + depthIndex32_within_frame;
  if (packedDepthU32_idx >= arrayLength(&depthBuffer)) {
    return;
  }
  let packedDepthU32 = depthBuffer[packedDepthU32_idx];
  var rawDepthU16: u32;
  if (pixel_idx_in_frame % 2u == 0u) {
    rawDepthU16 = packedDepthU32 & 0xFFFFu;
  } else {
    rawDepthU16 = (packedDepthU32 >> 16u) & 0xFFFFu;
  }
  let u = f32(pixel_idx_in_frame % depthWidth);
  let v = f32(pixel_idx_in_frame / depthWidth);
  var outputPoint: OutputPointData;
  outputPoint.x = 0.0;
  outputPoint.y = 0.0;
  outputPoint.z = 0.0;
  outputPoint.color = 0xFF000000u;
  let fx = intrinsics.fx;
  let fy = intrinsics.fy;
  let depthScale = intrinsics.depthScale;
  let maxSteroDepth = intrinsics.maxSteroDepth;
  if (rawDepthU16 > 0u && fx != 0.0 && fy != 0.0 && depthScale > 0.0) {
    let z = f32(rawDepthU16) * depthScale;
    let imageCenterX = f32(depthWidth) * 0.5;
    let imageCenterY = f32(depthHeight) * 0.5;
    let x = (u - imageCenterX) * z / fx;
    let y = (v - imageCenterY) * z / fy;
    
    // Filter out points beyond maxSteroDepth
    if (z <= maxSteroDepth) {
      outputPoint.x = x;
      outputPoint.y = y;
      outputPoint.z = z;
      
      if (i420WidthU > 0u && i420HeightU > 0u && bytesPerI420Frame > 0u) {
        let i420X = min(i420WidthU - 1u, u32(floor(u * (intrinsics.i420Width / intrinsics.depthWidth))));
        let i420Y = min(i420HeightU - 1u, u32(floor(v * (intrinsics.i420Height / intrinsics.depthHeight))));
        let frameSizeY = i420WidthU * i420HeightU;
        let uvWidth = (i420WidthU + 1u) / 2u;
        let uvHeight = (i420HeightU + 1u) / 2u;
        let frameSizeUV = uvWidth * uvHeight;
        let uOffsetBytes = frameSizeY;
        let vOffsetBytes = frameSizeY + frameSizeUV;
        let yIndexBytes = i420Y * i420WidthU + i420X;
        let uvX = i420X / 2u;
        let uvY = i420Y / 2u;
        let uvIndexBytes = uvY * uvWidth + uvX;
        let uIndexBytes = uOffsetBytes + uvIndexBytes;
        let vIndexBytes = vOffsetBytes + uvIndexBytes;
        let i420BufferTotalElements = arrayLength(&i420Buffer);
        if (yIndexBytes < bytesPerI420Frame && uIndexBytes < bytesPerI420Frame && vIndexBytes < bytesPerI420Frame && (i420_frame_offset + (vIndexBytes / 4u)) < i420BufferTotalElements) {
          let y_byte = getI420Byte(yIndexBytes, i420_frame_offset, i420BufferTotalElements);
          let u_byte = getI420Byte(uIndexBytes, i420_frame_offset, i420BufferTotalElements);
          let v_byte = getI420Byte(vIndexBytes, i420_frame_offset, i420BufferTotalElements);
          let y_f = f32(y_byte);
          let u_f = f32(u_byte) - 128.0;
          let v_f = f32(v_byte) - 128.0;
          var r_f = y_f + 1.402 * v_f;
          var g_f = y_f - 0.344136 * u_f - 0.714136 * v_f;
          var b_f = y_f + 1.772 * u_f;
          let r = u32(clamp(r_f, 0.0, 255.0)) & 0xFFu;
          let g = u32(clamp(g_f, 0.0, 255.0)) & 0xFFu;
          let b = u32(clamp(b_f, 0.0, 255.0)) & 0xFFu;
          let a = 255u;
          outputPoint.color = (r) | (g << 8u) | (b << 16u) | (a << 24u);
        }
      }
    }
  }
  if (global_output_idx < arrayLength(&xyzColorBuffer)) {
    xyzColorBuffer[global_output_idx] = outputPoint;
  }
}
`;
async function depthToPointcloudGPU(depthFrames, depthWidth, depthHeight, fx, fy, cx, cy, i420Frames, i420Width, i420Height, maxSteroDepth, { hasGPU }) {
    const numFrames = depthFrames.length;
    if (numFrames === 0 || numFrames !== i420Frames.length) {
        console.warn("GPU: Input frame arrays are empty or mismatched.");
        return [];
    }
    if (depthWidth <= 0 || depthHeight <= 0) {
        throw new Error("GPU: Invalid depth dimensions!");
    }
    if (depthFrames[0] && !(depthFrames[0] instanceof Uint16Array)) {
        throw new Error("GPU: depthFrames[0].data is NOT a Uint16Array!");
    }
    if (i420Frames[0] && !(i420Frames[0] instanceof Uint8Array)) {
        throw new Error("GPU: i420Frames[0].data is NOT a Uint8Array!");
    }
    const depthScale = 1.0;
    // eslint-disable-next-line @typescript-eslint/no-unnecessary-condition
    if (!hasGPU) {
        const results = [];
        for (let i = 0; i < numFrames; i++) {
            const depthFrame = depthFrames[i];
            const colorFrame = i420Frames[i];
            if (!depthFrame || !colorFrame) {
                console.warn(`CPU Fallback: Skipping frame ${i} due to missing data.`);
                results.push(new Uint8Array(0));
                continue;
            }
            try {
                const xyzData = depthToPointCloudBuffer(depthFrame, depthWidth, depthHeight, fx, fy, maxSteroDepth);
                results.push(i420ToRgbaToPointcloud(colorFrame, i420Width, i420Height, depthWidth, depthHeight, xyzData));
            }
            catch (cpuError) {
                console.error(`CPU Fallback: Error processing frame ${i}:`, cpuError);
                results.push(new Uint8Array(0));
            }
        }
        return results;
    }
    let device = null;
    let depthGpuBuffer = null;
    let intrinsicsBuffer = null;
    let i420GpuBuffer = null;
    let xyzColorBuffer = null;
    let readbackBuffer = null;
    try {
        const adapter = await navigator.gpu.requestAdapter();
        if (!adapter) {
            throw new Error("Failed to get GPU adapter.");
        }
        device = await adapter.requestDevice();
        const queue = device.queue;
        const numPixelsPerFrame = depthWidth * depthHeight;
        const outputFrameSize = numPixelsPerFrame * POINT_STRIDE_BYTES;
        const totalOutputBufferSize = outputFrameSize * numFrames;
        const depthBytesPerFrame = numPixelsPerFrame * 2;
        const depthElementsPerFrame = Math.ceil(depthBytesPerFrame / 4);
        const totalDepthBufferSize = depthElementsPerFrame * numFrames * 4;
        const i420BytesPerFrame = i420Frames[0]?.byteLength ?? 0;
        if (i420BytesPerFrame === 0 && numFrames > 0) {
            throw new Error("GPU: I420 frame data is empty or invalid for size calculation.");
        }
        const i420AlignedStrideBytes = Math.ceil(i420BytesPerFrame / 4) * 4;
        const totalI420BufferSize = i420AlignedStrideBytes * numFrames;
        const intrinsicsData = new Float32Array([
            fx, fy, cx, cy,
            depthScale,
            depthWidth, depthHeight,
            i420Width, i420Height,
            numFrames,
            i420BytesPerFrame,
            maxSteroDepth,
        ]);
        const intrinsicsBufferSize = Math.max(intrinsicsData.byteLength, 64);
        intrinsicsBuffer = device.createBuffer({
            label: "Intrinsics Buffer",
            size: intrinsicsBufferSize,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
        });
        queue.writeBuffer(intrinsicsBuffer, 0, intrinsicsData.buffer, intrinsicsData.byteOffset, intrinsicsData.byteLength);
        depthGpuBuffer = device.createBuffer({
            label: "Depth Buffer (Combined)",
            size: totalDepthBufferSize,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
            mappedAtCreation: true,
        });
        const mappedDepthRange = depthGpuBuffer.getMappedRange();
        const fullDepthArrayViewU32 = new Uint32Array(mappedDepthRange);
        for (let i = 0; i < numFrames; ++i) {
            const frameDataU16 = depthFrames[i];
            if (!frameDataU16 || !(frameDataU16 instanceof Uint16Array)) {
                console.warn(`GPU: Skipping depth frame ${i} due to missing/invalid data.`);
                const frameOffsetElements = i * depthElementsPerFrame;
                const elementsToWrite = depthElementsPerFrame;
                fullDepthArrayViewU32.fill(0, frameOffsetElements, frameOffsetElements + elementsToWrite);
                continue;
            }
            const frameOffsetElements = i * depthElementsPerFrame;
            for (let p = 0; p < numPixelsPerFrame; p++) {
                const u16_val = frameDataU16[p] ?? 0;
                const elementIdx = frameOffsetElements + Math.floor(p / 2);
                if (elementIdx >= fullDepthArrayViewU32.length) {
                    continue;
                }
                const isLowerHalf = (p % 2) === 0;
                if (isLowerHalf) {
                    // biome-ignore lint/style/noNonNullAssertion: <explanation>
                    fullDepthArrayViewU32[elementIdx] = (fullDepthArrayViewU32[elementIdx] & 0xFFFF0000) | u16_val;
                }
                else {
                    // biome-ignore lint/style/noNonNullAssertion: <explanation>
                    fullDepthArrayViewU32[elementIdx] = (fullDepthArrayViewU32[elementIdx] & 0x0000FFFF) | (u16_val << 16);
                }
            }
        }
        depthGpuBuffer.unmap();
        i420GpuBuffer = device.createBuffer({
            label: "I420 Buffer (Combined, Aligned)",
            size: totalI420BufferSize,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_DST,
            mappedAtCreation: true,
        });
        const mappedI420Range = i420GpuBuffer.getMappedRange();
        const fullI420ArrayViewU8 = new Uint8Array(mappedI420Range);
        for (let i = 0; i < numFrames; ++i) {
            const frameData = i420Frames[i];
            const byteOffset = i * i420AlignedStrideBytes;
            if (!frameData || !(frameData instanceof Uint8Array) || frameData.byteLength !== i420BytesPerFrame) {
                console.warn(`GPU: Skipping I420 frame ${i} due to missing/invalid data or size mismatch.`);
                fullI420ArrayViewU8.fill(0, byteOffset, byteOffset + i420AlignedStrideBytes);
                continue;
            }
            const frameTargetView = new Uint8Array(mappedI420Range, byteOffset, frameData.byteLength);
            frameTargetView.set(frameData);
            if (i420AlignedStrideBytes > frameData.byteLength) {
                const paddingStart = byteOffset + frameData.byteLength;
                const paddingLength = i420AlignedStrideBytes - frameData.byteLength;
                fullI420ArrayViewU8.fill(0, paddingStart, paddingStart + paddingLength);
            }
        }
        i420GpuBuffer.unmap();
        xyzColorBuffer = device.createBuffer({
            label: "Output XYZColor Buffer (Combined)",
            size: totalOutputBufferSize,
            usage: GPUBufferUsage.STORAGE | GPUBufferUsage.COPY_SRC,
        });
        readbackBuffer = device.createBuffer({
            label: "Readback Buffer",
            size: totalOutputBufferSize,
            usage: GPUBufferUsage.COPY_DST | GPUBufferUsage.MAP_READ,
        });
        const bindGroupLayout = device.createBindGroupLayout({
            label: "Compute Bind Group Layout",
            entries: [
                { binding: 0, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 1, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
                { binding: 2, visibility: GPUShaderStage.COMPUTE, buffer: { type: "storage" } },
                { binding: 3, visibility: GPUShaderStage.COMPUTE, buffer: { type: "read-only-storage" } },
            ],
        });
        const bindGroup = device.createBindGroup({
            label: "Compute Bind Group",
            layout: bindGroupLayout,
            entries: [
                { binding: 0, resource: { buffer: depthGpuBuffer } },
                { binding: 1, resource: { buffer: intrinsicsBuffer } },
                { binding: 2, resource: { buffer: xyzColorBuffer } },
                { binding: 3, resource: { buffer: i420GpuBuffer } },
            ],
        });
        const shaderModule = device.createShaderModule({ label: "Compute Shader Module", code: shader });
        const pipeline = device.createComputePipeline({
            label: "Compute Pipeline",
            layout: device.createPipelineLayout({ bindGroupLayouts: [bindGroupLayout] }),
            compute: { module: shaderModule, entryPoint: "main" },
        });
        const commandEncoder = device.createCommandEncoder({ label: "Compute Command Encoder" });
        const passEncoder = commandEncoder.beginComputePass({ label: "Compute Pass" });
        passEncoder.setPipeline(pipeline);
        passEncoder.setBindGroup(0, bindGroup);
        const numWorkgroupsX = Math.ceil(numPixelsPerFrame / WORKGROUP_SIZE_X);
        const numWorkgroupsY = numFrames;
        passEncoder.dispatchWorkgroups(numWorkgroupsX, numWorkgroupsY, 1);
        passEncoder.end();
        commandEncoder.copyBufferToBuffer(xyzColorBuffer, 0, readbackBuffer, 0, totalOutputBufferSize);
        queue.submit([commandEncoder.finish()]);
        await readbackBuffer.mapAsync(GPUMapMode.READ);
        const mappedRange = readbackBuffer.getMappedRange();
        const fullResultData = new Uint8Array(mappedRange.slice(0));
        readbackBuffer.unmap();
        readbackBuffer = null;
        const results = [];
        for (let i = 0; i < numFrames; i++) {
            const frameOffset = i * outputFrameSize;
            const frameEnd = Math.min(frameOffset + outputFrameSize, fullResultData.byteLength);
            const frameData = fullResultData.slice(frameOffset, frameEnd);
            if (frameData.byteLength !== outputFrameSize) {
                console.warn(`GPU: Result frame ${i} has incorrect size (${frameData.byteLength} vs ${outputFrameSize}). Padding or error occurred.`);
                results.push(new Uint8Array(outputFrameSize));
            }
            else {
                results.push(frameData);
            }
        }
        return results;
    }
    catch (error) {
        console.error("GPU: Error during WebGPU processing:", error);
        console.warn("GPU: Falling back to CPU due to error.");
        const results = [];
        for (let i = 0; i < numFrames; i++) {
            const depthFrame = depthFrames[i];
            const colorFrame = i420Frames[i];
            if (!depthFrame || !colorFrame) {
                console.warn(`CPU Fallback (error path): Skipping frame ${i} due to missing data.`);
                results.push(new Uint8Array(0));
                continue;
            }
            try {
                const xyzData = depthToPointCloudBuffer(depthFrame, depthWidth, depthHeight, fx, fy, maxSteroDepth);
                results.push(i420ToRgbaToPointcloud(colorFrame, i420Width, i420Height, depthWidth, depthHeight, xyzData));
            }
            catch (cpuError) {
                console.error(`CPU Fallback (error path): Error processing frame ${i}:`, cpuError);
                results.push(new Uint8Array(0));
            }
        }
        return results;
    }
    finally {
        depthGpuBuffer?.destroy();
        intrinsicsBuffer?.destroy();
        i420GpuBuffer?.destroy();
        xyzColorBuffer?.destroy();
        readbackBuffer?.destroy();
    }
}
expose({ depthToPointcloudGPU });
//# sourceMappingURL=pointcloudFromDepth.worker.js.map