//
// Abstract MCTS AI implementation
//

#include "AbstractMCTSAI.hpp"

#include <algorithm>
#include <chrono>
#include <fstream>
#include <future>
#include <iomanip>
#include <iostream>
#include <limits>
#include <mutex>
#include <random>
#include <sstream>
#include <stdexcept>
#include <thread>

#include "src/main/cpp/net/eagle0/common/mcts/util/TreeIndentUtil.hpp"

namespace shardok::mcts {

AbstractMCTSAI::AbstractMCTSAI(MCTSPlayerId playerId, MCTSConfig config)
    : playerId_(playerId),
      config_(config),
      randomEngine_(config.randomSeed == 0 ? std::random_device{}() : config.randomSeed) {}

void AbstractMCTSAI::SetConfig(const MCTSConfig& newConfig) {
    config_ = newConfig;
    std::lock_guard lock(randomEngineMutex_);
    randomEngine_.seed(newConfig.randomSeed == 0 ? std::random_device{}() : newConfig.randomSeed);
}

auto AbstractMCTSAI::Search(
        const MCTSGameEngine& engine,
        const MCTSGameState& initialState,
        std::chrono::milliseconds timeLimit) const -> SearchResult {
    const auto startTime = std::chrono::steady_clock::now();
    const auto deadline = startTime + timeLimit;

    // Build MCTS tree
    const auto rootNode = BuildMCTSTree(engine, initialState, deadline);

    SearchResult result;
    result.searchTime = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - startTime);

    if (!rootNode) {
        throw MCTSInternalError("MCTS search: BuildMCTSTree returned null root node");
    }

    if (rootNode->children.empty()) {
        // This can happen legitimately when:
        // 1. No legal actions available (terminal state) - return default
        // 2. Only one action and we early-exited without exploring - return index 0
        // 3. Multiple actions but none expanded - this is a bug
        if (rootNode->totalActions == 0) {
            // Terminal state - no actions available, return default result
            result.bestActionIndex = 0;
            result.bestScore = 0.0;
            result.nodesEvaluated = 0;
            return result;
        }
        // Single action case - should have been expanded in BuildMCTSTree
        if (rootNode->totalActions == 1) {
            result.bestActionIndex = 0;
            result.bestScore = 0.0;
            return result;
        }
        // Multiple actions but no children expanded - this shouldn't happen
        throw MCTSInternalError(
                "MCTS search: Root has " + std::to_string(rootNode->totalActions) +
                " actions but no children expanded - this indicates a bug in BuildMCTSTree");
    }

    // Find best child
    if (const auto* bestChild = rootNode->GetBestFinalChild();
        bestChild && bestChild->action && bestChild->actionIndex != SIZE_MAX) {
        // Use actionIndex which is the index into the filtered actions from
        // engine.getLegalActions()
        result.bestActionIndex = bestChild->actionIndex;
        result.bestScore = bestChild->lookaheadScore;  // Use minimax value, not poisoned average
        result.searchDepth = bestChild->depth;
        result.nodesEvaluated = rootNode->visitCount;

        // Check if we found a winning move
        if (bestChild->gameState && bestChild->gameState->isTerminal() &&
            bestChild->gameState->getWinner() == playerId_) {
            result.foundWinningMove = true;
        }

        // Log results
        LogSearchResults(rootNode.get(), bestChild, result);
    }

    // Dump tree if explicitly requested via config
    if (!config_.debugDumpPath.empty()) { DumpTreeToFile(rootNode.get(), config_.debugDumpPath); }

    return result;
}

