import { assert } from "../../../core/assert.js";
import { MoveEdge } from "./MoveEdge.js";

/**
 * Coarse classification of a node's terminal status, independent of payoff.
 * The actual per-player payoff lives in {@link StateNode#outcome}.
 * @enum {number}
 */
export const StateType = {
    /** Game can still continue from this state. */
    Undecided: 0,
    /** Game ended at this state per the game's rules. Payoff in `outcome`. */
    Terminal: 1,
    /** Search hit `maxExplorationDepth` here; heuristic stands in for payoff. */
    DepthCapped: 2,
    /** Expansion produced no legal moves. Treated as terminal; payoff in `outcome`. */
    NoMoves: 3
};

/**
 * Average outcome for `player` across all playouts that visited this child node.
 * Falls back to heuristic when the node has not been rolled out yet.
 * @param {MoveEdge} move
 * @param {number} player
 * @returns {number}
 */
function pickBestScore(move, player) {
    assert.isNonNegativeInteger(player, 'player');

    const node = move.target;

    if (node.playouts === 0) {
        return node.heuristicValue[player];
    }

    return node.totalScore[player] / node.playouts;
}


let stack_pointer = 0;

/**
 *
 * @type {StateNode[]}
 */
const stack = [];


/**
 * @template State, Action
 * @author Alex Goldring
 * @copyright Company Named Limited (c) 2026
 */
export class StateNode {

    /**
     * How deep is the node in the tree
     * @type {number}
     */
    depth = 0;

    /**
     * Total number of explored playouts through this node.
     * @type {number}
     */
    playouts = 0;

    /**
     * Per-player heuristic estimate of the value of this state.
     * For terminal nodes this mirrors {@link StateNode#outcome}.
     * For internal nodes this is refined by maxN backup as children materialize.
     * Index = player ID.
     * @type {Float64Array | null}
     */
    heuristicValue = null;

    /**
     * Per-player cumulative payoff aggregated by backpropagation.
     * Average value for player `p` is `totalScore[p] / playouts` once `playouts > 0`.
     * @type {Float64Array | null}
     */
    totalScore = null;

    /**
     * Per-player payoff at this terminal node, cached at materialization.
     * `null` for non-terminal nodes (uses {@link StateNode#heuristicValue} as a stand-in).
     * @type {Float64Array | null}
     */
    outcome = null;

    /**
     * Player ID of whoever chooses the next move from this state.
     * `-1` for terminal nodes (no decision to make).
     * @type {number}
     */
    activePlayer = -1;

    /**
     * @type {null|StateNode}
     */
    parent = null;

    /**
     *
     * @type {null|MoveEdge[]}
     */
    moves = null;

    /**
     *
     * @type {StateType}
     */
    type = StateType.Undecided;

    /**
     * Allocate the per-player buffers. Must be called before any backprop or selection.
     * @param {number} numPlayers
     */
    allocate(numPlayers) {
        assert.isNonNegativeInteger(numPlayers, 'numPlayers');

        this.heuristicValue = new Float64Array(numPlayers);
        this.totalScore = new Float64Array(numPlayers);
    }

    bubbleUpHeuristicScore() {
        let r = this.parent;

        while (r !== null) {
            r.aggregateHeuristicScore();

            r = r.parent;
        }
    }

    /**
     * maxN backup: the active player picks the child that maximizes their own
     * component; the chosen child's entire heuristic vector is copied here.
     *
     * This generalizes minimax across N players and asymmetric turn orders: the
     * non-active components ride along, representing "what would happen to player Y
     * if X plays the way X wants from this state".
     */
    aggregateHeuristicScore() {
        const moves = this.moves;

        if (moves === null) {
            return;
        }

        const p = this.activePlayer;

        if (p < 0) {
            // terminal node — no decision to make, heuristicValue already reflects outcome
            return;
        }

        const own = this.heuristicValue;
        let bestVec = own;
        let bestVal = own[p];

        const n = moves.length;

        for (let i = 0; i < n; i++) {
            /**
             * @type {MoveEdge}
             */
            const move = moves[i];

            if (!move.isTargetMaterialized()) {
                continue;
            }

            const childVec = move.target.heuristicValue;
            const val = childVec[p];

            if (val > bestVal) {
                bestVal = val;
                bestVec = childVec;
            }
        }

        if (bestVec !== own) {
            own.set(bestVec);
        }
    }

