src/routes/viewers/BabylonViewer.ts

import type { IViewer } from "./IViewer";
import * as BABYLON from "@babylonjs/core";
import "@babylonjs/loaders/glTF";
import "@babylonjs/loaders/OBJ";

export class BabylonViewer implements IViewer {
    canvas: HTMLCanvasElement;

    engine: BABYLON.Engine;
    scene: BABYLON.Scene;
    camera: BABYLON.ArcRotateCamera;

    vertexCount: number = 0;
    topoOnly: boolean = false;

    private _originalMaterials: Map<BABYLON.AbstractMesh, BABYLON.Material> = new Map();
    private _originalVertexColors: Map<BABYLON.AbstractMesh, BABYLON.Nullable<BABYLON.FloatArray>> = new Map();
    private _wireframes: Array<BABYLON.Mesh> = [];

    private _solidColor = new BABYLON.Color4(1, 1, 1, 1);
    private _wireframeMaterial: BABYLON.StandardMaterial;
    private _solidMaterial: BABYLON.StandardMaterial;

    constructor(canvas: HTMLCanvasElement) {
        this.canvas = canvas;

        this.engine = new BABYLON.Engine(canvas, true);

        this.scene = new BABYLON.Scene(this.engine);
        this.scene.clearColor = BABYLON.Color4.FromHexString("#000000FF");
        this.scene.imageProcessingConfiguration.exposure = 1.3;

        this.camera = new BABYLON.ArcRotateCamera(
            "camera",
            Math.PI / 3,
            Math.PI / 3,
            30,
            BABYLON.Vector3.Zero(),
            this.scene
        );
        this.camera.angularSensibilityY = 1000;
        this.camera.panningSensibility = 500;
        this.camera.wheelPrecision = 5;
        this.camera.inertia = 0.9;
        this.camera.panningInertia = 0.9;
        this.camera.lowerRadiusLimit = 3;
        this.camera.upperRadiusLimit = 100;
        this.camera.setTarget(BABYLON.Vector3.Zero());
        this.camera.attachControl(this.canvas, true);
        this.camera.onAfterCheckInputsObservable.add(() => {
            this.camera.wheelPrecision = 150 / this.camera.radius;
            this.camera.panningSensibility = 10000 / this.camera.radius;
        });

        this._wireframeMaterial = new BABYLON.StandardMaterial("wireframe", this.scene);
        this._wireframeMaterial.diffuseColor = new BABYLON.Color3(0, 0, 0);
        this._wireframeMaterial.emissiveColor = new BABYLON.Color3(0.7, 0.7, 0.7);
        this._wireframeMaterial.disableLighting = true;
        this._wireframeMaterial.wireframe = true;

        this._solidMaterial = new BABYLON.StandardMaterial("solid", this.scene);
        this._solidMaterial.diffuseColor = new BABYLON.Color3(0.8, 0.8, 0.8);
        this._solidMaterial.alphaMode = 0;

        this.handleResize = this.handleResize.bind(this);
        window.addEventListener("resize", this.handleResize);
        this.canvas.addEventListener("wheel", (e) => e.preventDefault());
    }

    handleResize() {
        this.engine.resize();
    }