auto AbstractMCTSAI::BuildMCTSTree(
        const MCTSGameEngine& engine,
        const MCTSGameState& initialState,
        const std::chrono::steady_clock::time_point deadline) const -> std::unique_ptr<MCTSNode> {
    // Clear transposition table for this search
    // Maps state hash -> minimum depth, used to detect redundant longer paths
    transpositionTable_.clear();

    // Create root node
    // IMPORTANT: Use the initial state's current player, not playerId_
    // node->playerId represents "whose turn it is", not "who we're searching for"
    // This is critical for correct player flip tracking
    auto root = std::make_unique<MCTSNode>(initialState.clone(), initialState.currentPlayerId(), 0);

    // Record root state in transposition table
    transpositionTable_[root->stateHash] = root->depth;

    // Set whether root is maximizing based on whether current player matches who we're searching
    // for
    root->isMaximizingPlayer = (initialState.currentPlayerId() == playerId_);

    // Get legal actions from engine for the root state
    // Root has 0 player flips
    const auto rootActions =
            engine.getLegalActions(initialState, playerId_, 0, config_.maxPlayerFlips);

    // Early exit if only one action available - no need to search
    if (rootActions.size() <= 1) {
        // Expand the single action so Search() can return it
        if (!rootActions.empty()) {
            root->totalActions = 1;
            [[maybe_unused]] auto* expanded = MCTSExpansion(root.get(), engine);
        }
        return root;
    }

    // Initialize action counter
    root->totalActions = rootActions.size();

    // CRITICAL: Do at least one expansion before entering the time-bounded loop.
    // This ensures we always have at least one child to return, even if the deadline
    // has already passed (e.g., due to debugger pause, system load, etc.)
    {
        auto* selected = MCTSSelection(root.get());
        const bool selectedIsRoot = (selected == root.get());
        const size_t childrenBeforeExpansion = root->children.size();

        if (selected) {
            auto* expanded = MCTSExpansion(selected, engine);
            const double reward =
                    MCTSSimulation(engine, *expanded->gameState, playerId_, expanded->playerFlips);
            MCTSBackpropagation(expanded, reward, config_.backpropagationPolicy);
        }

        // Verify we actually have at least one child after the initial expansion
        if (root->children.empty()) {
            throw MCTSInternalError(
                    "MCTS BuildMCTSTree: Initial expansion failed to produce any children. "
                    "totalActions=" +
                    std::to_string(root->totalActions) +
                    ", selected=" + (selected ? "non-null" : "null") +
                    ", selectedIsRoot=" + (selectedIsRoot ? "true" : "false") +
                    ", childrenBefore=" + std::to_string(childrenBeforeExpansion) +
                    ", childrenAfter=" + std::to_string(root->children.size()) +
                    ", root->CanExpand()=" + (root->CanExpand() ? "true" : "false") +
                    ", root->nextUntriedActionIndex=" +
                    std::to_string(root->nextUntriedActionIndex));
        }
    }

    std::atomic<int> iterations{0};

    if (config_.useMultithreading && config_.numThreads > 1) {
        // Multithreaded MCTS
        std::mutex treeMutex;
        std::vector<std::future<void>> futures;

        futures.reserve(config_.numThreads);
        for (int threadId = 0; threadId < config_.numThreads; ++threadId) {
            futures.push_back(std::async(std::launch::async, [&] {
                while (std::chrono::steady_clock::now() < deadline) {
                    // Selection and Expansion (with lock - tree modification must be serialized)
                    MCTSNode* expanded;
                    {
                        std::lock_guard lock(treeMutex);
                        auto* selected = MCTSSelection(root.get());
                        if (!selected) break;

                        // Expansion modifies tree structure - must be inside lock
                        expanded = MCTSExpansion(selected, engine);
                    }

                    // Simulation can run in parallel (doesn't modify tree)

                    // Backpropagation (with lock - modifies node statistics)
                    {
                        const double reward = MCTSSimulation(
                                engine,
                                *expanded->gameState,
                                playerId_,
                                expanded->playerFlips);
                        std::lock_guard lock(treeMutex);
                        MCTSBackpropagation(expanded, reward, config_.backpropagationPolicy);
                        iterations.fetch_add(1);
                    }
                }
            }));
        }

        // Wait for all threads to complete and propagate worker exceptions.
        for (auto& future : futures) { future.get(); }
    } else {
        // Single-threaded MCTS
        while (std::chrono::steady_clock::now() < deadline) {
            // Selection
            auto* selected = MCTSSelection(root.get());
            if (!selected) { break; }

            // Expansion
            auto* expanded = MCTSExpansion(selected, engine);

            // Simulation
            const double reward =
                    MCTSSimulation(engine, *expanded->gameState, playerId_, expanded->playerFlips);

            // Backpropagation
            MCTSBackpropagation(expanded, reward, config_.backpropagationPolicy);

            ++iterations;

            // Early termination check
            if (engine.shouldStopSearch(
                        *root->gameState,
                        iterations,
                        std::chrono::steady_clock::now())) {
                break;
            }
        }
    }

    return root;
}

auto AbstractMCTSAI::MCTSSelection(MCTSNode* root) const -> MCTSNode* {
    MCTSNode* current = root;

    while (current->depth < config_.maxTreeDepth) {
        // Check expansion FIRST - allows expanding "terminal" nodes that still have
        // untried actions (e.g., final round where we need to pick an action)
        if (current->CanExpand()) {
            return current;  // Node has untried actions/outcomes
        }

        // Only after expansion check: stop if terminal and fully expanded
        if (current->isTerminal) {
            break;  // Terminal and no more actions to try
        }

        if (!current->children.empty()) {
            // Choose child based on node type
            if (current->IsChanceNode()) {
                // Chance nodes: select outcome proportional to probability
                current = current->GetBestChanceChild();
            } else {
                // Decision nodes: select using UCB1
                current = current->GetBestChild(config_.explorationConstant);
            }
            if (!current) break;
        } else {
            break;  // Leaf node
        }
    }

    return current;
}

