//
// Abstract MCTS Node structure for game-agnostic implementation
//

#ifndef EAGLE0_ABSTRACT_MCTSNODE_HPP
#define EAGLE0_ABSTRACT_MCTSNODE_HPP

#include <cmath>
#include <limits>
#include <memory>
#include <vector>

#include "MCTSAction.hpp"
#include "MCTSGameState.hpp"
#include "MCTSTypes.hpp"

namespace shardok {
namespace mcts {

// Abstract MCTS Node structure
struct MCTSNode {
    // Action information
    std::unique_ptr<MCTSAction> action;  // The action that led to this node (null for root)
    size_t actionIndex = SIZE_MAX;       // Index in the original actions array (SIZE_MAX for root)

    // Score information
    double immediateScore = 0.0;
    double lookaheadScore = 0.0;

    // Game state after this action
    std::unique_ptr<MCTSGameState> gameState;

    // MCTS statistics
    int visitCount = 0;
    double totalReward = 0.0;
    double averageReward = 0.0;
    mutable double ucb1Value = 0.0;

    // Tree structure
    std::vector<std::unique_ptr<MCTSNode>> children;
    std::vector<size_t> untriedActionIndices;
    bool fullyExpanded = false;
    MCTSNode* parent = nullptr;

    // Game context
    MCTSPlayerId playerId;
    int depth = 0;
    bool isTerminal = false;

    // Transposition detection
    uint64_t stateHash = 0;
    bool isRedundant = false;  // True if this node represents a duplicate state

    // Constructor for root node
    MCTSNode(std::unique_ptr<MCTSGameState> state, MCTSPlayerId pid, int d)
        : gameState(std::move(state)),
          playerId(pid),
          depth(d) {
        if (gameState) {
            stateHash = gameState->hash();
            isTerminal = gameState->isTerminal();
        }
    }

    // Constructor for child node
    MCTSNode(
            std::unique_ptr<MCTSAction> act,
            std::unique_ptr<MCTSGameState> state,
            MCTSPlayerId pid,
            int d,
            size_t actIdx = SIZE_MAX)
        : action(std::move(act)),
          actionIndex(actIdx),
          gameState(std::move(state)),
          playerId(pid),
          depth(d) {
        if (gameState) {
            stateHash = gameState->hash();
            isTerminal = gameState->isTerminal();
        }
    }

    // Iterative destructor to avoid stack overflow with deep trees
    ~MCTSNode() {
        std::vector<std::unique_ptr<MCTSNode>> nodesToDestroy;
        nodesToDestroy.swap(children);

        while (!nodesToDestroy.empty()) {
            std::vector<std::unique_ptr<MCTSNode>> currentBatch;
            currentBatch.swap(nodesToDestroy);

            for (const auto& node : currentBatch) {
                if (node && !node->children.empty()) {
                    for (auto& child : node->children) {
                        nodesToDestroy.push_back(std::move(child));
                    }
                    node->children.clear();
                }
            }
        }
    }

    // Calculate UCB1 value for this node
    void CalculateUCB1(const double explorationConstant) const {
        if (visitCount == 0) {
            ucb1Value = std::numeric_limits<double>::max();
        } else if (parent && parent->visitCount > 0) {
            ucb1Value = averageReward +
                        explorationConstant * std::sqrt(std::log(parent->visitCount) / visitCount);
        } else {
            ucb1Value = averageReward;
        }
    }

    // Check if this node can be expanded
    [[nodiscard]] bool CanExpand() const { return !fullyExpanded && !untriedActionIndices.empty(); }

    // Get best child based on UCB1
    [[nodiscard]] MCTSNode* GetBestChild(const double explorationConstant) const {
        if (children.empty()) return nullptr;

        MCTSNode* bestChild = nullptr;
        double bestValue = -std::numeric_limits<double>::max();

        for (auto& child : children) {
            // Skip redundant nodes
            if (child->isRedundant) continue;

            child->CalculateUCB1(explorationConstant);

            if (child->ucb1Value > bestValue) {
                bestValue = child->ucb1Value;
                bestChild = child.get();
            }
        }

        return bestChild;
    }

    // Get best child based on visit count (for final selection)
    [[nodiscard]] MCTSNode* GetBestFinalChild() const {
        if (children.empty()) return nullptr;

        MCTSNode* bestChild = nullptr;
        int bestVisits = 0;
        double bestScore = -std::numeric_limits<double>::max();

        for (const auto& child : children) {
            // Skip redundant nodes
            if (child->isRedundant) continue;

            // Prefer most-visited node (robust child selection)
            if (child->visitCount > bestVisits) {
                bestVisits = child->visitCount;
                bestScore = child->averageReward;
                bestChild = child.get();
            } else if (child->visitCount == bestVisits && child->averageReward > bestScore) {
                // Tie-break on average reward
                bestScore = child->averageReward;
                bestChild = child.get();
            }
        }

        // If no child was visited, fall back to lookahead score
        if (!bestChild && !children.empty()) {
            for (const auto& child : children) {
                if (child->isRedundant) continue;

                if (child->lookaheadScore > bestScore) {
                    bestScore = child->lookaheadScore;
                    bestChild = child.get();
                }
            }
        }

        return bestChild;
    }
};

}  // namespace mcts
}  // namespace shardok

#endif  // EAGLE0_ABSTRACT_MCTSNODE_HPP