import { assert } from "../../../core/assert.js";
import Quaternion from "../../../core/geom/Quaternion.js";
import Vector3 from "../../../core/geom/Vector3.js";

/**
 * Visual smoothing for the divergence between predicted and authoritative state.
 *
 * After client reconciliation, the simulation position jumps to the
 * authoritative value. Without smoothing, this looks like a teleport. With
 * smoothing, we keep an `error` offset between the rendered position and the
 * simulation position; the error decays toward zero over a few frames, hiding
 * the snap.
 *
 * Algorithm (Glenn Fiedler, *State Synchronization*):
 *   - On reconciliation: `error += pre_correction_render_pos - post_correction_sim_pos`.
 *     The rendered position (simulation_pos + error) stays continuous across the jump.
 *   - Each frame: `error *= decay_factor(|error|)` where decay is adaptive —
 *     slow for small errors (smooth blend), fast for large errors (snappy).
 *
 * Per-entity state is stored sparsely in a `Map`. Entities not tracked have
 * implicit zero error.
 *
 * @author Alex Goldring
 * @copyright Company Named Limited (c) 2026
 */

const POSITION_SMALL_THRESHOLD = 0.25;
const POSITION_LARGE_THRESHOLD = 1.0;
const FACTOR_SMALL = 0.95;
const FACTOR_LARGE = 0.85;

function position_decay_factor(magnitude) {
    if (magnitude <= POSITION_SMALL_THRESHOLD) return FACTOR_SMALL;
    if (magnitude >= POSITION_LARGE_THRESHOLD) return FACTOR_LARGE;
    const t = (magnitude - POSITION_SMALL_THRESHOLD) / (POSITION_LARGE_THRESHOLD - POSITION_SMALL_THRESHOLD);
    return FACTOR_SMALL + t * (FACTOR_LARGE - FACTOR_SMALL);
}

class EntitySmoothing {
    constructor() {
        this.position_error = new Vector3(0, 0, 0);
        this.orientation_error = new Quaternion(0, 0, 0, 1);
    }
}

export class SmoothingState {

    /** @type {Map<number, EntitySmoothing>} @private */
    #per_entity = new Map();

    /**
     * Begin tracking an entity. Initial error is zero.
     * @param {number} entity_id
     */
    track(entity_id) {
        assert.isNonNegativeInteger(entity_id, 'entity_id');
        if (!this.#per_entity.has(entity_id)) {
            this.#per_entity.set(entity_id, new EntitySmoothing());
        }
    }

    /**
     * Stop tracking an entity. Releases the stored error state.
     * @param {number} entity_id
     */
    untrack(entity_id) {
        assert.isNonNegativeInteger(entity_id, 'entity_id');
        this.#per_entity.delete(entity_id);
    }

    /**
     * Returns true if `entity_id` has tracked smoothing state.
     * @param {number} entity_id
     * @returns {boolean}
     */
    is_tracked(entity_id) {
        return this.#per_entity.has(entity_id);
    }

    /**
     * Apply a correction. Called when the simulation position has been snapped
     * from `pre_render_pos` (the position the entity rendered at last frame,
     * i.e. `sim_old + error_old`) to `post_correction_pos` (the new
     * authoritative sim position). The error term absorbs the jump so the
     * rendered position (`sim + error`) stays continuous.
     *
     * Note this is an assignment, not an accumulation: the prior error is
     * already baked into `pre_render_pos`, so adding it again would
     * double-count and drift visibly across successive corrections.
     *
     * @param {number} entity_id
     * @param {Vector3|ArrayLike<number>} pre_render_pos position the entity was about to render at
     * @param {Vector3|ArrayLike<number>} post_correction_pos the new authoritative simulation position
     */
    apply_position_correction(entity_id, pre_render_pos, post_correction_pos) {
        assert.isNonNegativeInteger(entity_id, 'entity_id');
        let s = this.#per_entity.get(entity_id);
        if (s === undefined) {
            s = new EntitySmoothing();
            this.#per_entity.set(entity_id, s);
        }
        // error = pre - post (rendered = post + error = pre — continuous).
        s.position_error[0] = pre_render_pos[0] - post_correction_pos[0];
        s.position_error[1] = pre_render_pos[1] - post_correction_pos[1];
        s.position_error[2] = pre_render_pos[2] - post_correction_pos[2];
    }