auto AbstractMCTSAI::MCTSExpansion(MCTSNode* node, const MCTSGameEngine& engine) const
        -> MCTSNode* {
    // Only skip if we truly can't expand. Allow expansion even if "terminal" as long as
    // there are untried actions (e.g., final round where we need to pick an action).
    if (!node->CanExpand()) {
        return node;  // Nothing to expand
    }

    // Handle chance node expansion (expanding outcomes)
    if (node->IsChanceNode()) {
        // Chance nodes expand their outcome children
        // This should have been set up when the chance node was created
        if (node->outcomeProbabilities.empty()) {
            throw MCTSInternalError(
                    "Chance node has no outcome probabilities - this indicates a bug");
        }

        const size_t outcomeIndex = node->nextUntriedActionIndex++;
        if (outcomeIndex >= node->outcomeProbabilities.size()) {
            throw MCTSInternalError(
                    "Chance node outcomeIndex >= outcomeProbabilities.size() - bug in expansion");
        }

        // The chance node's action should be the binary action
        if (!node->action) {
            throw MCTSInternalError("Chance node has no action - this indicates a bug");
        }

        // Apply the action with the representative roll for this outcome
        // Outcome 0 = success, Outcome 1 = failure
        // Use the representative roll for this specific outcome
        const double representativeRoll = node->outcomeRolls[outcomeIndex];
        auto newState = engine.applyAction(*node->gameState, *node->action, representativeRoll);
        if (!newState) {
            throw MCTSInternalError(
                    "MCTS expansion: engine.applyAction() returned nullptr for chance node "
                    "outcome - this indicates a game engine error");
        }

        // Determine if player changed
        const MCTSPlayerId newPlayerId = newState->currentPlayerId();
        const bool playerChanged = (newPlayerId != node->playerId);

        // Calculate player flips and maximizing status
        const int newPlayerFlips = node->playerFlips + (playerChanged ? 1 : 0);
        const bool newIsMaximizing = (newPlayerId == playerId_);

        // Create outcome child (decision node)
        auto outcomeChild = std::make_unique<MCTSNode>(
                node->action->clone(),
                std::move(newState),
                newPlayerId,
                node->depth + 1,
                outcomeIndex,
                newPlayerFlips,
                newIsMaximizing,
                node->actionWeight);  // Inherit action weight from chance node

        // Set up outcome child's actions if not terminal
        const bool shouldExpand =
                !outcomeChild->isTerminal && node->playerFlips <= config_.maxPlayerFlips;
        if (shouldExpand) {
            const auto childActions = engine.getLegalActions(
                    *outcomeChild->gameState,
                    playerId_,
                    newPlayerFlips,
                    config_.maxPlayerFlips);
            outcomeChild->totalActions = childActions.size();
        }

        // Calculate scores
        outcomeChild->immediateScore = engine.evaluateState(*outcomeChild->gameState, playerId_);
        outcomeChild->lookaheadScore = outcomeChild->immediateScore;

        // Set parent and add to children
        outcomeChild->parent = node;
        node->children.push_back(std::move(outcomeChild));

        // Update chance node's immediate score to expected value of expanded outcomes
        // This corrects the initial value (which incorrectly used parent state) and ensures
        // fair UCB comparison with non-chance actions like END_TURN
        {
            double expectedImmediate = 0.0;
            for (size_t i = 0; i < node->outcomeProbabilities.size(); i++) {
                const double prob = node->outcomeProbabilities[i];
                const double childImmediate = i < node->children.size()
                                                      ? node->children[i]->immediateScore
                                                      : node->unvisitedOutcomeScore;
                expectedImmediate += prob * childImmediate;
            }
            node->immediateScore = expectedImmediate;
            // CRITICAL: Always update lookaheadScore to the expected value.
            // Without this, chance nodes keep their initial lookaheadScore from the parent
            // state (before the action), while regular actions use the child state (after).
            // This gives chance nodes an unfair initial UCB advantage.
            node->lookaheadScore = node->immediateScore;
        }

        return node->children.back().get();
    }

    // Handle decision node expansion (expanding actions)
    // Get next action to expand (sequential order)
    const size_t actionIndex = node->nextUntriedActionIndex++;

    // Get legal actions from engine (uses cached engine for performance)
    // Use the parameterized version to respect player flips
    const auto nodeActions = engine.getLegalActions(
            *node->gameState,
            playerId_,
            node->playerFlips,
            config_.maxPlayerFlips);

    // Create new child node
    if (actionIndex >= nodeActions.size()) {
        throw MCTSInternalError(
                "MCTS expansion: actionIndex (" + std::to_string(actionIndex) +
                ") >= nodeActions.size() (" + std::to_string(nodeActions.size()) +
                ") - this indicates a bug in action indexing");
    }

    // Get action weights from engine (for prior-weighted UCB)
    const auto actionWeights = engine.getActionWeights(nodeActions, *node->gameState);

    const auto& action = nodeActions[actionIndex];

    double totalActionWeight = 0.0;
    for (const double weight : actionWeights) { totalActionWeight += std::max(0.0, weight); }

    const double rawActionWeight =
            actionIndex < actionWeights.size() ? std::max(0.0, actionWeights[actionIndex]) : 1.0;
    const double actionWeight = totalActionWeight > 0.0
                                        ? rawActionWeight / totalActionWeight
                                        : 1.0 / static_cast<double>(nodeActions.size());

    // Check if this action requires a chance node
    if (action->requiresChanceNode()) {
        // Create intermediate chance node
        auto chanceNode = std::make_unique<MCTSNode>(
                action->clone(),
                node->gameState->clone(),  // Chance node has same state as parent
                node->playerId,
                node->depth + 1,
                actionIndex,
                node->playerFlips,
                node->isMaximizingPlayer,
                actionWeight);

        chanceNode->nodeType = NodeType::CHANCE;

        // Get outcome information from engine
        const auto outcomeInfo = engine.getBinaryOutcomeInfo(*node->gameState, *action);

        // Set up outcome metadata.
        chanceNode->outcomeProbabilities = outcomeInfo.getProbabilities();
        chanceNode->outcomeRolls = outcomeInfo.getRepresentativeRolls();
        chanceNode->totalActions = chanceNode->outcomeProbabilities.size();

        // Chance node immediate score will be computed as expected value during backpropagation
        // For now, initialize to parent's score as a reasonable default
        chanceNode->immediateScore = engine.evaluateState(*node->gameState, playerId_);
        chanceNode->lookaheadScore = chanceNode->immediateScore;
        chanceNode->unvisitedOutcomeScore = chanceNode->immediateScore;

        // Set parent and add to children
        chanceNode->parent = node;
        node->children.push_back(std::move(chanceNode));

        // CRITICAL: Immediately expand the first outcome and return that instead.
        // If we returned the chance node itself, MCTSSimulation would run on the parent state
        // (since chance nodes have parent's gameState), which is wrong. We need to simulate
        // from an actual outcome state.
        //
        // Note: This recursion is bounded because outcome children are decision nodes,
        // not chance nodes, so the recursion goes exactly one level deep.
        return MCTSExpansion(node->children.back().get(), engine);
    }

    // Regular (non-chance) action: create decision node directly
    auto newState = engine.applyAction(*node->gameState, *action);
    if (!newState) {
        throw MCTSInternalError(
                "MCTS expansion: engine.applyAction() returned nullptr for action " +
                action->getDescription() + " - this indicates a game engine error");
    }

    // Determine if player changed
    const MCTSPlayerId newPlayerId = newState->currentPlayerId();
    const bool playerChanged = (newPlayerId != node->playerId);

    // Calculate player flips and maximizing status
    const int newPlayerFlips = node->playerFlips + (playerChanged ? 1 : 0);
    // Node is maximizing if current player is the root player (playerId_)
    const bool newIsMaximizing = (newPlayerId == playerId_);

    auto child = std::make_unique<MCTSNode>(
            action->clone(),
            std::move(newState),
            newPlayerId,
            node->depth + 1,
            actionIndex,
            newPlayerFlips,
            newIsMaximizing,
            actionWeight);  // Pass the action weight for prior-weighted UCB

    // Check transposition table: mark as redundant if we've reached this state at a shallower depth
    // This prevents MCTS from exploring longer paths to the same game state
    // Works best with MINIMAX backpropagation (penalty propagates as min/max)
    // Also provides benefit with AVERAGING (penalty pulls average down significantly)
    const uint64_t childHash = child->stateHash;
    auto it = transpositionTable_.find(childHash);
    if (it != transpositionTable_.end()) {
        const int previousDepth = it->second;
        if (child->depth > previousDepth) {
            // Longer path to same state - mark as redundant and heavily penalize
            // Use -infinity to be unambiguously worse than any legitimate score
            child->isRedundant = true;
            child->immediateScore = -std::numeric_limits<double>::infinity();
            child->lookaheadScore = -std::numeric_limits<double>::infinity();
        } else {
            // Found shorter or equal path - update table
            transpositionTable_[childHash] = child->depth;
        }
    } else {
        // First time seeing this state - record it
        transpositionTable_[childHash] = child->depth;
    }

    // Set up child's untried actions if not terminal and parent hasn't exceeded player flips
    // playerFlips counts how many times the player has CHANGED from root
    // We expand children of nodes that are within the maxPlayerFlips limit
    // maxPlayerFlips=0: same player can take multiple sequential actions
    // maxPlayerFlips=1: can explore opponent's immediate responses
    const bool shouldExpand = !child->isTerminal && node->playerFlips <= config_.maxPlayerFlips;

    if (shouldExpand) {
        const auto childActions = engine.getLegalActions(
                *child->gameState,
                playerId_,
                newPlayerFlips,
                config_.maxPlayerFlips);
        child->totalActions = childActions.size();
    }

    // Calculate immediate and lookahead scores from root player's perspective
    // Skip for redundant nodes (already have penalty scores)
    if (!child->isRedundant) {
        child->immediateScore = engine.evaluateState(*child->gameState, playerId_);
        child->lookaheadScore = child->immediateScore;
    }

    // Set parent and add to children
    child->parent = node;
    node->children.push_back(std::move(child));

    return node->children.back().get();
}

