//
// Abstract base class for AI score calculator implementations
//

#include "AbstractAIScoreCalculator.hpp"

#include "AIScoreCalculatorSharedUtilities.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIScoreUtilities.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIStrategy.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIUnitScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/AIWaterCrossingCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/AIVictoryConditionScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexMapUtils.hpp"

namespace shardok {

using net::eagle0::shardok::storage::fb::BattalionTypeId;
using net::eagle0::shardok::storage::fb::Unit;

using score_calculator_internal::AttackerMultiplierForTargetDistance;
using score_calculator_internal::CAPTURED_UNIT_SCORE;
using score_calculator_internal::CAPTURED_VIP_SCORE;
using score_calculator_internal::EffectiveDistanceCache;
using score_calculator_internal::UnitsScoreComponents;

auto AbstractAIScoreCalculator::CalculateUnitsScoreComponents(
        const GameStateW &gameState,
        int roundsRemaining,
        bool attackerWantsCastles,
        bool defenderShouldScatter,
        const vector<TargetPriorityList> &attackerTargetPriorities,
        const MapId &mapId) const -> UnitsScoreComponents {
    // Cache frequently accessed FlatBuffer fields to avoid repeated offset calculations
    const auto *gameStateRawPtr = gameState.Get();
    const auto *cachedUnits = gameStateRawPtr->units();
    const auto *cachedHexMap = gameStateRawPtr->hex_map();

    const int16_t cachedRowCount = cachedHexMap->row_count();
    const int16_t cachedColumnCount = cachedHexMap->column_count();
    const int cachedCurrentRound = gameStateRawPtr->current_round();

    bool isLateGame = cachedCurrentRound > 18;  // Inline IsLateGame for efficiency

    // APDCache now has built-in thread-local caching - no need for PreCachedAPDs
    ActionPoints braveWaterCost = GetBraveWaterCost();

    // Memoization cache for EffectiveDistance calls
    EffectiveDistanceCache distanceCache;

    // Early return for empty game states
    const size_t estimatedUnitCount = cachedUnits->size();
    if (estimatedUnitCount == 0) { return UnitsScoreComponents{0.0, 0.0}; }

    std::vector<const Unit *> attackerUnits{};
    std::vector<const Unit *> defenderUnits{};
    // Pre-allocate vectors based on estimated unit ratios to avoid reallocations
    attackerUnits.reserve(estimatedUnitCount);
    defenderUnits.reserve(estimatedUnitCount);

    double attackerUnitsValue = 0;
    double defenderUnitsValue = 0;

    auto occupants = Occupants(*cachedUnits, cachedRowCount, cachedColumnCount);

    for (const Unit *unit : *cachedUnits) {
        const auto *pi = PlayerInfoForPid(gameState, unit->player_id());
        if (pi == nullptr) { continue; }

        switch (unit->status()) {
            case net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT: {
                if (pi->is_defender()) {
                    defenderUnits.push_back(unit);
                } else {
                    attackerUnits.push_back(unit);
                }
                break;
            }
            case net::eagle0::shardok::storage::fb::UnitStatus_CAPTURED_UNIT: {
                double thisScore = unit->has_attached_hero() && unit->attached_hero().is_vip()
                                           ? CAPTURED_VIP_SCORE
                                           : CAPTURED_UNIT_SCORE;
                if (pi->is_defender()) {
                    defenderUnitsValue += thisScore;
                } else {
                    attackerUnitsValue += thisScore;
                }
                break;
            }

            case net::eagle0::shardok::storage::fb::UnitStatus_DESTROYED_SUMMONED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_FLED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_NEVER_ENTERED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_OUTLAWED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_EVACUATED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_SECURED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_RETREATED_UNIT:
            case net::eagle0::shardok::storage::fb::UnitStatus_RESERVED_SLOT:
            case net::eagle0::shardok::storage::fb::UnitStatus_PENDING_REINFORCEMENT: break;

            case net::eagle0::shardok::storage::fb::UnitStatus_UNKNOWN_UNIT:
                throw ShardokInternalErrorException("Unknown unit status");
        }
    }

    double defenderAdvantage = 1.0 + static_cast<double>(cachedCurrentRound) / 31.0;

    // Can we cache this somehow, it won't usually change within your turn
    auto attackLocationsForAttacker = GetAlCache()->CachedLocations(defenderUnits, isLateGame);
    const auto &locationsCausingDanger = attackLocationsForAttacker.AllLocations();

    // Process attacker units using cached ActionPointDistances
    for (const Unit *unit : attackerUnits) {
        const int battTypeId = unit->battalion().type();
        // Cache battalion type reference to avoid shared_ptr atomic operations
        const auto &battalionType = GetBattalionType(static_cast<BattalionTypeId>(battTypeId));

        // Cache APD lookups - same battalion type is used multiple times below
        const auto *notBravingApd =
                GetApdCache()->GetRaw(cachedHexMap, mapId, battalionType, false);
        const auto *bravingApd = battalionType->allowsBraveWater ? GetApdCache()->GetRaw(
                                                                           cachedHexMap,
                                                                           mapId,
                                                                           battalionType,
                                                                           true,
                                                                           braveWaterCost)
                                                                 : nullptr;

        const auto &priorityList = std::ranges::find_if(
                attackerTargetPriorities,
                [&unit](const TargetPriorityList &tpl) {
                    return tpl.attackingUnitId == unit->unit_id();
                });

        // If there are any tiles being targeted, give this unit a multiplier based on how close
        // they are to being able to attack it
        double distanceMultiplier = priorityList == end(attackerTargetPriorities)
                                            ? 1.0
                                            : AttackerMultiplierForTargetDistance(
                                                      unit,
                                                      priorityList->priorityOrder,
                                                      occupants,
                                                      cachedHexMap,
                                                      battalionType,
                                                      notBravingApd,
                                                      bravingApd,
                                                      isLateGame);

        auto uv = UnitValue(
                unit,
                true,
                attackerUnits,
                attackerWantsCastles,
                /* includeCastleBonus=*/true,
                defenderUnits,
                cachedHexMap,
                roundsRemaining,
                attackLocationsForAttacker,
                locationsCausingDanger,
                notBravingApd,
                GetMeteorRange(),
                GetMeteorCastVigorCost(),
                GetArcheryActionPointCost());

        attackerUnitsValue += distanceMultiplier * uv;
    }

    auto attackLocationsForDefender = GetAlCache()->CachedLocations(attackerUnits, isLateGame);
    const auto &locationsCausingDangerForAttacker = attackLocationsForDefender.AllLocations();

    for (const Unit *unit : defenderUnits) {
        auto defenderUnitId = unit->unit_id();
        const int battTypeId = unit->battalion().type();
        // Cache battalion type reference to avoid shared_ptr atomic operations
        const auto &battalionType = GetBattalionType(static_cast<BattalionTypeId>(battTypeId));

        // Cache APD lookups for this defender unit
        const auto *defenderNotBravingApd =
                GetApdCache()->GetRaw(cachedHexMap, mapId, battalionType, false);

        auto dv = UnitValue(
                unit,
                false,
                attackerUnits,
                attackerWantsCastles,
                /* includeCastleBonus=*/!defenderShouldScatter,
                defenderUnits,
                cachedHexMap,
                roundsRemaining,
                attackLocationsForDefender,
                locationsCausingDangerForAttacker,
                defenderNotBravingApd,
                GetMeteorRange(),
                GetMeteorCastVigorCost(),
                GetArcheryActionPointCost());

        double distanceMultiplier = 1.0;

        // If the defender is trying to scatter, than we want to be as far away from the nearest
        // attacker as possible, AND as far away from the nearest friendly as possible
        if (unit->location().row() > -1 && defenderShouldScatter) {
            CoordsSet myLocationSet(cachedHexMap);
            myLocationSet.Add(unit->location());

            DIST_T closestDistanceToEnemy = 999;
            for (const auto &attackerUnit : attackerUnits) {
                const int attackerBattTypeId = attackerUnit->battalion().type();
                // Cache attacker battalion type reference in nested loop
                const auto &attackerBattalionType =
                        GetBattalionType(static_cast<BattalionTypeId>(attackerBattTypeId));
                const DIST_T thisDistance = distanceCache.GetOrCompute(
                        attackerUnit,
                        unit->location(),
                        GetApdCache()->GetRaw(cachedHexMap, mapId, attackerBattalionType, false),
                        attackerBattalionType->allowsBraveWater ? GetApdCache()->GetRaw(
                                                                          cachedHexMap,
                                                                          mapId,
                                                                          attackerBattalionType,
                                                                          true,
                                                                          braveWaterCost)
                                                                : nullptr,
                        cachedHexMap);
                if (thisDistance < closestDistanceToEnemy) {
                    closestDistanceToEnemy = thisDistance;
                }
            }

            // If the best we can do puts us very close to the enemy, and the unit is almost
            // destroyed, return a negative value; better to flee
            if (unit->can_flee() && closestDistanceToEnemy < 5 && unit->battalion().size() < 10) {
                distanceMultiplier = -1;
            } else {
                DIST_T closestDistanceToFriendly = 1;
                if (defenderUnits.size() > 1) {
                    for (const auto &defenderUnit : defenderUnits) {
                        if (defenderUnit->unit_id() != defenderUnitId) {
                            const int defenderBattTypeId = defenderUnit->battalion().type();
                            // Cache defender battalion type reference in nested loop
                            const auto &defenderBattalionType = GetBattalionType(
                                    static_cast<BattalionTypeId>(defenderBattTypeId));
                            const DIST_T thisDistance = distanceCache.GetOrCompute(
                                    defenderUnit,
                                    unit->location(),
                                    GetApdCache()->GetRaw(
                                            cachedHexMap,
                                            mapId,
                                            defenderBattalionType,
                                            false),
                                    defenderBattalionType->allowsBraveWater
                                            ? GetApdCache()->GetRaw(
                                                      cachedHexMap,
                                                      mapId,
                                                      defenderBattalionType,
                                                      true,
                                                      braveWaterCost)
                                            : nullptr,
                                    cachedHexMap);
                            if (thisDistance < closestDistanceToFriendly) {
                                closestDistanceToFriendly = thisDistance;
                            }
                        }
                    }
                }

                distanceMultiplier =
                        (closestDistanceToEnemy + closestDistanceToFriendly / 5.0) / 5.0;
            }
        }

        defenderUnitsValue += distanceMultiplier * dv;
    }

    defenderUnitsValue *= defenderAdvantage;

    return UnitsScoreComponents{attackerUnitsValue, defenderUnitsValue};
}

auto AbstractAIScoreCalculator::FindDefenderPlayerInfo(const GameStateW &gameState) const
        -> const net::eagle0::shardok::storage::fb::PlayerInfo * {
    const auto *playerInfos = gameState->player_infos();
    if (playerInfos == nullptr) { return nullptr; }

    for (const net::eagle0::shardok::storage::fb::PlayerInfo *pi : *playerInfos) {
        if (pi->is_defender()) return pi;
    }
    return nullptr;
}

auto AbstractAIScoreCalculator::CalculateAttackerVictoryConditionScore(
        const GameStateW &gameState,
        const AIStrategy &attackerStrategy,
        const CoordsSet &castleCoords) const -> ScoreValue {
    ScoreValue victoryConditionTotal = 0.0;

    const auto *playerInfos = gameState->player_infos();
    if (playerInfos == nullptr) { return 0.0; }

    for (const net::eagle0::shardok::storage::fb::PlayerInfo *pi : *playerInfos) {
        if (pi->is_defender()) continue;

        switch (attackerStrategy.strategyType) {
            case AIStrategy::STRATEGY_CROSS_RIVERS:
                victoryConditionTotal += WaterCrossingScore(
                        pi->player_id(),
                        [this](BattalionTypeId typeId) { return GetBattalionType(typeId); },
                        gameState,
                        castleCoords,
                        attackerStrategy.targetLocations,
                        GetApdCache());
                break;

            case AIStrategy::STRATEGY_ATTACK_CASTLES:
            case AIStrategy::STRATEGY_ATTACK_UNITS:
                // already factored into AttackerUnitsScore
                break;

            case AIStrategy::STRATEGY_HOLD_CASTLES:
                victoryConditionTotal += AttackerHoldsCriticalTilesVictoryScore(
                        gameState,
                        castleCoords,
                        pi,
                        GetApdCache(),
                        GetAlCache(),
                        [this](BattalionTypeId typeId) { return GetBattalionType(typeId); },
                        GetBraveWaterCost());
                break;

            case AIStrategy::STRATEGY_SCATTER:
                throw ShardokInternalErrorException("Attacker cannot use ScatterStrategy");

            case AIStrategy::STRATEGY_FLEE:
                // FLEE strategy is handled specially by each subclass
                // Return 0 here and let the caller handle it
                return 0.0;
        }
    }

    return victoryConditionTotal;
}

auto AbstractAIScoreCalculator::DefenderScoreForState(
        const GameStateW &gameState,
        const AIStrategy &defenderStrategy,
        const CoordsSet &castleCoords,
        const int roundsRemaining) const -> ScoreValue {
    // Check for terminal states
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_VICTORY) {
        const auto *winningIds = gameState->status()->winning_shardok_ids();
        const auto *playerInfos = gameState->player_infos();
        if (winningIds != nullptr && playerInfos != nullptr) {
            for (const PlayerId winningPid : *winningIds) {
                if (winningPid < 0) continue;
                if (defenderStrategy.strategyType == AIStrategy::STRATEGY_FLEE) {
                    return InterpretDefenderOutcome(GameOutcome::FLEE_OUTCOME);
                }
                if (playerInfos->Get(winningPid)->is_defender()) {
                    return InterpretDefenderOutcome(GameOutcome::DEFENDER_VICTORY);
                }
                return InterpretDefenderOutcome(GameOutcome::ATTACKER_VICTORY);
            }
        }
        return InterpretDefenderOutcome(GameOutcome::ATTACKER_VICTORY);
    }
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_DRAW) {
        return InterpretDefenderOutcome(GameOutcome::DRAW);
    }

    // Delegate to strategy-specific methods (implemented by subclasses)
    switch (defenderStrategy.strategyType) {
        case AIStrategy::STRATEGY_ATTACK_CASTLES:
            throw ShardokInternalErrorException("Defender cannot use AttackCastlesStrategy");
        case AIStrategy::STRATEGY_ATTACK_UNITS:
            throw ShardokInternalErrorException("Defender cannot use AttackUnitsStrategy");
        case AIStrategy::STRATEGY_CROSS_RIVERS:
            throw ShardokInternalErrorException("Defender cannot use CrossRiversStrategy");
        case AIStrategy::STRATEGY_HOLD_CASTLES:
            return DefenderHoldCastlesStrategyScoreForState(
                    gameState,
                    castleCoords,
                    roundsRemaining);
        case AIStrategy::STRATEGY_SCATTER:
            return DefenderScatterStrategyScoreForState(gameState, roundsRemaining);
        case AIStrategy::STRATEGY_FLEE: return DefenderFleeStrategyScoreForState(gameState);
    }
    throw ShardokInternalErrorException("Escaped AIStrategy switch");
}

