import { Geometry } from '@viamrobotics/sdk';
import { Color, Vector3, Vector4 } from 'three';
import { NURBSCurve } from 'three/addons/curves/NURBSCurve.js';
import {} from './buf/draw/v1/drawing_pb';
import { RenderArmModels } from './buf/draw/v1/scene_pb';
import { traits } from './ecs';
import { parseMetadata } from './metadata';
import { createBufferGeometry } from './attribute';
import { asColor, asFloat32Array, asOpacity, isPerVertexColors, STRIDE } from './buffer';
import { rgbaToHex } from './color';
const vec3 = new Vector3();
const colorUtil = new Color();
/**
 * Merges scene-level metadata (grid, camera, point/line settings) into the
 * current viewer settings. Millimetre values from the proto are converted
 * to metres.
 */
export const applySceneMetadata = (settings, metadata) => {
    const next = { ...settings };
    if (metadata.grid !== undefined) {
        next.grid = metadata.grid;
    }
    if (metadata.gridCellSize !== undefined) {
        next.gridCellSize = metadata.gridCellSize / 1000;
    }
    if (metadata.gridSectionSize !== undefined) {
        next.gridSectionSize = metadata.gridSectionSize / 1000;
    }
    if (metadata.gridFadeDistance !== undefined) {
        next.gridFadeDistance = metadata.gridFadeDistance / 1000;
    }
    if (metadata.pointSize !== undefined) {
        next.pointSize = metadata.pointSize / 1000;
    }
    if (metadata.pointColor !== undefined) {
        next.pointColor = rgbaToHex(metadata.pointColor);
    }
    if (metadata.lineWidth !== undefined) {
        next.lineWidth = metadata.lineWidth / 1000;
    }
    if (metadata.linePointSize !== undefined) {
        next.lineDotSize = metadata.linePointSize / 1000;
    }
    if (metadata.renderArmModels !== undefined) {
        next.renderArmModels = getRenderArmModels(metadata.renderArmModels);
    }
    if (metadata.sceneCamera?.cameraType.case === 'orthographicCamera') {
        next.cameraMode = 'orthographic';
    }
    else if (metadata.sceneCamera?.cameraType.case === 'perspectiveCamera') {
        next.cameraMode = 'perspective';
    }
    return next;
};
/**
 * Spawns ECS entities for every transform and drawing in a {@link Snapshot}.
 *
 * Each transform produces one entity with Name, Pose, Parent, Geometry, and
 * optional Color/Opacity traits. Each drawing produces one or more entities
 * depending on the geometry type (arrows, points, line, nurbs, model, or
 * simple shapes like box/sphere/capsule).
 *
 * @returns The spawned entities — pass them to {@link destroyEntities} to
 *          clean up before loading a new snapshot.
 */
export const spawnSnapshotEntities = (world, snapshot) => {
    const entities = [];
    for (const transform of snapshot.transforms) {
        entities.push(spawnTransformEntity(world, transform));
    }
    for (const drawing of snapshot.drawings) {
        const drawingEntities = spawnEntitiesFromDrawing(world, drawing);
        for (const e of drawingEntities) {
            entities.push(e);
        }
    }
    return entities;
};
/**
 * Destroys a list of entities that are still alive in the given world.
 * Silently skips entities that have already been removed.
 */