auto AbstractMCTSAI::MCTSSimulation(
        const MCTSGameEngine& engine,
        const MCTSGameState& state,
        const MCTSPlayerId startingPlayer,
        const int startingPlayerFlips) const -> double {
    if (state.isTerminal()) { return state.score(startingPlayer); }

    // If we've already exceeded the simulation horizon, don't simulate - just return immediate
    // score This ensures fair comparison: all leaves are evaluated at the same game phase Example:
    // maxSimulationFlips=1 means simulate THROUGH opponent's first response (i.e., allow one action
    // at playerFlips=1, then stop)
    if (startingPlayerFlips > config_.maxSimulationFlips) { return state.score(startingPlayer); }

    // Create a mutable copy for simulation
    auto currentState = state.clone();
    int depth = 0;
    int playerFlips = startingPlayerFlips;  // Start from the expanded node's flip count
    MCTSPlayerId previousPlayer = currentState->currentPlayerId();

    // Simulate until we exceed the horizon, hit terminal state, or max depth
    // Note: We allow one action AT maxSimulationFlips before stopping
    while (!currentState->isTerminal() && depth < config_.maxSimulationDepth &&
           playerFlips <= config_.maxSimulationFlips) {
        // Track player changes
        const MCTSPlayerId currentPlayer = currentState->currentPlayerId();
        if (currentPlayer != previousPlayer) {
            playerFlips++;
            previousPlayer = currentPlayer;
        }

        // Get legal actions with player flip tracking
        const auto actions = engine.getLegalActions(
                *currentState,
                playerId_,
                playerFlips,
                config_.maxSimulationFlips);
        if (actions.empty()) { break; }

        // Determine if current player is maximizing or minimizing
        // Maximizing: current player is root player (trying to maximize root player's score)
        // Minimizing: current player is opponent (trying to minimize root player's score)
        const bool isMaximizing = (currentPlayer == playerId_);

        // Select action based on simulation policy
        const size_t selectedIndex =
                SelectSimulationAction(engine, *currentState, actions, isMaximizing);
        if (selectedIndex >= actions.size()) { break; }

        // Apply action
        auto newState = engine.applyAction(*currentState, *actions[selectedIndex]);
        if (!newState) { break; }

        currentState = std::move(newState);
        depth++;
    }

    return currentState->score(startingPlayer);
}