auto AbstractAIScoreCalculator::DefenderScatterStrategyScoreForState(
        const GameStateW &gameState,
        const int roundsRemaining) const -> ScoreValue {
    // Check for terminal states
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_VICTORY) {
        const auto *winningIds = gameState->status()->winning_shardok_ids();
        const auto *playerInfos = gameState->player_infos();
        if (winningIds != nullptr && playerInfos != nullptr) {
            for (const PlayerId winningPid : *winningIds) {
                if (winningPid < 0) continue;
                if (playerInfos->Get(winningPid)->is_defender()) {
                    return InterpretDefenderOutcome(GameOutcome::DEFENDER_VICTORY);
                }
                return InterpretDefenderOutcome(GameOutcome::ATTACKER_VICTORY);
            }
        }
        return InterpretDefenderOutcome(GameOutcome::DEFENDER_VICTORY);
    }
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_DRAW) {
        return InterpretDefenderOutcome(GameOutcome::DRAW);
    }

    // Get units score components
    const auto mapId = ActionPointDistancesCache::GetMapId(gameState->hex_map());
    const auto components = CalculateUnitsScoreComponents(
            gameState,
            roundsRemaining,
            /* attackerWantsCastles=*/false,
            /* defenderShouldScatter=*/true,
            {},
            mapId);

    // Combine using subclass-specific logic
    return CombineDefenderScatterScores(components);
}

