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

#include "AITimeBudget.hpp"

#include "src/main/cpp/net/eagle0/shardok/library/fb_helpers/FlatbufferWrapper.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/settings/GameSettings.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexCubeUtils.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexMapUtils.hpp"
#include "src/main/flatbuffer/net/eagle0/shardok/storage/game_state.hpp"

namespace shardok {

// Static member definition
std::atomic<int> AIEvaluationCounter::activeCount{0};

AIEvaluationCounter::AIEvaluationCounter() { activeCount++; }

AIEvaluationCounter::~AIEvaluationCounter() { activeCount--; }

int AIEvaluationCounter::GetCurrentCount() { return activeCount.load(); }

auto CalculateTimeBudget(
        const PlayerId playerId,
        const GameSettingsSPtr &settings,
        const GameStateW &state,
        const size_t numCommands,
        const bool isAllAiBattle) -> AITimeBudget {
    const auto settingsGetter = settings->GetGetter();
    const auto castleCoords = AllCastleCoords(state->hex_map());

    // Check if we're in setup phase
    const bool isSetupPhase = state->status()->state() ==
                              net::eagle0::shardok::storage::fb::GameStatus_::State_SET_UP;

    // Get maximum budget cap from settings (in seconds)
    // For all-AI battles, use the faster budget limit
    const double maxBudgetSeconds =
            isAllAiBattle ? settingsGetter.Backing().all_ai_battle_time_budget_maximum()
                          : settingsGetter.Backing().lookahead_time_budget_maximum_seconds();
    const double maxBudgetMs = maxBudgetSeconds * 1000.0;

    // During setup, use the setup-specific time budget
    if (isSetupPhase) {
        // Dynamic budget: msPerCommand × numCommands
        const double msPerCommand =
                settingsGetter.Backing().lookahead_time_budget_per_command_setup_ms();
        const double budgetMs = msPerCommand * static_cast<double>(numCommands);

        // Clamp to reasonable bounds: 200ms minimum, maxBudgetMs maximum
        const auto clampedBudgetMs = std::clamp(budgetMs, 200.0, maxBudgetMs);
        const auto remainingBudget =
                std::chrono::milliseconds(static_cast<int64_t>(clampedBudgetMs));

        const size_t minDepth = settingsGetter.Backing().min_lookahead_turns();

        return AITimeBudget{
                .remainingBudget = remainingBudget,
                .minDepthRequired = minDepth,
                .isCloseToEnemy = false};  // Not relevant during setup
    }

    // Determine proximity (≤4 hex distance) - applies to both attackers and defenders
    bool isClose = false;
    const auto *units = state->units();

    for (size_t i = 0; i < units->size() && !isClose; ++i) {
        const auto *myUnit = units->Get(static_cast<unsigned int>(i));
        if (myUnit->player_id() != playerId) continue;

        const auto &myCoords = myUnit->location();
        // Skip units that haven't been placed on the map yet
        if (myCoords.row() == -1) continue;

        const Cube myCube = OffsetToCube(myCoords);

        // Check distance to enemy units
        for (size_t j = 0; j < units->size(); ++j) {
            const auto *enemyUnit = units->Get(static_cast<unsigned int>(j));
            if (enemyUnit->player_id() == playerId) continue;

            const auto &enemyCoords = enemyUnit->location();
            // Skip enemy units that haven't been placed on the map yet
            if (enemyCoords.row() == -1) continue;

            const Cube enemyCube = OffsetToCube(enemyCoords);

            if (const int hexDistance = CubeDistance(myCube, enemyCube); hexDistance <= 4) {
                isClose = true;
                break;
            }
        }

        // Check distance to castles
        if (!isClose) {
            for (const auto &castleCoord : castleCoords) {
                const Cube castleCube = OffsetToCube(castleCoord);

                if (const int hexDistance = CubeDistance(myCube, castleCube); hexDistance <= 4) {
                    isClose = true;
                    break;
                }
            }
        }
    }

    // Get time budget from settings - dynamic based on number of commands
    // Dynamic budget: msPerCommand × numCommands
    const double msPerCommand =
            isClose ? settingsGetter.Backing().lookahead_time_budget_per_command_close_ms()
                    : settingsGetter.Backing().lookahead_time_budget_per_command_far_ms();
    const double budgetMs = msPerCommand * static_cast<double>(numCommands);

    // Clamp to reasonable bounds: 200ms minimum, maxBudgetMs maximum
    const auto clampedBudgetMs = std::clamp(budgetMs, 200.0, maxBudgetMs);
    const auto remainingBudget = std::chrono::milliseconds(static_cast<int64_t>(clampedBudgetMs));

    // TEMPORARY DEBUG OUTPUT
    printf("[DEBUG CalculateTimeBudget] numCommands=%zu, msPerCommand=%.2f, budgetMs=%.2f, "
           "clampedBudgetMs=%.2f, isClose=%d\n",
           numCommands,
           msPerCommand,
           budgetMs,
           clampedBudgetMs,
           isClose);

    // Get minimum depth requirement
    const size_t minDepth = settingsGetter.Backing().min_lookahead_turns();

    return AITimeBudget{
            .remainingBudget = remainingBudget,
            .minDepthRequired = minDepth,
            .isCloseToEnemy = isClose};
}

}  // namespace shardok