auto AbstractMCTSAI::MCTSBackpropagation(
        MCTSNode* node,
        const double reward,
        const MCTSBackpropagationPolicy policy) const -> void {
    // Backpropagation strategy is configured via MCTSConfig:
    // - AVERAGING: Traditional MCTS averaging (for stochastic/single-player games)
    // - MINIMAX: Minimax backup (for deterministic adversarial games)

    const bool useMinimaxBackup = (policy == MCTSBackpropagationPolicy::MINIMAX);

    while (node) {
        node->visitCount++;

        // Always track average for UCB
        node->totalReward += reward;
        node->averageReward = node->totalReward / node->visitCount;

        // Update lookahead score based on node type and strategy
        if (node->IsChanceNode() && !node->children.empty()) {
            // Chance nodes: compute expected value (weighted average of outcomes)
            // lookaheadScore = sum(probability[i] * childValue[i])
            double expectedValue = 0.0;
            bool hasVisitedChild = false;

            for (size_t i = 0; i < node->outcomeProbabilities.size(); i++) {
                const double probability = node->outcomeProbabilities[i];
                double childValue = node->unvisitedOutcomeScore;
                if (i < node->children.size() && node->children[i]->visitCount > 0) {
                    childValue = node->children[i]->lookaheadScore;
                    hasVisitedChild = true;
                }
                expectedValue += probability * childValue;
            }

            // Use expected value if we have visited outcomes, else use average
            if (hasVisitedChild) {
                node->lookaheadScore = expectedValue;
            } else {
                // No outcomes visited yet, fall back to average
                if (node->visitCount == 1) {
                    node->lookaheadScore = reward;
                } else {
                    const double alpha = 1.0 / node->visitCount;
                    node->lookaheadScore = (1.0 - alpha) * node->lookaheadScore + alpha * reward;
                }
            }
        } else if (useMinimaxBackup && !node->children.empty()) {
            // Minimax backup: use best/worst child value for adversarial games
            // The operation (MAX or MIN) depends on whose turn it is at THIS node
            // - If this node is root player's turn: root chooses MAX (best for root)
            // - If this node is opponent's turn: opponent chooses MIN (best for opponent = worst
            // for root)
            //
            // Note: In setup phase, children can have different isMaximizingPlayer values:
            // - PLACE_UNIT keeps same player's turn
            // - END_PLAYER_SETUP flips to opponent's turn
            // So we must use the PARENT node's isMaximizingPlayer, not the child's.

            const bool thisNodeIsRootPlayer = node->isMaximizingPlayer;

            double minmaxValue = thisNodeIsRootPlayer ? -std::numeric_limits<double>::max()
                                                      : std::numeric_limits<double>::max();

            int visitedChildCount = 0;
            for (const auto& child : node->children) {
                if (child->visitCount == 0) continue;  // Unvisited children don't contribute

                const double childValue = child->lookaheadScore;
                visitedChildCount++;

                if (thisNodeIsRootPlayer) {
                    // Root player chooses: take MAX (best for root)
                    minmaxValue = std::max(minmaxValue, childValue);
                } else {
                    // Opponent chooses: take MIN (best for opponent = worst for root)
                    minmaxValue = std::min(minmaxValue, childValue);
                }
            }

            // Use minimax value if we found any visited children, else use average
            if (visitedChildCount > 0) {
                node->lookaheadScore = minmaxValue;
            } else {
                // No children visited yet, fall back to average
                if (node->visitCount == 1) {
                    node->lookaheadScore = reward;
                } else {
                    const double alpha = 1.0 / node->visitCount;
                    node->lookaheadScore = (1.0 - alpha) * node->lookaheadScore + alpha * reward;
                }
            }
        } else {
            // Standard MCTS averaging (for maxPlayerFlips=0 or leaf nodes)
            if (node->visitCount == 1) {
                node->lookaheadScore = reward;
            } else {
                const double alpha = 1.0 / node->visitCount;
                node->lookaheadScore = (1.0 - alpha) * node->lookaheadScore + alpha * reward;
            }
        }

        node = node->parent;
    }
}