    async loadScene(url: string, loadingBarCallback?: (progress: number) => void, topoOnly?: boolean) {
        this.topoOnly = topoOnly ?? false;

        // Load scene
        await BABYLON.SceneLoader.AppendAsync("", url, this.scene, (event) => {
            const progress = event.loaded / event.total;
            loadingBarCallback?.(progress);
        });

        // Dispose of all cameras and lights
        this.scene.cameras.forEach((camera) => {
            if (camera !== this.camera) {
                camera.dispose();
            }
        });
        this.scene.lights.forEach((light) => {
            light.dispose();
        });

        // Add lights
        const hemi = new BABYLON.HemisphericLight("hemi", new BABYLON.Vector3(0, 1, 0), this.scene);
        hemi.intensity = 0.5;
        hemi.diffuse = new BABYLON.Color3(1, 1, 1);
        hemi.groundColor = new BABYLON.Color3(1, 1, 1);

        const sun = new BABYLON.DirectionalLight("sun", new BABYLON.Vector3(-0.5, -1, -0.5), this.scene);
        sun.intensity = 1;
        sun.diffuse = new BABYLON.Color3(1, 1, 1);

        // Center and scale model
        const parentNode = new BABYLON.TransformNode("parent", this.scene);
        const standardSize = 10;
        let scaleFactor = 1;
        let center = BABYLON.Vector3.Zero();
        if (this.scene.meshes.length > 0) {
            let bounds = this.scene.meshes[0].getBoundingInfo().boundingBox;
            let min = bounds.minimumWorld;
            let max = bounds.maximumWorld;

            for (let i = 1; i < this.scene.meshes.length; i++) {
                bounds = this.scene.meshes[i].getBoundingInfo().boundingBox;
                min = BABYLON.Vector3.Minimize(min, bounds.minimumWorld);
                max = BABYLON.Vector3.Maximize(max, bounds.maximumWorld);
            }

            const extent = max.subtract(min).scale(0.5);
            const size = extent.length();

            center = BABYLON.Vector3.Center(min, max);

            scaleFactor = standardSize / size;
        }
        this.vertexCount = 0;
        this.scene.meshes.forEach((mesh) => {
            mesh.setParent(parentNode);
            if (mesh.getTotalVertices() > 0) {
                this.vertexCount += mesh.getTotalVertices();
            }
        });
        parentNode.position = center.scale(-1 * scaleFactor);
        parentNode.scaling.scaleInPlace(scaleFactor);

        if (this.topoOnly) {
            this.setRenderMode("wireframe");
        }

        // Run render loop
        this.engine.runRenderLoop(() => {
            this.scene.render();
        });
    }

    dispose() {
        if (this.scene) {
            this.scene.dispose();
        }
        if (this.engine) {
            this.engine.dispose();
        }
        this._originalMaterials.clear();
        window.removeEventListener("resize", this.handleResize);
        this.canvas.removeEventListener("wheel", (e) => e.preventDefault());
    }

    async capture(): Promise<string | null> {
        if (!this.engine || !this.camera) return null;
        const cachedColor = this.scene.clearColor;
        this.scene.clearColor = BABYLON.Color4.FromHexString("#00000000");
        let data = await new Promise<string>((resolve) => {
            BABYLON.Tools.CreateScreenshotUsingRenderTarget(this.engine, this.camera, 512, (result) => {
                resolve(result);
            });
        });
        this.scene.clearColor = cachedColor;
        return data;
    }

    setRenderMode(mode: string) {
        if (mode === "wireframe") {
            this.scene.meshes.forEach((mesh) => {
                const vertexData = mesh.getVerticesData(BABYLON.VertexBuffer.ColorKind);
                if (vertexData) {
                    this._originalVertexColors.set(mesh, vertexData);

                    const newVertexData = new Array<number>(vertexData.length);
                    for (let i = 0; i < vertexData.length; i += 4) {
                        newVertexData[i] = this._solidColor.r;
                        newVertexData[i + 1] = this._solidColor.g;
                        newVertexData[i + 2] = this._solidColor.b;
                        newVertexData[i + 3] = this._solidColor.a;
                    }
                    mesh.setVerticesData(BABYLON.VertexBuffer.ColorKind, newVertexData);
                }

                const material = mesh.material as BABYLON.StandardMaterial;
                if (material) {
                    this._originalMaterials.set(mesh, material);
                    mesh.material = this._solidMaterial;
                }

                const wireframeMesh = mesh.clone(mesh.name + "_wireframe", null) as BABYLON.Mesh;
                wireframeMesh.material = this._wireframeMaterial;
                this._wireframes.push(wireframeMesh);
            });
        } else {
            this._wireframes.forEach((mesh) => {
                mesh.dispose();
            });
            this._wireframes = [];

            this.scene.meshes.forEach((mesh) => {
                const vertexData = this._originalVertexColors.get(mesh);
                if (vertexData) {
                    mesh.setVerticesData(BABYLON.VertexBuffer.ColorKind, vertexData);
                }

                const material = this._originalMaterials.get(mesh);
                if (material) {
                    mesh.material = material;
                }
            });
        }
    }
}