// 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/
import { wrap } from "comlink";
import { NumericType } from "@foxglove/schemas";
import { estimateObjectSize } from "@foxglove/studio-base/players/messageMemoryEstimation";
import { deserializeDepthFrame } from "./depth";
import { convertNv12ToI420, deserializeEncodedFrameH264 } from "./h264";
import { deserializeEncodedFrameH265 } from "./h265";
import { deserializeMJPEGFrame } from "./mjpeg";
import { Profile } from "../../../protobuf.generated/EncodedFrame";
import { Type } from "../../../protobuf.generated/ImgFrame";
import { parseMessage } from "../../utils";
import { Logger } from "../../../utils/logger";
import { uint8ArrayToUint16Array } from "../pointcloud/utils";
const RGBD_POINTCLOUD_WORKER_COUNT = 4;
let cachedRGBDHasGPU;
let rgbdHasGPUPromise;
let rgbdPointCloudWorkers;
let rgbdPointCloudWorkerIndex = 0;
const logger = Logger.getLogger();
function ensureEvenByteLength(data) {
    if (data.length % 2 === 0) {
        return data;
    }
    const padded = new Uint8Array(data.length + 1);
    padded.set(data);
    return padded;
}
async function detectRGBDGPUSupport() {
    if (cachedRGBDHasGPU != undefined) {
        return cachedRGBDHasGPU;
    }
    if (rgbdHasGPUPromise) {
        return await rgbdHasGPUPromise;
    }
    rgbdHasGPUPromise = (async () => {
        if (typeof navigator.gpu === "undefined") {
            cachedRGBDHasGPU = false;
            return false;
        }
        try {
            const adapter = await navigator.gpu.requestAdapter();
            cachedRGBDHasGPU = adapter != undefined;
            return cachedRGBDHasGPU;
        }
        catch {
            cachedRGBDHasGPU = false;
            return false;
        }
        finally {
            rgbdHasGPUPromise = undefined;
        }
    })();
    return await rgbdHasGPUPromise;
}
function getRGBDPointCloudWorkerSlot() {
    if (!rgbdPointCloudWorkers) {
        rgbdPointCloudWorkers = Array.from({ length: RGBD_POINTCLOUD_WORKER_COUNT }, () => {
            const worker = new Worker(new URL("../pointcloud/pointcloudFromDepth.worker", import.meta.url), {
                type: "module",
            });
            return {
                worker,
                remote: wrap(worker),
                activeJobs: 0,
            };
        });
    }
    const slots = rgbdPointCloudWorkers;
    const candidateSlots = slots
        .map((slot, index) => ({ slot, index }))
        .sort((left, right) => {
        if (left.slot.activeJobs !== right.slot.activeJobs) {
            return left.slot.activeJobs - right.slot.activeJobs;
        }
        const leftDistance = (left.index - rgbdPointCloudWorkerIndex + slots.length) % slots.length;
        const rightDistance = (right.index - rgbdPointCloudWorkerIndex + slots.length) % slots.length;
        return leftDistance - rightDistance;
    });
    const selected = candidateSlots[0];
    rgbdPointCloudWorkerIndex = (selected.index + 1) % slots.length;
    return selected.slot;
}
function repackYuv420pToI420(imageFrame, width, height) {
    const specs = imageFrame.fb ?? imageFrame.sourceFb;
    const source = ensureEvenByteLength(imageFrame.data);
    const yStride = specs?.stride || width;
    const uvWidth = Math.floor(width / 2);
    const uvHeight = Math.floor(height / 2);
    const uvStride = yStride > width ? Math.floor(yStride / 2) : uvWidth;
    const yPlaneSize = width * height;
    const uvPlaneSize = uvWidth * uvHeight;
    const uOffset = specs?.p1Offset || yStride * height;
    const vOffset = specs?.p2Offset || uOffset + uvStride * uvHeight;
    const i420 = new Uint8Array(yPlaneSize + 2 * uvPlaneSize);
    for (let row = 0; row < height; row++) {
        const srcOffset = row * yStride;
        const dstOffset = row * width;
        i420.set(source.subarray(srcOffset, srcOffset + width), dstOffset);
    }
    for (let row = 0; row < uvHeight; row++) {
        const srcUOffset = uOffset + row * uvStride;
        const srcVOffset = vOffset + row * uvStride;
        const dstOffset = row * uvWidth;
        i420.set(source.subarray(srcUOffset, srcUOffset + uvWidth), yPlaneSize + dstOffset);
        i420.set(source.subarray(srcVOffset, srcVOffset + uvWidth), yPlaneSize + uvPlaneSize + dstOffset);
    }
    return i420;
}
function toMatrix3x3(matrix) {
    const rows = matrix?.arrays;
    if (!rows || rows.length < 3) {
        return undefined;
    }
    const normalizedRows = rows
        .slice(0, 3)
        .map((row) => row?.values.slice(0, 3).map((value) => value ?? 0));
    if (normalizedRows.some((row) => !row || row.length < 3)) {
        return undefined;
    }
    return normalizedRows;
}
function multiply3x3(left, right) {
    return [
        [
            left[0][0] * right[0][0] + left[0][1] * right[1][0] + left[0][2] * right[2][0],
            left[0][0] * right[0][1] + left[0][1] * right[1][1] + left[0][2] * right[2][1],
            left[0][0] * right[0][2] + left[0][1] * right[1][2] + left[0][2] * right[2][2],
        ],
        [
            left[1][0] * right[0][0] + left[1][1] * right[1][0] + left[1][2] * right[2][0],
            left[1][0] * right[0][1] + left[1][1] * right[1][1] + left[1][2] * right[2][1],
            left[1][0] * right[0][2] + left[1][1] * right[1][2] + left[1][2] * right[2][2],
        ],
        [
            left[2][0] * right[0][0] + left[2][1] * right[1][0] + left[2][2] * right[2][0],
            left[2][0] * right[0][1] + left[2][1] * right[1][1] + left[2][2] * right[2][1],
            left[2][0] * right[0][2] + left[2][1] * right[1][2] + left[2][2] * right[2][2],
        ],
    ];
}
function extractIntrinsicsFromTransformation(transformation) {
    const sourceIntrinsicMatrix = toMatrix3x3(transformation?.sourceIntrinsicMatrix);
    if (!sourceIntrinsicMatrix) {
        return undefined;
    }
    const transformationMatrix = toMatrix3x3(transformation?.transformationMatrix);
    const effectiveIntrinsicMatrix = transformationMatrix
        ? multiply3x3(transformationMatrix, sourceIntrinsicMatrix)
        : sourceIntrinsicMatrix;
    const fx = effectiveIntrinsicMatrix[0][0];
    const fy = effectiveIntrinsicMatrix[1][1];
    const cx = effectiveIntrinsicMatrix[0][2];
    const cy = effectiveIntrinsicMatrix[1][2];
    if (![fx, fy, cx, cy].every((value) => Number.isFinite(value))) {
        return undefined;
    }
    return {
        left: {
            focalLenght: { x: fx, y: fy },
            principalPoint: { x: cx, y: cy },
        },
        right: {
            focalLenght: { x: fx, y: fy },
            principalPoint: { x: cx, y: cy },
        },
    };
}
function resolveRGBDIntrinsics(message, fallbackIntrinsics) {
    return (extractIntrinsicsFromTransformation(message.depthImgFrame?.transformation) ??
        extractIntrinsicsFromTransformation(message.depthEncodedFrame?.transformation) ??
        extractIntrinsicsFromTransformation(message.colorImgFrame?.transformation) ??
        extractIntrinsicsFromTransformation(message.colorEncodedFrame?.transformation) ??
        fallbackIntrinsics);
}
async function extractEncodedColorFrame(encodedFrame, topicDecoders, topic) {
    const decodeFrame = async () => await new Promise((resolve, reject) => {
        const callback = (event) => {
            resolve(event.message);
        };
        const rejectDecode = (error) => {
            reject(error);
        };
        const args = {
            topic,
            message: encodedFrame,
            callback,
        };
        if (encodedFrame.profile === Profile.AVC) {
            void deserializeEncodedFrameH264({ ...args, topicDecoders, onError: rejectDecode }).catch(rejectDecode);
        }
        else if (encodedFrame.profile === Profile.JPEG) {
            void deserializeMJPEGFrame(args).catch(rejectDecode);
        }
        else if (encodedFrame.profile === Profile.HEVC) {
            void deserializeEncodedFrameH265({ ...args, topicDecoders, onError: rejectDecode }).catch(rejectDecode);
        }
        else {
            reject(new Error(`Unsupported RGBD color profile ${Profile[encodedFrame.profile] ?? encodedFrame.profile}`));
        }
    });
    const decoded = await decodeFrame();
    return {
        data: ensureEvenByteLength(decoded.data),
        width: decoded.width,
        height: decoded.height,
    };
}
function extractColorFrame(encodedFrame, imageFrame, topicDecoders, topic) {
    if (encodedFrame) {
        return extractEncodedColorFrame(encodedFrame, topicDecoders, topic);
    }
    if (imageFrame) {
        const { width, height, type } = parseMessage(imageFrame);
        const specs = imageFrame.fb ?? imageFrame.sourceFb;
        if (type === Type.NV12) {
            const stride = specs?.stride || width;
            const uvOffset = specs?.p1Offset || stride * height;
            return Promise.resolve({
                data: convertNv12ToI420(ensureEvenByteLength(imageFrame.data), width, height, stride, stride, 0, uvOffset),
                width,
                height,
            });
        }
        if (type === Type.YUV420p) {
            return Promise.resolve({
                data: repackYuv420pToI420(imageFrame, width, height),
                width,
                height,
            });
        }
        return Promise.reject(new Error(`Unsupported RGBD color ImgFrame type ${Type[type]}`));
    }
    return Promise.resolve({
        data: new Uint8Array(),
        width: 0,
        height: 0,
    });
}
function extractDepthFrame(encodedFrame, imageFrame) {
    if (encodedFrame) {
        return ensureEvenByteLength(encodedFrame.data);
    }
    if (imageFrame) {
        return ensureEvenByteLength(imageFrame.data);
    }
    return new Uint8Array();
}
function parseDimensions(message) {
    const obj = {
        depthWidth: 0,
        depthHeight: 0,
    };
    if (message.depthEncodedFrame) {
        const { width, height } = message.depthEncodedFrame ?? { width: 1920, height: 1200 };
        obj.depthWidth = width;
        obj.depthHeight = height;
    }
    else if (message.depthImgFrame) {
        const { width, height } = parseMessage(message.depthImgFrame);
        obj.depthWidth = width;
        obj.depthHeight = height;
    }
    return obj;
}
export async function deserializeRGBDFrameToPointCloud({ topic, topicDecoders, cameraIntrinsics, maxStereoDepth, message, callback, }) {
    const workerSlot = getRGBDPointCloudWorkerSlot();
    workerSlot.activeJobs += 1;
    const { receiveTime } = parseMessage(message);
    const { depthWidth, depthHeight } = parseDimensions(message);
    const rgbdIntrinsics = resolveRGBDIntrinsics(message, cameraIntrinsics);
    const hasGPU = await detectRGBDGPUSupport();
    const colorFrame = await extractColorFrame(message.colorEncodedFrame, message.colorImgFrame, topicDecoders, topic);
    const uint8DepthFrames = extractDepthFrame(message.depthEncodedFrame, message.depthImgFrame);
    const depthFrames = uint8ArrayToUint16Array(uint8DepthFrames);
    try {
        const pointCloud = await workerSlot.remote.depthToPointcloudGPU([depthFrames], depthWidth, depthHeight, rgbdIntrinsics.right.focalLenght.x, rgbdIntrinsics.right.focalLenght.y, rgbdIntrinsics.right.principalPoint.x, rgbdIntrinsics.right.principalPoint.y, [colorFrame.data], colorFrame.width, colorFrame.height, maxStereoDepth, { hasGPU });
        const foxgloveMessage = {
            timestamp: message.tsDevice ?? receiveTime,
            frame_id: `pointcloud-${topic}-frame`,
            point_stride: 16,
            pose: {
                position: { x: 0, y: 0, z: 0 },
                orientation: { x: 0, y: 0, z: 1, w: 0 }
            },
            width: depthWidth,
            height: depthHeight,
            fields: [
                { name: "x", offset: 0, type: NumericType.FLOAT32 },
                { name: "y", offset: 4, type: NumericType.FLOAT32 },
                { name: "z", offset: 8, type: NumericType.FLOAT32 },
                { name: "red", offset: 12, type: NumericType.UINT8 },
                { name: "green", offset: 13, type: NumericType.UINT8 },
                { name: "blue", offset: 14, type: NumericType.UINT8 },
                { name: "alpha", offset: 15, type: NumericType.UINT8 },
            ],
            data: pointCloud[0]
        };
        callback({
            topic,
            receiveTime,
            message: foxgloveMessage,
            sizeInBytes: estimateObjectSize(foxgloveMessage),
            schemaName: "foxglove.PointCloud.Colored"
        });
        if (message.depthImgFrame) {
            void deserializeDepthFrame({
                topic,
                message: message.depthImgFrame,
                callback,
                maxStereoDepth: maxStereoDepth / 1000,
            });
        }
    }
    finally {
        workerSlot.activeJobs = Math.max(0, workerSlot.activeJobs - 1);
    }
}
//# sourceMappingURL=rgbd-to-pointcloud.js.map