auto AbstractMCTSAI::SelectSimulationAction(
        const MCTSGameEngine& engine,
        const MCTSGameState& state,
        const std::vector<std::unique_ptr<MCTSAction>>& actions,
        const bool isMaximizing) const -> size_t {
    if (actions.empty()) {
        throw MCTSInternalError(
                "MCTSSimulation called with empty actions list - this indicates a bug in the "
                "MCTS tree building or game state");
    }

    switch (config_.simulationPolicy) {
        case MCTSSimulationPolicy::RANDOM: {
            std::uniform_int_distribution<size_t> dis(0, actions.size() - 1);
            return SampleDistribution(dis);
        }

        case MCTSSimulationPolicy::FILTERED_RANDOM: {
            if (const auto filteredIndices = engine.filterActions(actions, state);
                !filteredIndices.empty()) {
                std::uniform_int_distribution<size_t> dis(0, filteredIndices.size() - 1);
                return filteredIndices[SampleDistribution(dis)];
            }
            // Fall back to random
            std::uniform_int_distribution<size_t> dis(0, actions.size() - 1);
            return SampleDistribution(dis);
        }

        case MCTSSimulationPolicy::BEST_IMMEDIATE: {
            // For adversarial search:
            // - Maximizing nodes select action with HIGHEST score (best for root player)
            // - Minimizing nodes select action with LOWEST score (worst for root player)

            // First, filter out obviously bad moves (e.g., BECOME_OUTLAW)
            const auto filteredIndices = engine.filterActions(actions, state);
            if (filteredIndices.empty()) {
                // If all actions filtered out, fall back to first action
                return 0;
            }

            double bestScore = isMaximizing ? -std::numeric_limits<double>::max()
                                            : std::numeric_limits<double>::max();
            size_t bestIndex = filteredIndices[0];

            for (const size_t i : filteredIndices) {
                // CRITICAL: Always get score from ROOT player's perspective for adversarial search
                // If we use currentPlayerId, opponent actions would be scored from their
                // perspective, causing them to select moves that help themselves instead of hurt
                // us!
                const double score = engine.getActionScore(state, *actions[i], playerId_);

                const bool shouldSelect = isMaximizing ? (score > bestScore) : (score < bestScore);

                if (shouldSelect) {
                    bestScore = score;
                    bestIndex = i;
                }
            }
            return bestIndex;
        }

        case MCTSSimulationPolicy::WEIGHTED_BEST_IMMEDIATE: {
            // Score all actions and weight by ranking
            std::vector<std::pair<size_t, double>> scores;
            scores.reserve(actions.size());

            for (size_t i = 0; i < actions.size(); ++i) {
                // CRITICAL: Always get score from ROOT player's perspective for adversarial search
                const double score = engine.getActionScore(state, *actions[i], playerId_);
                scores.emplace_back(i, score);
            }

            // Sort by score
            // - Maximizing: highest scores first (prefer actions that maximize root player's score)
            // - Minimizing: lowest scores first (prefer actions that minimize root player's score)
            if (isMaximizing) {
                std::ranges::sort(scores, [](const auto& a, const auto& b) {
                    return a.second > b.second;  // Descending
                });
            } else {
                std::ranges::sort(scores, [](const auto& a, const auto& b) {
                    return a.second < b.second;  // Ascending
                });
            }

            // Create weights based on ranking
            std::vector<double> weights;
            weights.reserve(scores.size());
            for (size_t i = 0; i < scores.size(); ++i) {
                weights.push_back(1.0 / (static_cast<double>(i) + 1.0));
            }

            // Select based on weights
            std::discrete_distribution<> dis(weights.begin(), weights.end());
            return scores[SampleDistribution(dis)].first;
        }

        case MCTSSimulationPolicy::WEIGHTED_HEURISTIC: {
            // Get heuristic weights from engine (fast O(1) per action)
            const auto weights = engine.getActionWeights(actions, state);

            // Filter out zero-weight actions
            std::vector<size_t> validIndices;
            std::vector<double> validWeights;
            validIndices.reserve(actions.size());
            validWeights.reserve(actions.size());

            for (size_t i = 0; i < weights.size() && i < actions.size(); ++i) {
                if (weights[i] > 0.0) {
                    validIndices.push_back(i);
                    validWeights.push_back(weights[i]);
                }
            }

            // If all actions filtered out, this is a bug in the weighting logic
            if (validWeights.empty()) {
                throw MCTSInternalError(
                        "MCTS simulation: All actions have zero weight in WEIGHTED_HEURISTIC "
                        "policy (action count: " +
                        std::to_string(actions.size()) + ") - this indicates incorrect weighting");
            }

            // Select based on heuristic weights using discrete_distribution
            std::discrete_distribution<> dis(validWeights.begin(), validWeights.end());
            return validIndices[SampleDistribution(dis)];
        }
    }

    // Default to random
    std::uniform_int_distribution<size_t> dis(0, actions.size() - 1);
    return SampleDistribution(dis);
}

auto AbstractMCTSAI::FindNodeAtDepthWithHash(
        const MCTSNode* root,
        const int maxDepth,
        const uint64_t targetHash) -> const MCTSNode* {
    if (!root || root->depth >= maxDepth || root->stateHash == targetHash) { return root; }

    // Breadth-first search for matching hash at specific depth
    std::vector<const MCTSNode*> currentLevel = {root};
    for (int d = 0; d < maxDepth && !currentLevel.empty(); ++d) {
        std::vector<const MCTSNode*> nextLevel;

        for (const auto* node : currentLevel) {
            for (const auto& child : node->children) {
                if (child->stateHash == targetHash && child->depth <= maxDepth) {
                    return child.get();
                }
                if (child->depth < maxDepth) { nextLevel.push_back(child.get()); }
            }
        }

        currentLevel = std::move(nextLevel);
    }

    return nullptr;
}

