// Axis helpers — 3d-viewer-design skill §8 (non-negotiable).
//
// Three layers, all attached at scene-init time:
//   §8a world-origin axis helper   (default ON, toggleable)
//   §8b mesh-local axis helper      (default OFF, toggleable, parented to root)
//   §8c screen-space corner triad   (always on, no toggle, viewport overlay)
//
// R/G/B for X/Y/Z — never brand-tinted. Lengths scale to scene extent so
// they neither dominate small chip previews nor disappear in big assemblies.

import { AxesViewer } from '@babylonjs/core/Debug/axesViewer';
import { ArcRotateCamera } from '@babylonjs/core/Cameras/arcRotateCamera';
import { Vector3 } from '@babylonjs/core/Maths/math.vector';
import { Viewport } from '@babylonjs/core/Maths/math.viewport';
import type { Scene } from '@babylonjs/core/scene';
import type { Camera } from '@babylonjs/core/Cameras/camera';
import type { TransformNode } from '@babylonjs/core/Meshes/transformNode';
import type { AbstractMesh } from '@babylonjs/core/Meshes/abstractMesh';

// Per-scene state stored on the scene object so consumers can grab it via
// scene.metadata.adomAxes.
export interface AdomAxesState {
    world: AxesViewer | null;
    meshLocal: Map<TransformNode | AbstractMesh, AxesViewer>;
    meshLocalEnabled: boolean;
    cornerTriad: { cam: ArcRotateCamera; axes: AxesViewer } | null;
    sceneExtent: number;
}

const STATE_KEY = '__adomAxes';

/** Get (or create) the per-scene axis state. */
function getState(scene: Scene): AdomAxesState {
    const meta = scene.metadata || (scene.metadata = {});
    if (!meta[STATE_KEY]) {
        meta[STATE_KEY] = {
            world: null,
            meshLocal: new Map(),
            meshLocalEnabled: false,
            cornerTriad: null,
            sceneExtent: 1,
        } as AdomAxesState;
    }
    return meta[STATE_KEY] as AdomAxesState;
}

/**
 * Initialize §8a + §8c at scene boot. Call once from SceneBuilder.
 * §8b helpers are added per-loaded-mesh via addMeshLocalAxes().
 */
export function initAxisHelpers(scene: Scene, mainCamera: Camera, opts?: { worldDefaultOn?: boolean }): void {
    const state = getState(scene);
    const worldOn = opts?.worldDefaultOn !== false;

    // §8a — world-origin helper. Length = 15% of scene extent (refreshed
    // by frameModel via refreshScale()). Initial guess of 1m; consumers
    // call refreshScale(scene, sceneRadius) after frameModel completes.
    const worldAxes = new AxesViewer(scene, 0.15);
    state.world = worldAxes;
    setEnabled(worldAxes, worldOn);
    // AxesViewer creates X/Y/Z meshes with .isPickable = true by default;
    // disable so they don't intercept Shift+Alt+Click etc.
    [worldAxes.xAxis, worldAxes.yAxis, worldAxes.zAxis].forEach(n => {
        n.getChildMeshes().forEach(m => { m.isPickable = false; });
        if ((n as AbstractMesh).isPickable !== undefined) (n as AbstractMesh).isPickable = false;
    });

    // §8c — screen-space corner triad: NOT YET IMPLEMENTED.
    // Tried the second-camera-viewport approach (multi-camera in
    // scene.activeCameras with a 10% bottom-left viewport on a triad
    // camera); the viewport is correctly set and the meshes carry the
    // right layerMask, but the viewer's runRenderLoop only renders the
    // primary scene.activeCamera and skips the activeCameras array.
    // Likely needs a UtilityLayerRenderer with explicit shouldRender +
    // an attached secondary engine.render() pass. Tracked as a v0.3.0
    // follow-up; for now §8a + §8b cover the bulk of "where is origin"
    // debugging value.
    state.cornerTriad = null;
}

/** Refresh world-axes scale to 15% of current scene radius (§8a). Call
 *  after frameModel(). */
export function refreshAxisScale(scene: Scene, sceneRadius: number): void {
    const state = getState(scene);
    state.sceneExtent = sceneRadius;
    const worldLen = Math.max(sceneRadius * 0.15, 0.001);
    if (state.world) {
        // AxesViewer doesn't expose a setSize after construction, so we
        // dispose + recreate when scale changes meaningfully (>10%).
        const old = state.world;
        const enabled = isEnabled(old);
        old.dispose();
        const fresh = new AxesViewer(scene, worldLen);
        [fresh.xAxis, fresh.yAxis, fresh.zAxis].forEach(n => {
            n.getChildMeshes().forEach(m => { m.isPickable = false; });
        });
        setEnabled(fresh, enabled);
        state.world = fresh;
    }
}

/** §8a — toggle visibility of the world-origin helper. */
export function setWorldAxesVisible(scene: Scene, visible: boolean): void {
    const state = getState(scene);
    if (state.world) setEnabled(state.world, visible);
}

export function isWorldAxesVisible(scene: Scene): boolean {
    const state = getState(scene);
    return state.world ? isEnabled(state.world) : false;
}

/**
 * §8b — attach a local-origin helper to a loaded root node and return it.
 * Helper is parented to `root` so it follows the mesh's transform.
 * Visibility tracks the global mesh-local toggle.
 */
export function addMeshLocalAxes(scene: Scene, root: TransformNode | AbstractMesh): AxesViewer {
    const state = getState(scene);
    if (state.meshLocal.has(root)) return state.meshLocal.get(root)!;
    // Sized to the root's bounding sphere if it has one, else fall back to
    // 40% of scene extent.
    let r = 0;
    if ((root as AbstractMesh).getBoundingInfo) {
        try { r = (root as AbstractMesh).getBoundingInfo().boundingSphere.radius; }
        catch { /* TransformNode without bounds */ }
    }
    if (!r) r = state.sceneExtent;
    const axes = new AxesViewer(scene, r * 0.4);
    [axes.xAxis, axes.yAxis, axes.zAxis].forEach(n => {
        n.parent = root;
        n.getChildMeshes().forEach(m => { m.isPickable = false; });
    });
    setEnabled(axes, state.meshLocalEnabled);
    state.meshLocal.set(root, axes);
    return axes;
}

/** §8b — toggle visibility of every mesh-local helper. */
export function setMeshLocalAxesVisible(scene: Scene, visible: boolean): void {
    const state = getState(scene);
    state.meshLocalEnabled = visible;
    for (const axes of state.meshLocal.values()) setEnabled(axes, visible);
}

export function isMeshLocalAxesVisible(scene: Scene): boolean {
    return getState(scene).meshLocalEnabled;
}

/** Drop axis helpers attached to a no-longer-present root. */
export function removeMeshLocalAxes(scene: Scene, root: TransformNode | AbstractMesh): void {
    const state = getState(scene);
    const axes = state.meshLocal.get(root);
    if (axes) { axes.dispose(); state.meshLocal.delete(root); }
}

function setEnabled(axes: AxesViewer, on: boolean): void {
    [axes.xAxis, axes.yAxis, axes.zAxis].forEach(n => n.setEnabled(on));
}
function isEnabled(axes: AxesViewer): boolean {
    return axes.xAxis.isEnabled();
}