//
// Abstract MCTS AI implementation
//

#include "AbstractMCTSAI.hpp"

#include <algorithm>
#include <atomic>
#include <future>
#include <mutex>
#include <random>
#include <thread>

namespace shardok::mcts {

AbstractMCTSAI::AbstractMCTSAI(MCTSPlayerId playerId, MCTSConfig config)
    : playerId_(playerId),
      config_(config) {}

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 || rootNode->children.empty()) {
        // Fallback to first action if no tree was built
        result.bestActionIndex = 0;
        result.bestScore = 0.0;
        return result;
    }

    // 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->averageReward;
        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);
    }

    return result;
}

auto AbstractMCTSAI::BuildMCTSTree(
        const MCTSGameEngine& engine,
        const MCTSGameState& initialState,
        const std::chrono::steady_clock::time_point deadline) const -> std::unique_ptr<MCTSNode> {
    // Create root node
    auto root = std::make_unique<MCTSNode>(initialState.clone(), playerId_, 0);

    // Get legal actions from engine for the root state
    const auto rootActions = engine.getLegalActions(initialState);

    // Initialize untried actions from the root actions
    root->untriedActionIndices.reserve(rootActions.size());
    for (size_t i = 0; i < rootActions.size(); ++i) { root->untriedActionIndices.push_back(i); }

    std::atomic<int> iterations{0};
    constexpr int maxIterations = 100000;

    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 &&
                       iterations.load() < maxIterations) {
                    // 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_);
                        std::lock_guard lock(treeMutex);
                        MCTSBackpropagation(expanded, reward);
                        iterations.fetch_add(1);
                    }
                }
            }));
        }

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

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

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

            // Backpropagation
            MCTSBackpropagation(expanded, reward);

            ++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->isTerminal && current->depth < config_.maxTreeDepth) {
        if (current->CanExpand()) {
            return current;  // Node has untried actions
        } else if (!current->children.empty()) {
            current = current->GetBestChild(config_.explorationConstant);
            if (!current) break;
        } else {
            break;  // Leaf node
        }
    }

    return current;
}

auto AbstractMCTSAI::MCTSExpansion(MCTSNode* node, const MCTSGameEngine& engine) const
        -> MCTSNode* {
    if (node->untriedActionIndices.empty() || node->isTerminal) {
        return node;  // Nothing to expand
    }

    // Select a random untried action
    thread_local std::mt19937 gen(std::random_device{}());
    std::uniform_int_distribution<size_t> dis(0, node->untriedActionIndices.size() - 1);
    const size_t randomIndex = dis(gen);
    const size_t actionIndex = node->untriedActionIndices[randomIndex];

    // Remove from untried list
    node->untriedActionIndices.erase(std::next(
            node->untriedActionIndices.begin(),
            static_cast<std::vector<size_t>::difference_type>(randomIndex)));

    if (node->untriedActionIndices.empty()) { node->fullyExpanded = true; }

    // Get legal actions from engine (uses cached engine for performance)
    const auto nodeActions = engine.getLegalActions(*node->gameState);

    // Create new child node
    if (actionIndex >= nodeActions.size()) {
        return node;  // Invalid action index
    }

    const auto& action = nodeActions[actionIndex];
    auto newState = engine.applyAction(*node->gameState, *action);
    if (!newState) {
        return node;  // Failed to apply action
    }

    auto child = std::make_unique<MCTSNode>(
            action->clone(),
            std::move(newState),
            node->gameState->currentPlayerId(),
            node->depth + 1,
            actionIndex);

    // Set up child's untried actions if not terminal
    if (!child->isTerminal) {
        const auto childActions = engine.getLegalActions(*child->gameState);
        child->untriedActionIndices.reserve(childActions.size());
        for (size_t i = 0; i < childActions.size(); ++i) {
            child->untriedActionIndices.push_back(i);
        }
    }

    // Calculate immediate and lookahead scores
    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 -> double {
    if (state.isTerminal()) { return state.score(startingPlayer); }

    // Create a mutable copy for simulation
    auto currentState = state.clone();
    int depth = 0;

    // Simulate until terminal or max depth
    while (!currentState->isTerminal() && depth < config_.maxSimulationDepth) {
        const auto actions = engine.getLegalActions(*currentState);
        if (actions.empty()) { break; }

        // Select action based on simulation policy
        const size_t selectedIndex = SelectSimulationAction(engine, *currentState, actions);
        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) -> void {
    while (node) {
        node->visitCount++;
        node->totalReward += reward;
        node->averageReward = node->totalReward / node->visitCount;

        // Update lookahead score as weighted 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;
        }

        node = node->parent;
    }
}

auto AbstractMCTSAI::SelectSimulationAction(
        const MCTSGameEngine& engine,
        const MCTSGameState& state,
        const std::vector<std::unique_ptr<MCTSAction>>& actions) const -> size_t {
    if (actions.empty()) { return 0; }

    thread_local std::mt19937 gen(std::random_device{}());

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

        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[dis(gen)];
            }
            // Fall back to random
            std::uniform_int_distribution<size_t> dis(0, actions.size() - 1);
            return dis(gen);
        }

        case MCTSSimulationPolicy::BEST_IMMEDIATE: {
            double bestScore = -std::numeric_limits<double>::max();
            size_t bestIndex = 0;

            for (size_t i = 0; i < actions.size(); ++i) {
                const double score =
                        engine.getActionScore(state, *actions[i], state.currentPlayerId());
                if (score > bestScore) {
                    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) {
                const double score =
                        engine.getActionScore(state, *actions[i], state.currentPlayerId());
                scores.emplace_back(i, score);
            }

            // Sort by score
            std::ranges::sort(scores, [](const auto& a, const auto& b) {
                return a.second > b.second;
            });

            // 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[dis(gen)].first;
        }
    }

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

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";
    printf("MCTS: Selected %s (visit:%d, reward:%.2f, lookahead:%.2f) from %zu options\n",
           selectedDesc.c_str(),
           bestChild->visitCount,
           bestChild->averageReward,
           bestChild->lookaheadScore,
           rootNode->children.size());

    // 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;
    });

    printf("MCTS: Top actions by visits:\n");
    for (size_t i = 0; i < std::min(static_cast<size_t>(3), sortedChildren.size()); ++i) {
        const auto* child = sortedChildren[i];
        printf("  [%zu] visits:%d immediate:%.2f backprop:%.2f",
               i,
               child->visitCount,
               child->immediateScore,
               child->averageReward);

        // Show the action's own description
        if (child->action) { printf(" %s", child->action->getDescription().c_str()); }

        // 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) {
                printf(" -> %s", bestNext->action->getDescription().c_str());
            }
        }
        printf("\n");
    }

    printf("MCTS: Tree stats - max depth:%d, total nodes:%d, root visits:%d\n",
           result.searchDepth,
           result.nodesEvaluated,
           rootNode->visitCount);

    // 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()) {
        printf("MCTS: Best sequence from chosen action (final: %.2f):\n", sequenceScore);
        for (size_t i = 0; i < bestSequence.size(); ++i) {
            const auto* node = bestSequence[i];
            printf("  %zu.", i + 1);
            if (node->action) { printf(" %s", node->action->getDescription().c_str()); }
            printf(" (visits:%d, immediate:%.2f, backprop:%.2f)\n",
                   node->visitCount,
                   node->immediateScore,
                   node->averageReward);
        }
    }
}

}  // namespace shardok::mcts