auto AbstractMCTSAI::LogSearchResults(
        const MCTSNode* rootNode,
        const MCTSNode* bestChild,
        const SearchResult& result) -> void {
    // Log selected command
    const std::string selectedDesc =
            bestChild->action ? bestChild->action->getDescription() : "Unknown";
    std::ostringstream selectedLine;
    selectedLine << std::fixed << std::setprecision(2) << "MCTS: Selected " << selectedDesc
                 << " (visit:" << bestChild->visitCount << ", reward:" << bestChild->averageReward
                 << ", lookahead:" << bestChild->lookaheadScore << ") from "
                 << rootNode->children.size() << " options";
    std::cout << selectedLine.str() << '\n';

    // Log top 3 actions by visits
    std::vector<const MCTSNode*> sortedChildren;
    sortedChildren.reserve(rootNode->children.size());
    for (const auto& child : rootNode->children) { sortedChildren.push_back(child.get()); }
    std::ranges::sort(sortedChildren, [](const auto* a, const auto* b) {
        return a->visitCount > b->visitCount;
    });

    // Show all actions if there are <= 10, otherwise top 5
    const size_t numToShow = sortedChildren.size() <= 10 ? sortedChildren.size() : 5;
    std::cout << "MCTS: Top " << numToShow << " actions by visits (out of " << sortedChildren.size()
              << " total):\n";
    for (size_t i = 0; i < numToShow; ++i) {
        const auto* child = sortedChildren[i];
        std::ostringstream line;
        line << std::fixed << std::setprecision(2) << "  [" << i << "] visits:" << child->visitCount
             << " avgReward:" << child->averageReward << " immediate:" << child->immediateScore
             << " lookahead:" << child->lookaheadScore;

        // Show the action's own description
        if (child->action) { line << ' ' << child->action->getDescription(); }

        // Show sequence preview
        if (!child->children.empty()) {
            const MCTSNode* bestNext = nullptr;
            int maxVisits = 0;
            for (const auto& grandchild : child->children) {
                if (grandchild->visitCount > maxVisits) {
                    maxVisits = grandchild->visitCount;
                    bestNext = grandchild.get();
                }
            }
            if (bestNext && bestNext->action) {
                line << " -> " << bestNext->action->getDescription();
            }
        }
        std::cout << line.str() << '\n';
    }

    std::cout << "MCTS: Tree stats - max depth:" << result.searchDepth
              << ", total nodes:" << result.nodesEvaluated
              << ", root visits:" << rootNode->visitCount << '\n';

    // Log best sequence from chosen action (following most-visited path)
    std::vector<const MCTSNode*> bestSequence;
    bestSequence.reserve(10);
    const MCTSNode* current = bestChild;
    double sequenceScore = bestChild->averageReward;

    while (current && bestSequence.size() < 10) {
        bestSequence.push_back(current);
        if (current->children.empty()) break;

        // Find most visited child (standard MCTS principal variation)
        const MCTSNode* bestChildNode = nullptr;
        int maxVisits = 0;
        for (const auto& child : current->children) {
            if (child->visitCount > maxVisits) {
                maxVisits = child->visitCount;
                bestChildNode = child.get();
            }
        }
        current = bestChildNode;
        if (current) { sequenceScore = current->averageReward; }
    }

    if (!bestSequence.empty()) {
        std::ostringstream sequenceLine;
        sequenceLine << std::fixed << std::setprecision(2)
                     << "MCTS: Best sequence from chosen action (final: " << sequenceScore << "):";
        std::cout << sequenceLine.str() << '\n';
        int displayedStep = 0;
        for (size_t i = 0; i < bestSequence.size(); ++i) {
            const auto* node = bestSequence[i];

            // Skip outcome nodes (children of chance nodes) - they're displayed with their parent
            if (i > 0 && node->parent && node->parent->IsChanceNode()) { continue; }

            displayedStep++;
            std::ostringstream line;
            line << "  " << displayedStep << ".";
            if (node->action) { line << ' ' << node->action->getDescription(); }

            // If this is a chance node, display outcome probabilities and scores
            if (node->IsChanceNode() && !node->outcomeProbabilities.empty()) {
                line << " [";
                for (size_t j = 0; j < node->outcomeProbabilities.size(); ++j) {
                    if (j > 0) { line << ", "; }
                    const double prob = node->outcomeProbabilities[j] * 100;
                    // Show lookahead score for each outcome if child exists
                    if (j < node->children.size() && node->children[j]->visitCount > 0) {
                        line << std::fixed << std::setprecision(0) << prob << "%->"
                             << std::setprecision(1) << node->children[j]->lookaheadScore;
                    } else {
                        line << std::fixed << std::setprecision(0) << prob << "%->?";
                    }
                }
                line << "]";
            }

            line << std::fixed << std::setprecision(2) << " (visits:" << node->visitCount
                 << ", immediate:" << node->immediateScore << ", lookahead:" << node->lookaheadScore
                 << ")";
            std::cout << line.str() << '\n';

            // For non-root nodes in the sequence, show what the top alternatives were
            // Skip showing alternatives for chance nodes (they have outcome children, not action
            // alternatives)
            if (i > 0 && node->parent && !node->parent->children.empty() &&
                !node->parent->IsChanceNode()) {
                // Collect all siblings (including this node) and sort by visit count
                std::vector<const MCTSNode*> siblings;
                siblings.reserve(node->parent->children.size());
                for (const auto& child : node->parent->children) {
                    if (!child->isRedundant) { siblings.push_back(child.get()); }
                }

                // Sort by visit count (descending)
                std::ranges::sort(siblings, [](const MCTSNode* a, const MCTSNode* b) {
                    return a->visitCount > b->visitCount;
                });

                // Show top 3 alternatives at this decision point
                std::cout << "      Alternatives at this node (" << siblings.size() << " total):\n";
                const size_t topN = std::min(siblings.size(), size_t(3));
                for (size_t j = 0; j < topN; ++j) {
                    const auto* alt = siblings[j];
                    std::ostringstream altLine;
                    altLine << std::fixed << std::setprecision(2) << "        [" << j
                            << "] visits:" << alt->visitCount
                            << " immediate:" << alt->immediateScore
                            << " lookahead:" << alt->lookaheadScore;
                    if (alt->action) { altLine << ' ' << alt->action->getDescription(); }
                    std::cout << altLine.str() << '\n';
                }
            }
        }
    }
}

