//
// Created by Dan Crosby on 07/04/25.
//

#include "IterativeDeepeningAI.hpp"

#include <algorithm>
#include <limits>
#include <numeric>
#include <utility>

#include "AIAttackerStrategySelector.hpp"
#include "AIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/common/TimeUtils.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"

namespace shardok {

#define DEBUG_ITERATIVE_DEEPENING_TIMINGS 1

IterativeDeepeningAI::IterativeDeepeningAI(
        const PlayerId playerId,
        const bool isDefender,
        AIStrategy strategy,
        const CoordsSet& castleCoords,
        const APDCache& apdCache,
        const ALCache& alCache)
    : playerId(playerId),
      isDefender(isDefender),
      strategy(std::move(strategy)),
      castleCoords(castleCoords),
      apdCache(apdCache),
      alCache(alCache) {}

auto IterativeDeepeningAI::IterativeSearch(
        const GameSettingsSPtr& settings,
        const GameStateW& state,
        const std::vector<CommandProto>& commands,
        const AITimeBudget& initialBudget) const -> SearchResult {
    // Make a mutable copy of the time budget to track remaining time
    AITimeBudget timeBudget = initialBudget;
    const auto startTime = std::chrono::steady_clock::now();
    const auto initialBudgetMs = initialBudget.remainingBudget;
    SearchResult result;

    if (commands.empty()) {
#if DEBUG_ITERATIVE_DEEPENING_TIMINGS
        printf("ID AI: Commands are empty, returning early\n");
#endif
        result.searchCompleted = true;
        return result;
    }

    // Check if we're in SET_UP phase
    bool isSetupPhase =
            (state->status()->state() ==
             net::eagle0::shardok::storage::fb::GameStatus_::State_SET_UP);
    int maxDepth = isSetupPhase ? 2 : std::numeric_limits<int>::max();

    // Calculate current utility and create engine once for all command evaluations
    const auto& settingsGetter = settings->GetGetter();
    const auto guessedEngine = ShardokEngine(settings, state);
    const auto maxRepeatCount = settingsGetter.Backing().ai_utility_repeat_count();
    const ScoreValue currentUtility = AIScoreCalculator::GuessedStateScore(
            isDefender,
            state,
            strategy,
            castleCoords,
            settingsGetter,
            apdCache,
            alCache);

    // Initialize data structures for tracking scores at each depth
    scoresByDepth.clear();
    scoresByDepth.resize(commands.size());
    highestDepthCompleted.clear();
    highestDepthCompleted.resize(commands.size(), 0);

    int currentDepth = 1;
    size_t previousBestCommand = 0;  // Track best command from previous depth
    size_t evaluatedCountAtHighestDepth = 0;
    EvaluationCompletionReason completionReason = EvaluationCompletionReason::RAN_OUT_OF_TIME;

    // Main iterative deepening loop
    while ((currentDepth == 1 || !IsTimeExpired(timeBudget)) && currentDepth <= maxDepth) {
        // Get command indices sorted by best score from previous depth
        std::vector<size_t> sortedIndices = GetCommandsSortedByPreviousDepth(
                currentDepth,
                scoresByDepth,
                highestDepthCompleted);

        int evaluatedCount = 0;
        bool allEvaluated = true;
        bool allEndTurnCommands = true;  // Track if all commands are END_TURN

        // Try to evaluate all commands at this depth, within budget constraints
        for (size_t cmdIndex : sortedIndices) {
            if (currentDepth > 1 && IsTimeExpired(timeBudget)) {
                allEvaluated = false;
                break;
            }

            auto cmdResult = SearchCommandAtDepthWithEngine(
                    guessedEngine,
                    settingsGetter,
                    maxRepeatCount,
                    commands,
                    cmdIndex,
                    currentDepth,
                    currentUtility,
                    timeBudget);

            // Ensure scoresByDepth[cmdIndex] has enough space
            if (scoresByDepth[cmdIndex].size() <= currentDepth) {
                scoresByDepth[cmdIndex].resize(currentDepth + 1);
            }
            scoresByDepth[cmdIndex][currentDepth] = cmdResult.bestScore;
            highestDepthCompleted[cmdIndex] = currentDepth;
            evaluatedCount++;

            // Check if this command is not END_TURN_COMMAND
            if (commands[cmdIndex].type() != net::eagle0::shardok::common::END_TURN_COMMAND) {
                allEndTurnCommands = false;
            }
        }

        // Find the best command at current depth and check if it changed
        if (evaluatedCount > 0) {
            evaluatedCountAtHighestDepth = evaluatedCount;
            size_t currentBestCommand = 0;
            ScoreValue currentBestScore = -std::numeric_limits<ScoreValue>::infinity();

            for (size_t i = 0; i < commands.size(); ++i) {
                if (highestDepthCompleted[i] >= currentDepth) {
                    if (scoresByDepth[i][currentDepth] > currentBestScore) {
                        currentBestScore = scoresByDepth[i][currentDepth];
                        currentBestCommand = i;
                    }
                }
            }

            // Log if best command changed from previous depth
            if (currentDepth > 1 && currentBestCommand != previousBestCommand) {
#if DEBUG_ITERATIVE_DEEPENING_TIMINGS
                printf("ID AI: Best command changed at depth %d:\n", currentDepth);
                printf("  Depth %d best: command %zu (score %.2f) - %s\n",
                       currentDepth - 1,
                       previousBestCommand,
                       scoresByDepth[previousBestCommand][currentDepth - 1],
                       commands[previousBestCommand].DebugString().c_str());
                printf("  Depth %d best: command %zu (score %.2f) - %s\n",
                       currentDepth,
                       currentBestCommand,
                       currentBestScore,
                       commands[currentBestCommand].DebugString().c_str());
#endif
            }

            previousBestCommand = currentBestCommand;
        }

        // Only proceed to next depth if we completed all commands at current depth
        if (!allEvaluated) {
            completionReason = EvaluationCompletionReason::RAN_OUT_OF_TIME;
            break;
        }

        // Stop if all evaluated commands were END_TURN at the root - no point going deeper
        if (allEndTurnCommands && evaluatedCount > 0) {
            completionReason = EvaluationCompletionReason::RAN_OUT_OF_COMMANDS;
            break;
        }

        // Also check if scores haven't changed from previous depth
        // This indicates we've hit END_TURN in the lookahead
        if (currentDepth > 1 && evaluatedCount > 0) {
            bool scoresUnchanged = true;
            int unchangedCount = 0;

            for (size_t i = 0; i < sortedIndices.size() && i < evaluatedCount; ++i) {
                size_t cmdIndex = sortedIndices[i];
                // This command was evaluated at both current and previous depth
                if (scoresByDepth[cmdIndex].size() > currentDepth &&
                    scoresByDepth[cmdIndex].size() > currentDepth - 1) {
                    // Check if score changed between depth N-1 and depth N
                    if (std::abs(
                                scoresByDepth[cmdIndex][currentDepth] -
                                scoresByDepth[cmdIndex][currentDepth - 1]) < 1e-9) {
                        unchangedCount++;
                    } else {
                        scoresUnchanged = false;
                        break;
                    }
                }
            }

            // If all evaluated commands had unchanged scores, we've hit END_TURN in lookahead
            if (scoresUnchanged && unchangedCount == evaluatedCount) {
                completionReason = EvaluationCompletionReason::RAN_OUT_OF_COMMANDS;
                break;
            }
        }

        // Check if we've used more than 50% of total budget
        auto totalElapsed = std::chrono::steady_clock::now() - startTime;
        auto totalElapsedMs = std::chrono::duration_cast<std::chrono::milliseconds>(totalElapsed);
        double budgetUsedPercent = (double)totalElapsedMs.count() / initialBudgetMs.count();

        if (budgetUsedPercent > 0.5) {
            printf("ID AI: Stopping after depth %d - used %.1f%% of time budget\n",
                   currentDepth,
                   budgetUsedPercent * 100);
            completionReason = EvaluationCompletionReason::NOT_ENOUGH_TIME_TO_CONTINUE;
            break;
        }

        currentDepth++;
    }

    // If we completed the loop without any breaks, we successfully exhausted meaningful search
    if (completionReason == EvaluationCompletionReason::RAN_OUT_OF_TIME &&
        currentDepth > maxDepth) {
        // We hit the depth limit rather than running out of time
        completionReason = EvaluationCompletionReason::RAN_OUT_OF_COMMANDS;
    }

    // Select best result from highest depth achieved for each command
    result = SelectBestResult(scoresByDepth, highestDepthCompleted);
    result.minimumDepthCompleted = result.depthAchieved >= timeBudget.minDepthRequired;
    result.searchCompleted = result.minimumDepthCompleted;
    result.timeUsed = std::chrono::duration_cast<std::chrono::milliseconds>(
            std::chrono::steady_clock::now() - startTime);
    result.availableCommandCount = commands.size();
    result.commandCountEvaluated = evaluatedCountAtHighestDepth;
    result.completionReason = completionReason;

    // Validation: if completion reason is RAN_OUT_OF_COMMANDS, evaluation should be 100%
    if (completionReason == EvaluationCompletionReason::RAN_OUT_OF_COMMANDS &&
        result.commandCountEvaluated < result.availableCommandCount) {
        printf("ERROR: Completion reason RAN_OUT_OF_COMMANDS but evaluation %lu/%zu < 100%%\n",
               result.commandCountEvaluated,
               result.availableCommandCount);
    }

    return result;
}

bool IterativeDeepeningAI::IsTimeExpired(const AITimeBudget& budget) {
    return budget.remainingBudget <= std::chrono::milliseconds(0);
}

auto IterativeDeepeningAI::SearchCommandAtDepthWithEngine(
        const ShardokEngine& guessedEngine,
        const GameSettings::Getter& settingsGetter,
        const int maxRepeatCount,
        const std::vector<CommandProto>& commands,
        const size_t commandIndex,
        const int depth,
        const ScoreValue currentUtility,
        AITimeBudget& timeBudget) const -> SearchResult {
    SearchResult result;
    result.bestCommandIndex = commandIndex;
    result.depthAchieved = depth;
    result.searchCompleted = true;
    result.minimumDepthCompleted = true;
    result.availableCommandCount = commands.size();
    result.commandCountEvaluated = 1;  // We're evaluating just this command

    if (commandIndex >= commands.size()) {
        result.bestScore = 0.0;
        return result;
    }

    try {
        // Track concurrent evaluations and adjust time accounting
        AIEvaluationCounter counter;
        const auto startTime = std::chrono::steady_clock::now();

        // Use CommandScore to evaluate the specific command at the given depth
        const auto commandScore = AIScoreCalculator::CommandScore(
                playerId,
                isDefender,
                depth,
                maxRepeatCount,
                guessedEngine,
                strategy,
                currentUtility,
                settingsGetter,
                castleCoords,
                apdCache,
                alCache,
                commandIndex);

        // Calculate time used and adjust based on concurrent evaluations
        const auto elapsed = std::chrono::steady_clock::now() - startTime;
        const int concurrentCount = counter.GetCurrentCount();
        const auto adjustedElapsed = elapsed / std::max(1, concurrentCount);
        const auto adjustedElapsedMs =
                std::chrono::duration_cast<std::chrono::milliseconds>(adjustedElapsed);

        // Deduct adjusted time from remaining budget
        timeBudget.remainingBudget -= adjustedElapsedMs;

        result.bestScore = commandScore;
    } catch (const std::exception& e) {
        // If evaluation fails, return a neutral score rather than crashing
#if DEBUG_ITERATIVE_DEEPENING_TIMINGS
        printf("SearchCommandAtDepthWithEngine: evaluation failed with exception: %s\n", e.what());
#endif
        result.bestScore = 0.0;
    }

    return result;
}

auto IterativeDeepeningAI::GetCommandsSortedByPreviousDepth(
        int currentDepth,
        const std::vector<std::vector<ScoreValue>>& scoresByDepth,
        const std::vector<int>& highestDepthCompleted) const -> std::vector<size_t> {
    std::vector<size_t> indices(scoresByDepth.size());
    std::iota(indices.begin(), indices.end(), 0);

    if (currentDepth == 1) {
        // For depth 1, return natural order
        return indices;
    }

    // Sort by score at previous depth
    int prevDepth = currentDepth - 1;
    std::sort(indices.begin(), indices.end(), [&](size_t a, size_t b) {
        // Only consider commands that were evaluated at previous depth
        if (highestDepthCompleted[a] >= prevDepth && highestDepthCompleted[b] >= prevDepth) {
            return scoresByDepth[a][prevDepth] > scoresByDepth[b][prevDepth];
        }
        // Commands not evaluated at prev depth go to the end
        return highestDepthCompleted[a] >= prevDepth;
    });

    return indices;
}

auto IterativeDeepeningAI::SelectBestResult(
        const std::vector<std::vector<ScoreValue>>& scoresByDepth,
        const std::vector<int>& highestDepthCompleted) const -> SearchResult {
    SearchResult result;
    result.bestScore = -std::numeric_limits<ScoreValue>::infinity();
    result.searchCompleted = false;

    // Find the command with best score at its highest evaluated depth
    for (size_t i = 0; i < scoresByDepth.size(); ++i) {
        if (highestDepthCompleted[i] > 0) {
            int depth = highestDepthCompleted[i];
            ScoreValue score = scoresByDepth[i][depth];
            if (score > result.bestScore) {
                result.bestScore = score;
                result.bestCommandIndex = i;
                result.depthAchieved = depth;
            }
        }
    }

    return result;
}

}  // namespace shardok