/**
 * spz-v4-adapter.js
 * SPZ v4 (NGSP / ZSTD) → SPZ v3 (gzip) 前端转换适配器
 *
 * 背景：
 *   SparkJS (spark.module.js) 内置的 SpzReader 只支持 gzip 单流的 SPZ v1~v3，
 *   遇到后端新产出的 NGSP v4（明文 32 字节头 + TOC + 多路 ZSTD 流）会抛
 *   "Invalid gzip header"。
 *
 *   本适配器在浏览器端把 v4 文件解包、按 v3 布局重组并用 gzip 重新压缩，
 *   再交给 SparkJS 原有解码路径，无需改动 spark.module.js。
 *
 * v4 文件布局（依据 nianticlabs/spz src/cc/load-spz.cc）：
 *   Bytes 0–31                NgspFileHeader（明文，小端）
 *   Bytes 32..(tbo-1)         扩展区（flags & 0x2 时存在，本适配器跳过）
 *   Bytes tbo..(tbo+N*16-1)   TOC：N × [compressedSize u64, uncompressedSize u64]
 *   之后                      N 路独立 ZSTD 流
 *   流顺序固定为：positions, alphas, colors, scales, rotations, sh
 *   （长度为 0 的流会被跳过，例如 shDegree=0 时没有 sh 流）
 *
 * v3 负载布局：16 字节头 + 上述流按相同顺序直接拼接，整体 gzip。
 *
 * 依赖：fzstd（纯 JS zstd 解码器，本地文件 ./fzstd.esm.js）
 * 浏览器要求：支持 CompressionStream("gzip")（Chrome/Edge 80+，Firefox 113+，Safari 16.4+）
 */

import { decompress as zstdDecompress } from "./fzstd.esm.js";

const NGSP_MAGIC = 0x5053474e; // "NGSP"
const V4_HEADER_SIZE = 32;
const TOC_ENTRY_SIZE = 16;
// 各 SH 阶数对应的每点系数向量数（vec3 数）
const SH_DIM_FOR_DEGREE = [0, 3, 8, 15, 24];

/** 是否为 gzip 文件（SPZ v1~v3） */
export function isGzipSpz(bytes) {
    return bytes.length >= 2 && bytes[0] === 0x1f && bytes[1] === 0x8b;
}

/** 是否为 NGSP v4 文件 */
export function isNgspV4(bytes) {
    if (bytes.length < V4_HEADER_SIZE) return false;
    const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
    return view.getUint32(0, true) === NGSP_MAGIC && view.getUint32(4, true) >= 4;
}

/**
 * 解析 v4 头部
 * @returns {{numPoints:number, shDegree:number, fractionalBits:number, flags:number, numStreams:number, tocByteOffset:number}}
 */
export function parseNgspHeader(bytes) {
    const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);
    if (view.getUint32(0, true) !== NGSP_MAGIC) {
        throw new Error("[spz-v4-adapter] 不是 NGSP 文件（magic 不匹配）");
    }
    const version = view.getUint32(4, true);
    if (version < 4) {
        throw new Error(`[spz-v4-adapter] NGSP 版本 ${version} 不是 v4，请走 gzip 路径`);
    }
    return {
        version,
        numPoints: view.getUint32(8, true),
        shDegree: view.getUint8(12),
        fractionalBits: view.getUint8(13),
        flags: view.getUint8(14),
        numStreams: view.getUint8(15),
        tocByteOffset: view.getUint32(16, true),
    };
}

/**
 * 读取 TOC 并逐路 zstd 解压
 * @returns {Uint8Array[]} 解压后的流，顺序为 positions/alphas/colors/scales/rotations/sh
 */
function decompressStreams(bytes, header) {
    const { numStreams, tocByteOffset } = header;
    const view = new DataView(bytes.buffer, bytes.byteOffset, bytes.byteLength);

    if (tocByteOffset < V4_HEADER_SIZE || tocByteOffset + numStreams * TOC_ENTRY_SIZE > bytes.length) {
        throw new Error("[spz-v4-adapter] TOC 偏移越界，文件损坏");
    }

    // 每路流预期的解压后大小（与官方加载器 checkSizes 一致，用于校验）
    const n = header.numPoints;
    const shDim = SH_DIM_FOR_DEGREE[header.shDegree] ?? 0;
    const expectedSizes = [n * 9, n * 1, n * 3, n * 3, n * 4, n * shDim * 3].filter((s) => s > 0);
    if (expectedSizes.length !== numStreams) {
        throw new Error(
            `[spz-v4-adapter] 流数量不匹配：头部声明 ${numStreams} 路，按 shDegree=${header.shDegree} 预期 ${expectedSizes.length} 路`
        );
    }

    const streams = new Array(numStreams);
    let offset = tocByteOffset + numStreams * TOC_ENTRY_SIZE;
    for (let i = 0; i < numStreams; i++) {
        const e = tocByteOffset + i * TOC_ENTRY_SIZE;
        // u64 小端；实际文件不会超 4GB，取低 32 位并校验高位
        const compressedSize = view.getUint32(e, true) + view.getUint32(e + 4, true) * 4294967296;
        const uncompressedSize = view.getUint32(e + 8, true) + view.getUint32(e + 12, true) * 4294967296;

        if (uncompressedSize !== expectedSizes[i]) {
            throw new Error(
                `[spz-v4-adapter] 第 ${i} 路流大小不匹配：文件声明 ${uncompressedSize}，预期 ${expectedSizes[i]}`
            );
        }
        if (offset + compressedSize > bytes.length) {
            throw new Error(`[spz-v4-adapter] 第 ${i} 路流超出文件末尾，文件损坏`);
        }

        streams[i] = zstdDecompress(bytes.subarray(offset, offset + compressedSize));
        if (streams[i].length !== uncompressedSize) {
            throw new Error(`[spz-v4-adapter] 第 ${i} 路流解压后大小异常`);
        }
        offset += compressedSize;
    }
    return streams;
}