auto AbstractMCTSAI::DumpTreeToFile(const MCTSNode* root, const std::string& filepath) -> void {
    if (!root) return;

    std::ofstream out(filepath);
    if (!out) {
        std::cerr << "Failed to open dump file: " << filepath << '\n';
        return;
    }

    out << "MCTS Tree Dump\n";
    out << "==============\n\n";
    out << "Root Node:\n";
    out << "  Visits: " << root->visitCount << "\n";
    out << "  Immediate Score: " << root->immediateScore << "\n";
    out << "  Lookahead Score: " << root->lookaheadScore << "\n";
    out << "  Average Reward: " << root->averageReward << "\n";
    out << "  Player ID: " << root->playerId << "\n";
    out << "  Depth: " << root->depth << "\n";
    out << "  Is Maximizing: " << (root->isMaximizingPlayer ? "true" : "false") << "\n";
    out << "  State Hash: " << std::hex << root->stateHash << std::dec << "\n";
    out << "\n";

    if (!root->children.empty()) {
        out << "Children:\n";
        for (size_t i = 0; i < root->children.size(); ++i) {
            const auto& child = root->children[i];
            const bool isLast = (i == root->children.size() - 1);
            DumpNodeRecursive(child.get(), out, 1, isLast);
        }
    }

    out << "\n=== End of Tree Dump ===\n";
    out.close();

    std::cout << "MCTS: Tree dumped to " << filepath << '\n';
}

auto AbstractMCTSAI::DumpNodeRecursive(
        const MCTSNode* node,
        std::ostream& out,
        const int indentLevel,
        const bool isLastChild) -> void {
    if (!node) return;

    // Create indent string
    const std::string indent = ::mcts::util::BuildTreeIndent(indentLevel, isLastChild);

    // Write node information
    out << indent;

    // Show node type for chance nodes
    if (node->IsChanceNode()) { out << "[CHANCE] "; }

    if (node->action) {
        out << node->action->getDescription();
    } else {
        out << "[ROOT]";
    }
    out << " (visits:" << node->visitCount;
    out << ", immediate:" << std::fixed << std::setprecision(2) << node->immediateScore;
    out << ", lookahead:" << node->lookaheadScore;
    out << ", avgReward:" << node->averageReward;
    out << ", weight:" << node->actionWeight;
    out << ", depth:" << node->depth;
    out << ", flips:" << node->playerFlips;
    out << ", player:" << node->playerId;
    out << ", max:" << (node->isMaximizingPlayer ? "T" : "F");
    if (node->isRedundant) { out << ", REDUNDANT"; }
    if (node->isTerminal) { out << ", TERMINAL"; }
    out << ")\n";

    // Show outcome probabilities and rolls for chance nodes
    if (node->IsChanceNode() && !node->outcomeProbabilities.empty()) {
        const std::string outcomeIndent = ::mcts::util::ConvertBranchToContinuation(indent);
        out << outcomeIndent << "  Outcomes: ";
        for (size_t i = 0; i < node->outcomeProbabilities.size(); ++i) {
            if (i > 0) out << ", ";
            out << "[" << i << "] p=" << std::fixed << std::setprecision(3)
                << node->outcomeProbabilities[i];
            if (i < node->outcomeRolls.size()) {
                out << " roll=" << std::fixed << std::setprecision(1) << node->outcomeRolls[i];
            }
        }
        out << "\n";
    }

    // Recursively dump children
    if (!node->children.empty()) {
        for (size_t i = 0; i < node->children.size(); ++i) {
            const auto& child = node->children[i];
            const bool isLast = (i == node->children.size() - 1);
            DumpNodeRecursive(child.get(), out, indentLevel + 1, isLast);
        }
    }
}

}  // namespace shardok::mcts