    /**
     * Same as {@link apply_position_correction} but for orientation. Stores the
     * relative rotation `pre_render_rot * inverse(post_correction_rot)` as the
     * error quaternion so that the render path's `error * post == pre`
     * invariant holds.
     *
     * Like position, this is an assignment: `pre_render_rot` already
     * incorporates the prior error (it's the orientation the entity rendered
     * at last frame), so composing with `error_old` again would double-count
     * and drift across successive corrections.
     *
     * @param {number} entity_id
     * @param {Quaternion|ArrayLike<number>} pre_render_rot
     * @param {Quaternion|ArrayLike<number>} post_correction_rot
     */
    apply_orientation_correction(entity_id, pre_render_rot, post_correction_rot) {
        assert.isNonNegativeInteger(entity_id, 'entity_id');
        let s = this.#per_entity.get(entity_id);
        if (s === undefined) {
            s = new EntitySmoothing();
            this.#per_entity.set(entity_id, s);
        }
        // error = pre * inverse(post). With render = error * post, this
        // satisfies render = pre — continuous across the correction.
        const px = pre_render_rot[0], py = pre_render_rot[1], pz = pre_render_rot[2], pw = pre_render_rot[3];
        const ix = -post_correction_rot[0], iy = -post_correction_rot[1], iz = -post_correction_rot[2], iw = post_correction_rot[3];
        s.orientation_error[0] = pw*ix + px*iw + py*iz - pz*iy;
        s.orientation_error[1] = pw*iy - px*iz + py*iw + pz*ix;
        s.orientation_error[2] = pw*iz + px*iy - py*ix + pz*iw;
        s.orientation_error[3] = pw*iw - px*ix - py*iy - pz*iz;
    }

    /**
     * Decay both error terms toward zero. Position uses adaptive decay
     * (slow for small errors, fast for large); orientation uses a flat
     * decay matching the position decay at the current position-error magnitude.
     *
     * Should be called once per render frame for each tracked entity.
     *
     * @param {number} entity_id
     */
    decay(entity_id) {
        assert.isNonNegativeInteger(entity_id, 'entity_id');
        const s = this.#per_entity.get(entity_id);
        if (s === undefined) return;

        const ex = s.position_error[0], ey = s.position_error[1], ez = s.position_error[2];
        const magnitude = Math.sqrt(ex*ex + ey*ey + ez*ez);
        const factor = position_decay_factor(magnitude);

        s.position_error[0] = ex * factor;
        s.position_error[1] = ey * factor;
        s.position_error[2] = ez * factor;

        // Orientation: slerp toward identity by (1 - factor). The existing
        // Quaternion.slerp mutates `this`; use the static helper variant if
        // available. For now, do a manual lerp+normalize since identity is
        // (0,0,0,1) and we want to bias toward it.
        const qx = s.orientation_error[0], qy = s.orientation_error[1], qz = s.orientation_error[2], qw = s.orientation_error[3];
        // Choose the shorter arc (q vs -q both represent the same rotation).
        const sign = qw < 0 ? -1 : 1;
        const ax = qx * sign, ay = qy * sign, az = qz * sign, aw = qw * sign;
        // Lerp toward identity by `1 - factor` (factor stays, identity gets the rest).
        const nx = ax * factor;
        const ny = ay * factor;
        const nz = az * factor;
        const nw = aw * factor + (1 - factor);
        // Normalize.
        const len = Math.sqrt(nx*nx + ny*ny + nz*nz + nw*nw);
        if (len > 0) {
            const inv = 1 / len;
            s.orientation_error[0] = nx * inv;
            s.orientation_error[1] = ny * inv;
            s.orientation_error[2] = nz * inv;
            s.orientation_error[3] = nw * inv;
        }
    }

    /**
     * Get the render position by adding the error term to the simulation position.
     *
     * @param {number} entity_id
     * @param {Vector3|ArrayLike<number>} sim_pos input
     * @param {Vector3|ArrayLike<number>} out_pos output (mutated)
     */
    render_position(entity_id, sim_pos, out_pos) {
        const s = this.#per_entity.get(entity_id);
        if (s === undefined) {
            out_pos[0] = sim_pos[0];
            out_pos[1] = sim_pos[1];
            out_pos[2] = sim_pos[2];
        } else {
            out_pos[0] = sim_pos[0] + s.position_error[0];
            out_pos[1] = sim_pos[1] + s.position_error[1];
            out_pos[2] = sim_pos[2] + s.position_error[2];
        }
    }

    /**
     * Direct access to the per-entity error terms (or undefined if not tracked).
     *
     * @param {number} entity_id
     * @returns {{position_error: Vector3, orientation_error: Quaternion}|undefined}
     */
    state_for(entity_id) {
        return this.#per_entity.get(entity_id);
    }

    /**
     * Number of tracked entities.
     * @returns {number}
     */
    tracked_count() {
        return this.#per_entity.size;
    }
}
