//
// Command evaluator for AI lookahead search.
// Extracted from AIScoreCalculator to separate concerns.
//

#include "AICommandEvaluator.hpp"

#include <chrono>
#include <cmath>
#include <future>
#include <limits>

#include "AICommandFilter.hpp"
#include "TranspositionTable.hpp"
#include "src/main/cpp/net/eagle0/common/SequenceRandomGenerator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/AIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexCubeUtils.hpp"

namespace shardok {

// No need to forward declare internal functions - use the public interface instead

// Helper constants and static variables
static const std::vector<double> _averageSequence = {0.5};
static const auto _averageGenerator = std::make_shared<SequenceRandomGenerator>(_averageSequence);

#define MULTITHREAD true
#define LOGGING_ 0

// Helper function to determine if a command type is deterministic
static auto IsDeterministic(const CommandType type) -> bool {
    switch (type) {
        case net::eagle0::shardok::common::MOVE_COMMAND:
        case net::eagle0::shardok::common::CONTROL_COMMAND:
        case net::eagle0::shardok::common::METEOR_START_COMMAND:
        case net::eagle0::shardok::common::METEOR_TARGET_COMMAND:
        case net::eagle0::shardok::common::METEOR_CANCEL_COMMAND:
        case net::eagle0::shardok::common::END_TURN_COMMAND:
        case net::eagle0::shardok::common::PLACE_UNIT_COMMAND:
        case net::eagle0::shardok::common::PLACE_HIDDEN_UNIT_COMMAND:
        case net::eagle0::shardok::common::UNIT_STOP_COMMAND:
        case net::eagle0::shardok::common::UNIT_REST_COMMAND:
        case net::eagle0::shardok::common::FLEE_COMMAND:
        case net::eagle0::shardok::common::REINFORCE_COMMAND:
        case net::eagle0::shardok::common::RETREAT_COMMAND:
        case net::eagle0::shardok::common::END_PLAYER_SETUP_COMMAND:
        case net::eagle0::shardok::common::HIDE_COMMAND:
        case net::eagle0::shardok::common::FORTIFY_COMMAND:
        case net::eagle0::shardok::common::BECOME_OUTLAW_COMMAND:
        case net::eagle0::shardok::common::HOLY_WAVE_COMMAND:
        case net::eagle0::shardok::common::REPAIR_COMMAND: return true;
        default: return false;
    }
}

// Helper function to sort commands by score
static auto CommandSorter(
        const AICommandEvaluator::IndexAndScore& l,
        const AICommandEvaluator::IndexAndScore& r) -> bool {
    if (l.lookaheadScore < r.lookaheadScore) return true;
    if (l.lookaheadScore > r.lookaheadScore) return false;

    // At this point the scores are tied
    if (l.immediateScore < r.immediateScore) return true;
    if (l.immediateScore > r.immediateScore) return false;

    return false;
}

AICommandEvaluator::AICommandEvaluator(
        const AIScoreCalculator& scorer,
        const APDCache& apdCache,
        BattalionTypeGetter battalionTypeGetter)
    : scorer_(scorer),
      apdCache_(apdCache),
      battalionTypeGetter_(std::move(battalionTypeGetter)) {}  // Move the function object

auto AICommandEvaluator::PerformLookahead(
        const PlayerId pid,
        const bool isDefender,
        const int remainingLookahead,
        const int maxRepeatCount,
        const std::shared_ptr<ShardokEngine>& innerEngine,
        const ScoreValue currentUtility,
        const AIStrategy& attackerStrategy,
        const CoordsSet& allCastleCoords,
        std::chrono::steady_clock::time_point deadline) const -> std::future<ScoreValue> {
    // Check transposition table before expensive computation
    auto cachedScore =
            g_transpositionTable.probe(innerEngine->GetCurrentGameState(), remainingLookahead, pid);

    if (cachedScore.has_value()) {
        // Return cached result immediately
        std::promise<ScoreValue> p;
        p.set_value(*cachedScore);
        return p.get_future();
    }
    const auto nextUtility = currentUtility;

    // Check if we've reached the depth limit before making recursive calls
    if (remainingLookahead <= 0) {
        // Store the current utility in the transposition table and return it
        // Note: Store with depth 1 since depth 0 indicates an empty entry in the transposition
        // table
        g_transpositionTable.store(innerEngine->GetCurrentGameState(), 1, pid, nextUtility);

        std::promise<ScoreValue> p;
        p.set_value(nextUtility);
        return p.get_future();
    }

    if (const CommandListSPtr nextCommands = innerEngine->GetAvailableCommandsForAIPlayer(pid);
        nextCommands && !nextCommands->empty()) {
        // Get the future from FindBestCommand without calling .get()
        auto bestCommandFuture = FindBestCommand(
                pid,
                isDefender,
                remainingLookahead - 1,
                maxRepeatCount,
                *innerEngine,
                attackerStrategy,
                nextUtility,
                allCastleCoords,
                deadline);

        // Return a future that chains the best command evaluation
        return std::async(
                std::launch::deferred,
                [bestCommandFuture = std::move(bestCommandFuture),
                 innerEngine,
                 pid,
                 nextUtility,
                 remainingLookahead]() mutable -> ScoreValue {
                    const auto [index, type, lookaheadScore, immediateScore] =
                            bestCommandFuture.get();

                    ScoreValue resultScore;
                    if (auto& nextCommand =
                                innerEngine->GetAvailableCommandsForAIPlayer(pid)->at(index);
                        nextCommand->GetCommandType() !=
                        net::eagle0::shardok::common::END_TURN_COMMAND) {
                        resultScore = immediateScore;
                    } else {
                        resultScore = nextUtility;
                    }

                    // Store in transposition table before returning
                    g_transpositionTable.store(
                            innerEngine->GetCurrentGameState(),
                            remainingLookahead,
                            pid,
                            resultScore);

                    return resultScore;
                });
    }

    // No commands available, store and return the current utility as a future
    g_transpositionTable
            .store(innerEngine->GetCurrentGameState(), remainingLookahead, pid, nextUtility);

    std::promise<ScoreValue> p;
    p.set_value(nextUtility);
    return p.get_future();
}

auto AICommandEvaluator::EvaluateWithRandomness(
        const PlayerId pid,
        const bool isDefender,
        const uint32_t commandIndex,
        const int remainingLookahead,
        const int maxRepeatCount,
        const std::shared_ptr<RandomGenerator>& randomGenerator,
        const ShardokEngine& guessedEngine,
        const AIStrategy& attackerStrategy,
        const CoordsSet& allCastleCoords,
        std::chrono::steady_clock::time_point deadline) const -> ImmediateAndLookaheadScore {
    ImmediateAndLookaheadScore returnValue{};

    // Check if we've exceeded the deadline
    if (std::chrono::steady_clock::now() > deadline) {
        // Return with a default score and an empty future that resolves immediately
        std::promise<ScoreValue> p;
        p.set_value(0.0);  // Default timeout score
        returnValue.immediateScore = 0.0;
        returnValue.lookaheadScore = p.get_future();
        return returnValue;
    }

    auto innerEngine = std::make_shared<ShardokEngine>(guessedEngine, false);
    innerEngine->PostCommand(pid, commandIndex, randomGenerator);

    auto innerUtility = scorer_.GuessedStateScore(
            isDefender,
            innerEngine->GetCurrentGameState(),
            attackerStrategy,
            allCastleCoords);

    returnValue.immediateScore = innerUtility;

    if (remainingLookahead <= 0) {
        std::promise<ScoreValue> p;
        returnValue.lookaheadScore = p.get_future();
        p.set_value(innerUtility);
    } else {
        auto lookaheadLambda = [this,
                                pid,
                                isDefender,
                                remainingLookahead,
                                maxRepeatCount,
                                innerEngine,
                                attackerStrategy,
                                innerUtility,
                                &allCastleCoords,
                                deadline]() -> ScoreValue {
            auto lookaheadFuture = PerformLookahead(
                    pid,
                    isDefender,
                    remainingLookahead,
                    maxRepeatCount,
                    innerEngine,
                    innerUtility,
                    attackerStrategy,
                    allCastleCoords,
                    deadline);
            return lookaheadFuture.get();
        };

#if MULTITHREAD
        auto launchPolicy = remainingLookahead == 1 ? std::launch::async : std::launch::deferred;
        returnValue.lookaheadScore = std::async(launchPolicy, lookaheadLambda);
#else
        std::promise<ScoreValue> p;
        returnValue.lookaheadScore = p.get_future();
        auto lambdaResult = lookaheadLambda();
        p.set_value(lambdaResult);
#endif
    }

    return returnValue;
}

auto AICommandEvaluator::FindBestCommand(
        const PlayerId pid,
        const bool isDefender,
        const int remainingLookahead,
        const int maxRepeatCount,
        const ShardokEngine& guessedEngine,
        const AIStrategy& attackerStrategy,
        const ScoreValue currentUtility,
        const CoordsSet& allCastleCoords,
        std::chrono::steady_clock::time_point deadline) const -> std::future<IndexAndScore> {
    const CommandListSPtr guessedDescriptors = guessedEngine.GetAvailableCommandsForAIPlayer(pid);

    // Filter out obviously bad commands to reduce search space
    const std::vector<size_t> filteredIndices = AICommandFilter::FilterCommands(
            guessedDescriptors,
            pid,
            isDefender,
            guessedEngine.GetCurrentGameState(),
            apdCache_,
            battalionTypeGetter_);

    const auto& gameState = guessedEngine.GetCurrentGameState();
    // Calculate minimum hex distance to enemies for this player
    double minDistToEnemies = std::numeric_limits<double>::max();
    const auto* units = gameState->units();

    for (size_t i = 0; i < units->size(); ++i) {
        if (const auto* playerUnit = units->Get(static_cast<unsigned int>(i));
            playerUnit->player_id() == pid) {
            const auto& playerCoords = playerUnit->location();

            for (size_t j = 0; j < units->size(); ++j) {
                if (const auto* enemyUnit = units->Get(static_cast<unsigned int>(j));
                    enemyUnit->player_id() != pid) {
                    const auto& enemyCoords = enemyUnit->location();

                    // Proper hex distance calculation using cube coordinates
                    const Cube playerCube = OffsetToCube(playerCoords);
                    const Cube enemyCube = OffsetToCube(enemyCoords);
                    const int hexDistance = CubeDistance(playerCube, enemyCube);

                    minDistToEnemies = std::min(minDistToEnemies, static_cast<double>(hexDistance));
                }
            }
        }
    }

    if (minDistToEnemies == std::numeric_limits<double>::max()) {
        minDistToEnemies = 0.0;  // No enemies found
    }

#if LOGGING_
    // Log command count and distance metrics for performance analysis
    const auto allCommandCount = guessedDescriptors->size();
    const auto filteredCommandCount = filteredIndices.size();
    const int currentRound = gameState->current_round();

    printf("AI_COMMAND_COUNT: Round %d, Player %d, Defender %d, MinDist %.1f, Commands %zu -> %zu "
           "(%.1f%% filtered)\n",
           currentRound,
           static_cast<int>(pid),
           isDefender ? 1 : 0,
           minDistToEnemies,
           allCommandCount,
           filteredCommandCount,
           100.0 * (allCommandCount - filteredCommandCount) / allCommandCount);
#endif

    const auto commandCount = filteredIndices.size();

    // Structure to hold all command evaluation data
    struct CommandEvaluation {
        size_t index;
        CommandType type;
        ScoreValue immediateScore;
        std::vector<std::future<ScoreValue>> lookaheadFutures;
    };

    std::vector<CommandEvaluation> commandEvaluations(commandCount);

    for (uint32_t index = 0; index < commandCount; index++) {
        const auto originalIndex = filteredIndices[index];
        const auto& guessedDescriptor = guessedDescriptors->at(originalIndex);
        const auto guessedCommandType = guessedDescriptor->GetCommandType();

        commandEvaluations[index].index = originalIndex;
        commandEvaluations[index].type = guessedCommandType;

        if (guessedCommandType == net::eagle0::shardok::common::END_TURN_COMMAND) {
            std::promise<ScoreValue> p;
            commandEvaluations[index].lookaheadFutures.push_back(p.get_future());
            p.set_value(currentUtility);
            commandEvaluations[index].immediateScore = currentUtility;
        } else if (IsDeterministic(guessedCommandType)) {
            auto [immediateScore, lookaheadScore] = EvaluateWithRandomness(
                    pid,
                    isDefender,
                    originalIndex,
                    remainingLookahead,
                    maxRepeatCount,
                    _averageGenerator,
                    guessedEngine,
                    attackerStrategy,
                    allCastleCoords,
                    deadline);

            commandEvaluations[index].immediateScore = immediateScore;
            commandEvaluations[index].lookaheadFutures.push_back(std::move(lookaheadScore));
        } else if (guessedDescriptor->HasOdds()) {
            const auto successChancePercentile = guessedDescriptor->GetOddsPercentile();
            const double successChance = static_cast<double>(successChancePercentile) / 100.0;

            // Success attempt uses 1.0 - (successChance / 2) as the roll
            auto [successImmediateScore, successLookaheadScore] = EvaluateWithRandomness(
                    pid,
                    isDefender,
                    originalIndex,
                    remainingLookahead,
                    maxRepeatCount,
                    std::make_shared<SequenceRandomGenerator>(
                            std::vector{1.0 - successChance / 2.0}),
                    guessedEngine,
                    attackerStrategy,
                    allCastleCoords,
                    deadline);

            // Failure attempt uses the average of (1 - successChance) and 0 as the roll
            auto [failureImmediateScore, failureLookaheadScore] = EvaluateWithRandomness(
                    pid,
                    isDefender,
                    originalIndex,
                    remainingLookahead,
                    maxRepeatCount,
                    std::make_shared<SequenceRandomGenerator>(
                            std::vector{(1.0 - successChance) / 2.0}),
                    guessedEngine,
                    attackerStrategy,
                    allCastleCoords,
                    deadline);

            commandEvaluations[index].immediateScore =
                    std::lerp(failureImmediateScore, successImmediateScore, successChance);

            auto successSF = successLookaheadScore.share();
            auto failureSF = failureLookaheadScore.share();
            commandEvaluations[index].lookaheadFutures.push_back(std::async(
                    std::launch::deferred,
                    [successSF, failureSF, successChance]() -> double {
                        return std::lerp(failureSF.get(), successSF.get(), successChance);
                    }));
        } else {
            ScoreValue sum = 0.0;
            for (int repeatIteration = 0; repeatIteration < maxRepeatCount; repeatIteration++) {
                // In each iteration, use a double from [0, 1] as the random roll
                auto sequence = std::vector{
                        static_cast<double>(repeatIteration) /
                        static_cast<double>(maxRepeatCount - 1)};
                auto [immediateScore, lookaheadScore] = EvaluateWithRandomness(
                        pid,
                        isDefender,
                        originalIndex,
                        remainingLookahead,
                        maxRepeatCount,
                        std::make_shared<SequenceRandomGenerator>(sequence),
                        guessedEngine,
                        attackerStrategy,
                        allCastleCoords,
                        deadline);

                sum += immediateScore;
                commandEvaluations[index].lookaheadFutures.push_back(std::move(lookaheadScore));
            }
            commandEvaluations[index].immediateScore = sum / maxRepeatCount;
        }
    }

    // Return a future that will wait for all evaluations and find the best one
    return std::async(
            std::launch::deferred,
            [evals = std::move(commandEvaluations)]() mutable -> IndexAndScore {
                std::vector<IndexAndScore> allResults;
                allResults.reserve(evals.size());

                // Wait for all futures and compute final scores
                for (auto& eval : evals) {
                    ScoreValue totalLookaheadScore = 0.0;
                    for (auto& future : eval.lookaheadFutures) {
                        totalLookaheadScore += future.get();
                    }
                    ScoreValue avgLookaheadScore =
                            eval.lookaheadFutures.empty()
                                    ? eval.immediateScore
                                    : totalLookaheadScore / eval.lookaheadFutures.size();

                    allResults.push_back(IndexAndScore{
                            .index = eval.index,
                            .type = eval.type,
                            .lookaheadScore = avgLookaheadScore,
                            .immediateScore = eval.immediateScore});
                }
                // Find the best command using the existing sorter
                auto bestIt = std::ranges::max_element(allResults, CommandSorter);
                return *bestIt;
            });
}

auto AICommandEvaluator::EvaluateCommand(
        const PlayerId pid,
        const bool isDefender,
        const int remainingLookahead,
        const int maxRepeatCount,
        const ShardokEngine& guessedEngine,
        const AIStrategy& attackerStrategy,
        const ScoreValue currentUtility,
        const CoordsSet& allCastleCoords,
        const size_t commandIndex,
        std::chrono::steady_clock::time_point deadline) const -> std::future<ScoreValue> {
    const CommandListSPtr guessedDescriptors = guessedEngine.GetAvailableCommandsForAIPlayer(pid);

    if (commandIndex >= guessedDescriptors->size()) {
        std::promise<ScoreValue> p;
        p.set_value(currentUtility);
        return p.get_future();
    }

    const auto& guessedDescriptor = guessedDescriptors->at(commandIndex);

    if (const auto guessedCommandType = guessedDescriptor->GetCommandType();
        guessedCommandType == net::eagle0::shardok::common::END_TURN_COMMAND) {
        std::promise<ScoreValue> p;
        p.set_value(currentUtility);
        return p.get_future();
    } else if (IsDeterministic(guessedCommandType)) {
        auto [immediateScore, lookaheadScore] = EvaluateWithRandomness(
                pid,
                isDefender,
                commandIndex,
                remainingLookahead,
                maxRepeatCount,
                _averageGenerator,
                guessedEngine,
                attackerStrategy,
                allCastleCoords,
                deadline);
        return std::move(lookaheadScore);
    } else if (guessedDescriptor->HasOdds()) {
        const auto successChancePercentile = guessedDescriptor->GetOddsPercentile();
        const double successChance = static_cast<double>(successChancePercentile) / 100.0;

        // Success attempt
        auto [successImmediateScore, successLookaheadScore] = EvaluateWithRandomness(
                pid,
                isDefender,
                commandIndex,
                remainingLookahead,
                maxRepeatCount,
                std::make_shared<SequenceRandomGenerator>(std::vector{1.0 - successChance / 2.0}),
                guessedEngine,
                attackerStrategy,
                allCastleCoords,
                deadline);

        // Failure attempt
        auto [failureImmediateScore, failureLookaheadScore] = EvaluateWithRandomness(
                pid,
                isDefender,
                commandIndex,
                remainingLookahead,
                maxRepeatCount,
                std::make_shared<SequenceRandomGenerator>(std::vector{(1.0 - successChance) / 2.0}),
                guessedEngine,
                attackerStrategy,
                allCastleCoords,
                deadline);

        // Return weighted average of success and failure
        auto successSF = successLookaheadScore.share();
        auto failureSF = failureLookaheadScore.share();
        return std::async(std::launch::deferred, [successSF, failureSF, successChance]() -> double {
            return std::lerp(failureSF.get(), successSF.get(), successChance);
        });
    } else {
        // For non-deterministic commands without odds, use multiple attempts
        std::vector<std::future<ScoreValue>> lookaheadFutures;
        lookaheadFutures.reserve(maxRepeatCount);

        for (int repeatIteration = 0; repeatIteration < maxRepeatCount; repeatIteration++) {
            auto sequence = std::vector{
                    static_cast<double>(repeatIteration) / static_cast<double>(maxRepeatCount - 1)};
            auto [immediateScore, lookaheadScore] = EvaluateWithRandomness(
                    pid,
                    isDefender,
                    commandIndex,
                    remainingLookahead,
                    maxRepeatCount,
                    std::make_shared<SequenceRandomGenerator>(sequence),
                    guessedEngine,
                    attackerStrategy,
                    allCastleCoords,
                    deadline);

            lookaheadFutures.push_back(std::move(lookaheadScore));
        }

        // Return a future that computes the average when needed
        return std::async(
                std::launch::deferred,
                [lookaheadFutures = std::move(lookaheadFutures),
                 maxRepeatCount]() mutable -> double {
                    ScoreValue total = 0.0;
                    for (auto& future : lookaheadFutures) { total += future.get(); }
                    return total / maxRepeatCount;
                });
    }
}

}  // namespace shardok