    /**
     * @param state
     * @param {function(State, source:StateNode):MoveEdge[]} computeValidMoves
     * @param {function(State):number} computeActivePlayer
     * @param {function(State):Float64Array} computeOutcome called when expansion produces no moves
     * @returns {number} number of children
     */
    expand(state, computeValidMoves, computeActivePlayer, computeOutcome) {
        assert.defined(state, 'state');
        assert.isFunction(computeValidMoves, 'computeValidMoves');
        assert.isFunction(computeActivePlayer, 'computeActivePlayer');
        assert.isFunction(computeOutcome, 'computeOutcome');


        /**
         *
         * @type {MoveEdge[]}
         */
        const moves = computeValidMoves(state, this);

        assert.notNull(moves, 'moves');
        assert.defined(moves, 'moves');
        assert.isArray(moves, 'moves');

        const numMoves = moves.length;

        this.moves = moves;

        if (numMoves === 0) {
            // No legal moves → terminal by exhaustion. Cache outcome so subsequent
            // visits don't recompute, and seed heuristicValue with it.
            // Copy into our own buffer so the user is free to reuse the source;
            // outcome and heuristicValue share storage within the node (safe since
            // terminal nodes are skipped by aggregateHeuristicScore).
            this.type = StateType.NoMoves;
            this.activePlayer = -1;

            const out = computeOutcome(state);

            assert.defined(out, 'out');
            assert.isArrayLike(out, 'out');

            this.heuristicValue.set(out);
            this.outcome = this.heuristicValue;
        } else {

            this.activePlayer = computeActivePlayer(state);

            assert.isNumber(this.activePlayer, 'activePlayer');
        }

        return numMoves;
    }


    /**
     * Walk up from this node to the root, adding the per-player outcome to each
     * ancestor's cumulative score and incrementing their visit count.
     * @param {Float64Array} outcomeVector indexed by player ID
     */
    backpropagate(outcomeVector) {
        assert.defined(outcomeVector, 'outcomeVector');
        assert.isArrayLike(outcomeVector, 'outcomeVector');
        assert(outcomeVector.length > 0, 'outcomeVector must not be empty');

        const len = outcomeVector.length;
        let node = this;

        do {
            node.playouts += 1;

            const score = node.totalScore;
            for (let p = 0; p < len; p++) {
                score[p] += outcomeVector[p];
            }

            node = node.parent;

        } while (node !== null);
    }


    /**
     * Whenever this is a terminal state or not (win/loss/tie/no-moves/depth-cap)
     * @returns {boolean}
     */
    isTerminal() {
        return this.type !== StateType.Undecided;
    }

    /**
     *
     * @returns {boolean}
     */
    isExpanded() {
        return this.moves !== null;
    }

    /**
     * Pick the best moves from this node's perspective — i.e. for whichever
     * player owns the decision here.
     * @returns {MoveEdge[]}
     */
    pickBestMoves() {
        const moves = this.moves;

        if (moves === null) {
            return [];
        }

        const numMoves = moves.length;

        if (numMoves === 0) {
            return [];
        }

        // For a decision node, score by the active player's component.
        // For a terminal node (no active player), fall back to player 0 — caller
        // shouldn't normally ask for best moves on a terminal but we don't crash.
        const player = this.activePlayer >= 0 ? this.activePlayer : 0;

        const firstMove = moves[0];

        let result = [firstMove];
        let bestScore = pickBestScore(firstMove, player);

        for (let i = 1; i < numMoves; i++) {

            const move = moves[i];

            const score = pickBestScore(move, player);

            if (score > bestScore) {
                bestScore = score;
                result = [move];
            } else if (score === bestScore) {
                result.push(move);
            }
        }

        return result;
    }

    /**
     *
     * @param {function(StateNode)} visitor
     */
    traverse(visitor) {
        const stackOffset = stack_pointer;

        stack[stack_pointer++] = this;

        let n;

        while (stack_pointer-- > stackOffset) {

            n = stack[stack_pointer];

            visitor(n);

            if (n.isExpanded()) {

                const moves = n.moves;
                const numMoves = moves.length;

                for (let i = 0; i < numMoves; i++) {
                    const moveEdge = moves[i];

                    if (moveEdge.isTargetMaterialized()) {
                        stack[stack_pointer++] = moveEdge.target;
                    }
                }
            }
        }

    }
}