auto AbstractAIScoreCalculator::DefenderHoldCastlesStrategyScoreForState(
        const GameStateW &gameState,
        const CoordsSet &castleCoords,
        const int roundsRemaining) const -> ScoreValue {
    // Get units score components
    const auto components = CalculateUnitsScoreComponents(
            gameState,
            roundsRemaining,
            /* attackerWantsCastles=*/true,
            /* defenderShouldScatter=*/false,
            {},
            ActionPointDistancesCache::GetMapId(gameState->hex_map()));

    // Get victory condition score from defender's perspective
    const PlayerInfo *defenderPi = FindDefenderPlayerInfo(gameState);

    // Handle null defenderPi gracefully
    if (defenderPi == nullptr) {
        // No defender player found - return neutral score using components only
        return CombineDefenderHoldCastlesScores(components, 0.0, roundsRemaining);
    }

    const double victoryConditionScore =
            DefenderHoldsCriticalTilesVictoryScore(gameState, castleCoords, defenderPi);

    // Combine using subclass-specific logic
    return CombineDefenderHoldCastlesScores(components, victoryConditionScore, roundsRemaining);
}

auto AbstractAIScoreCalculator::AttackerScoreForState(
        const GameStateW &gameState,
        const AIStrategy &attackerStrategy,
        const CoordsSet &castleCoords,
        const int roundsRemaining) const -> ScoreValue {
    // Check for terminal states
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_VICTORY) {
        const auto *winningIds = gameState->status()->winning_shardok_ids();
        const auto *playerInfos = gameState->player_infos();
        if (winningIds != nullptr && playerInfos != nullptr) {
            for (const PlayerId winningPid : *winningIds) {
                if (winningPid < 0) continue;
                if (attackerStrategy.strategyType == AIStrategy::STRATEGY_FLEE) {
                    return InterpretAttackerOutcome(GameOutcome::FLEE_OUTCOME);
                }
                if (playerInfos->Get(winningPid)->is_defender()) {
                    return InterpretAttackerOutcome(GameOutcome::DEFENDER_VICTORY);
                }
                return InterpretAttackerOutcome(GameOutcome::ATTACKER_VICTORY);
            }
        }
        return InterpretAttackerOutcome(GameOutcome::ATTACKER_VICTORY);
    }
    if (gameState->status()->state() ==
        net::eagle0::shardok::storage::fb::GameStatus_::State_DRAW) {
        return InterpretAttackerOutcome(GameOutcome::DRAW);
    }

    // Handle FLEE strategy specially
    if (attackerStrategy.strategyType == AIStrategy::STRATEGY_FLEE) {
        return AttackerFleeStrategyScoreForState(gameState);
    }

    // Edge case: no units means call subclass method to handle neutral state
    // This is defensive - some subclasses may want special handling
    const auto *units = gameState->units();
    if (units == nullptr || units->size() == 0) {
        // Call CombineAttackerScores with all zeros
        return CombineAttackerScores(UnitsScoreComponents{0.0, 0.0}, 0.0, roundsRemaining);
    }

    // Get units score components
    const auto mapId = ActionPointDistancesCache::GetMapId(gameState->hex_map());
    const auto components = CalculateUnitsScoreComponents(
            gameState,
            roundsRemaining,
            attackerStrategy.strategyType == AIStrategy::STRATEGY_HOLD_CASTLES,
            /* defenderShouldScatter=*/false,
            attackerStrategy.targetPriorities,
            mapId);

    // Calculate victory condition score
    const ScoreValue victoryConditionTotal =
            CalculateAttackerVictoryConditionScore(gameState, attackerStrategy, castleCoords);

    // Combine scores using subclass-specific logic.
    return CombineAttackerScores(components, victoryConditionTotal, roundsRemaining);
}

}  // namespace shardok