/**
 * SH degree 4 → 3 截断：v4 每点 sh 占 24*3=72 字节，v3 最多 15*3=45 字节，
 * 取每点前 45 字节（低阶系数排前面），高阶光泽细节会丢失，但模型可正常显示。
 */
function truncateSh4To3(shStream, numPoints) {
    const out = new Uint8Array(numPoints * 45);
    for (let i = 0; i < numPoints; i++) {
        out.set(shStream.subarray(i * 72, i * 72 + 45), i * 45);
    }
    return out;
}

/** 用浏览器原生 CompressionStream 做 gzip 压缩 */
async function gzipCompress(bytes) {
    if (typeof CompressionStream === "undefined") {
        throw new Error("[spz-v4-adapter] 当前浏览器不支持 CompressionStream，无法重打包 gzip");
    }
    const stream = new Blob([bytes]).stream().pipeThrough(new CompressionStream("gzip"));
    return new Uint8Array(await new Response(stream).arrayBuffer());
}

/**
 * 核心转换：NGSP v4 (zstd) 字节 → SPZ v3 (gzip) 字节
 * @param {Uint8Array} bytes - v4 文件完整字节
 * @returns {Promise<Uint8Array>} 可直接喂给 SparkJS SplatMesh(fileBytes) 的 gzip SPZ 字节
 */
export async function convertNgspV4ToGzipSpz(bytes) {
    const header = parseNgspHeader(bytes);
    const streams = decompressStreams(bytes, header);

    let shDegree = header.shDegree;
    // SparkJS 只支持 SH degree ≤ 3，degree 4 截断处理
    if (shDegree === 4) {
        streams[5] = truncateSh4To3(streams[5], header.numPoints);
        shDegree = 3;
        console.warn("[spz-v4-adapter] SH degree 4 已截断为 3（SparkJS 上限），会损失少量高光细节");
    }

    // 拼 v3 负载：16 字节头 + 各流顺序拼接
    const totalSize = 16 + streams.reduce((acc, s) => acc + s.length, 0);
    const payload = new Uint8Array(totalSize);
    const view = new DataView(payload.buffer);
    view.setUint32(0, NGSP_MAGIC, true);
    view.setUint32(4, 3, true); // version = 3
    view.setUint32(8, header.numPoints, true);
    view.setUint8(12, shDegree);
    view.setUint8(13, header.fractionalBits);
    view.setUint8(14, header.flags & 0x1); // 只保留 antialiased 位，去掉扩展标志
    view.setUint8(15, 0);

    let offset = 16;
    for (const s of streams) {
        payload.set(s, offset);
        offset += s.length;
    }

    return gzipCompress(payload);
}

/**
 * 统一入口：给任意 SPZ 字节，返回 SparkJS 可读的 gzip SPZ 字节
 * @param {Uint8Array} bytes
 * @param {{onConvertStart?:Function, onConvertDone?:Function}} [hooks]
 * @returns {Promise<Uint8Array>}
 */
export async function ensureSparkCompatibleSpz(bytes, hooks = {}) {
    if (isGzipSpz(bytes)) return bytes; // v1~v3，原样返回
    if (isNgspV4(bytes)) {
        console.log("[spz-v4-adapter] 检测到 NGSP v4 (ZSTD)，开始转换…");
        hooks.onConvertStart?.();
        const t0 = performance.now();
        const out = await convertNgspV4ToGzipSpz(bytes);
        console.log(
            `[spz-v4-adapter] 转换完成：${(bytes.length / 1024).toFixed(0)}KB → ${(out.length / 1024).toFixed(0)}KB，` +
            `耗时 ${(performance.now() - t0).toFixed(0)}ms`
        );
        hooks.onConvertDone?.(out);
        return out;
    }
    throw new Error("[spz-v4-adapter] 无法识别的 SPZ 格式（既不是 gzip v1~v3，也不是 NGSP v4）");
}