export const destroyEntities = (world, entities) => {
    for (const entity of entities) {
        if (world.has(entity)) {
            entity.destroy();
        }
    }
};
const getRenderArmModels = (renderArmModels) => {
    switch (renderArmModels) {
        case RenderArmModels.COLLIDERS: {
            return 'colliders';
        }
        case RenderArmModels.UNSPECIFIED:
        case RenderArmModels.COLLIDERS_AND_MODEL: {
            return 'colliders+model';
        }
        case RenderArmModels.MODEL: {
            return 'model';
        }
    }
};
const spawnTransformEntity = (world, transform) => {
    const entityTraits = [
        traits.Name(transform.referenceFrame),
        traits.Geometry(transform.physicalObject ?? Geometry.fromJson({})),
        traits.Center(transform.physicalObject?.center),
        traits.SnapshotAPI,
        traits.Removable,
    ];
    const poseInFrame = transform.poseInObserverFrame;
    entityTraits.push(traits.Pose(poseInFrame?.pose), traits.Parent(poseInFrame?.referenceFrame));
    if (transform.metadata) {
        const metadata = parseMetadata(transform.metadata.fields);
        if (metadata.colors)
            addColorTraits(entityTraits, metadata.colors);
    }
    return world.spawn(...entityTraits);
};
const spawnEntitiesFromDrawing = (world, drawing) => {
    const entities = [];
    const poseInFrame = drawing.poseInObserverFrame;
    const parent = poseInFrame?.referenceFrame;
    const { geometryType } = drawing.physicalObject ?? {};
    if (geometryType?.case === 'arrows') {
        const poses = asFloat32Array(geometryType.value.poses);
        const colors = drawing.metadata?.colors;
        const entityTraits = [
            traits.Name(drawing.referenceFrame),
            traits.Pose(poseInFrame?.pose),
            traits.Positions(poses),
        ];
        if (parent) {
            entityTraits.push(traits.Parent(parent));
        }
        if (colors) {
            entityTraits.push(traits.Colors(colors));
        }
        const entity = world.spawn(...entityTraits, traits.Arrows({ headAtPose: true }), traits.Instances({ count: poses.length / STRIDE.ARROWS }), traits.SnapshotAPI, traits.Removable);
        entities.push(entity);
    }
    else if (geometryType?.case === 'model') {
        const rootEntityTraits = [
            traits.Name(drawing.referenceFrame),
            traits.Pose(poseInFrame?.pose),
            traits.ReferenceFrame,
        ];
        if (parent) {
            rootEntityTraits.push(traits.Parent(parent));
        }
        const rootEntity = world.spawn(...rootEntityTraits, traits.SnapshotAPI, traits.Removable);
        entities.push(rootEntity);
        let i = 1;
        for (const asset of geometryType.value.assets) {
            const entityTraits = [
                traits.Name(`${drawing.referenceFrame} model ${i++}`),
                traits.Parent(drawing.referenceFrame),
            ];
            if (geometryType.value.scale) {
                entityTraits.push(traits.Scale(geometryType.value.scale));
            }
            if (asset.content.case === 'url') {
                entityTraits.push(traits.GLTF({
                    source: { url: asset.content.value },
                    animationName: geometryType.value.animationName ?? '',
                }));
            }
            else if (asset.content.value) {
                entityTraits.push(traits.GLTF({
                    source: { glb: asset.content.value },
                    animationName: geometryType.value.animationName ?? '',
                }));
            }
            const entity = world.spawn(...entityTraits, traits.SnapshotAPI, traits.Removable);
            entities.push(entity);
        }
    }
    else {
        const entityTraits = [
            traits.Name(drawing.referenceFrame),
            traits.Pose(poseInFrame?.pose),
        ];
        if (parent && parent !== 'world') {
            entityTraits.push(traits.Parent);
        }
        if (drawing.physicalObject?.center) {
            entityTraits.push(traits.Center(drawing.physicalObject.center));
        }
        if (geometryType?.case === 'line') {
            const positions = asFloat32Array(geometryType.value.positions);
            for (let i = 0, l = positions.length; i < l; i += 1) {
                positions[i] *= 0.001;
            }
            entityTraits.push(traits.LinePositions(positions), traits.LineWidth(geometryType.value.lineWidth), traits.PointSize((geometryType.value.pointSize ?? 0) * 0.001));
            // Lines pack exactly 2 colors: [lineColor, pointColor]
            const colors = drawing.metadata?.colors;
            if (colors && colors.length >= STRIDE.COLORS_RGB) {
                const stride = colors.length % STRIDE.COLORS_RGBA === 0 ? STRIDE.COLORS_RGBA : STRIDE.COLORS_RGB;
                asColor(colors, colorUtil, 0);
                entityTraits.push(traits.Color({ r: colorUtil.r, g: colorUtil.g, b: colorUtil.b }));
                if (colors.length >= stride * 2) {
                    asColor(colors, colorUtil, stride);
                    entityTraits.push(traits.PointColor({ r: colorUtil.r, g: colorUtil.g, b: colorUtil.b }));
                    if (stride === STRIDE.COLORS_RGBA) {
                        entityTraits.push(traits.Opacity(asOpacity(colors, 1, 3)));
                    }
                }
            }
        }
        else if (geometryType?.case === 'points') {
            const positions = asFloat32Array(geometryType.value.positions);
            for (let i = 0, l = positions.length; i < l; i += 1) {
                positions[i] *= 0.001;
            }
            const colors = drawing.metadata?.colors;
            const numPoints = positions.length / STRIDE.POSITIONS;
            const vertexColors = colors && isPerVertexColors(colors, numPoints) ? colors : undefined;
            const geometry = createBufferGeometry(positions, vertexColors);
            entityTraits.push(traits.BufferGeometry(geometry));
            if (colors && !vertexColors) {
                addColorTraits(entityTraits, colors);
            }
            if (geometryType.value.pointSize) {
                entityTraits.push(traits.PointSize(geometryType.value.pointSize * 0.001));
            }
            entityTraits.push(traits.Points);
        }
        else if (geometryType?.case === 'nurbs') {
            const { degree = 3, knots: knotsBuffer, weights: weightsBuffer, controlPoints: controlPointsBuffer, } = geometryType.value;
            const knots = [...asFloat32Array(knotsBuffer)];
            const weights = weightsBuffer
                ? [...asFloat32Array(weightsBuffer)]
                : [];
            const controlPointsArray = [...asFloat32Array(controlPointsBuffer)];
            const controlPoints = [];
            for (let i = 0, j = 0, l = controlPointsArray.length / STRIDE.NURBS_CONTROL_POINTS; i < l; i += STRIDE.NURBS_CONTROL_POINTS, j += 1) {
                vec3
                    .set(controlPointsArray[0], controlPointsArray[1], controlPointsArray[2])
                    .multiplyScalar(0.001);
                controlPoints.push(new Vector4(vec3.x, vec3.y, vec3.z, weights[j] ?? 0));
            }
            const curve = new NURBSCurve(degree, knots, controlPoints);
            const numPoints = 600;
            const points = new Float32Array(numPoints * 3);
            const l = numPoints * 3;
            for (let i = 0; i < l; i += 3) {
                curve.getPointAt(i / (l - 1), vec3);
                points[i + 0] = vec3.x;
                points[i + 1] = vec3.y;
                points[i + 2] = vec3.z;
            }
            entityTraits.push(traits.LinePositions(points));
            const colors = drawing.metadata?.colors;
            if (colors) {
                addColorTraits(entityTraits, colors);
            }
        }
        else {
            // Box, sphere, capsule, and other geometry shapes with a single color
            const colors = drawing.metadata?.colors;
            if (colors) {
                addColorTraits(entityTraits, colors);
            }
        }
        const entity = world.spawn(...entityTraits, traits.SnapshotAPI, traits.Removable);
        entities.push(entity);
    }
    return entities;
};
const addColorTraits = (entityTraits, bytes) => {
    asColor(bytes, colorUtil);
    entityTraits.push(traits.Color(colorUtil));
    const isRgba = bytes.length % STRIDE.COLORS_RGBA === 0;
    if (isRgba)
        entityTraits.push(traits.Opacity(asOpacity(bytes)));
};
