//
// Created by Dan Crosby on 12/5/21.
//

#include "src/main/cpp/net/eagle0/shardok/library/view_filters/GameStateGuesser.hpp"

#include <cstdint>
#include <optional>
#include <vector>

#include "src/main/cpp/net/eagle0/shardok/library/fb_helpers/GameStateHelpers.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/fb_helpers/HexMapHelpers.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/map/Coordinates.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexCubeUtils.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexMapUtils.hpp"

namespace shardok {

using flatbuffers::FlatBufferBuilder;
using flatbuffers::Offset;

using BattalionViewProto = net::eagle0::shardok::api::BattalionView;
using HeroViewProto = net::eagle0::shardok::api::HeroView;
using PlayerInfoProto = net::eagle0::shardok::common::PlayerInfo;
using UnitViewProto = net::eagle0::shardok::api::UnitView;

using Coords = net::eagle0::shardok::storage::fb::Coords;
using Unit = net::eagle0::shardok::storage::fb::Unit;
using GameState = net::eagle0::shardok::storage::fb::GameState;

// Helper to get a mutable unit from the units vector.
// const_cast is safe because we're accessing through a mutable GameState pointer.
inline auto GetMutableUnit(GameState *state, UnitId unitId) -> Unit * {
    return const_cast<Unit *>(state->mutable_units()->GetMutableObject(unitId));
}

constexpr int8_t kGuessedHeroStat = 75;
constexpr int8_t kGuessedBattalionStat = 0;
constexpr int8_t kGuessedMorale = 50;

// Helper to check if coords are valid within the hex map bounds
auto CoordsAreValidProto(const HexMapProto &hexMap, const Coords &coords) -> bool {
    return coords.row() >= 0 && coords.row() < hexMap.row_count() && coords.column() >= 0 &&
           coords.column() < hexMap.column_count();
}

// Guess where an enemy mage might be targeting with their meteor.
// Priority: (1) largest unit of targetPlayerId, (2) any unit of targetPlayerId,
// (3) any castle not occupied by the caster's player, (4) random valid tile in range.
auto GuessMeteorTarget(
        const SettingsGetter &settings,
        const HexMapProto &hexMap,
        const Coords &mageLocation,
        PlayerId targetPlayerId,
        PlayerId casterPlayerId,
        const google::protobuf::RepeatedPtrField<UnitViewProto> &units) -> std::optional<Coords> {
    const int meteorRange = settings.Backing().meteor_range();
    const Cube mageCube = OffsetToCube(mageLocation);

    // Build a set of valid coords within meteor range
    std::vector<Coords> validCoordsInRange;
    validCoordsInRange.reserve(1 + 3 * meteorRange * (meteorRange + 1));
    for (const Cube &cube : CubesWithinDistance(mageCube, meteorRange)) {
        Coords coords = CubeToOffset(cube);
        if (CoordsAreValidProto(hexMap, coords) && coords != mageLocation) {
            validCoordsInRange.push_back(coords);
        }
    }

    if (validCoordsInRange.empty()) { return std::nullopt; }

    // Priority 1 & 2: Find the largest unit of the target player within range
    const UnitViewProto *largestTargetUnit = nullptr;
    int largestSize = 0;

    for (const auto &unit : units) {
        if (unit.player_id() != targetPlayerId) continue;

        Coords unitCoords = FromCoordsProto(unit.location());
        if (IsUnplaced(unitCoords)) continue;

        // Check if unit is within meteor range
        const int distance = CubeDistance(mageCube, OffsetToCube(unitCoords));
        if (distance > meteorRange) continue;

        const int size = unit.battalion().size();
        if (size > largestSize) {
            largestSize = size;
            largestTargetUnit = &unit;
        }
    }

    if (largestTargetUnit != nullptr) { return FromCoordsProto(largestTargetUnit->location()); }

    // Priority 3: Find a castle not occupied by the caster's player
    // Build set of occupied coords by caster's player
    std::unordered_set<int> casterOccupiedIndices;
    casterOccupiedIndices.reserve(units.size());
    for (const auto &unit : units) {
        if (unit.player_id() == casterPlayerId) {
            Coords loc = FromCoordsProto(unit.location());
            if (!IsUnplaced(loc)) {
                casterOccupiedIndices.insert(loc.row() * hexMap.column_count() + loc.column());
            }
        }
    }

    for (const auto &coords : validCoordsInRange) {
        const int index = coords.row() * hexMap.column_count() + coords.column();
        if (index < 0 || index >= hexMap.terrain_size()) continue;

        const auto &terrain = hexMap.terrain(index);
        if (terrain.modifier().has_castle() &&
            casterOccupiedIndices.find(index) == casterOccupiedIndices.end()) {
            return coords;
        }
    }

    // Priority 4: Return the first valid coord in range (deterministic "random")
    return validCoordsInRange[0];
}

auto GuessedUnit(
        const SettingsGetter &settings,
        const UnitViewProto &uv,
        PlayerId toPid,
        Coords location,
        net::eagle0::shardok::storage::fb::UnitStatus status,
        const std::unordered_map<PlayerId, AverageStats> &playerAverages) -> Unit;

auto GuessedHero(const HeroViewProto &hv) -> net::eagle0::shardok::storage::fb::Hero;

auto GuessedBattalion(
        const BattalionViewProto &bv,
        PlayerId playerId,
        const std::unordered_map<PlayerId, AverageStats> &playerAverages)
        -> net::eagle0::shardok::storage::fb::Battalion;

auto GameStateGuesser::GuessedState(
        PlayerId pid,
        const SettingsGetter &settings,
        const GameStateView &gameStateView) -> GameStateW {
    // Get known unit stats up front to use for guessing
    auto playerAverages = KnownUnitStats(gameStateView);

    bool isDefender = false;
    bool defenderExists = false;
    for (const auto &pip : gameStateView.player_infos()) {
        defenderExists = defenderExists || pip.is_defender();
        isDefender = isDefender || (pip.is_defender() && pip.player_id() == pid);
    }

    FlatBufferBuilder fbb;
    fbb.ForceDefaults(true);

    auto gameIdOffset = fbb.CreateString(gameStateView.game_id());

    std::vector<Offset<net::eagle0::shardok::storage::fb::PlayerInfo>> piVec{};
    piVec.reserve(gameStateView.player_infos_size());
    for (const PlayerInfoProto &pi : gameStateView.player_infos()) {
        piVec.push_back(fb::FromPlayerInfoProto(fbb, pi));
    }
    auto piOffset = fbb.CreateVectorOfSortedTables(&piVec);

    auto statusDescription = fbb.CreateString(gameStateView.status().description());

    std::vector<PlayerId> winningShardokIdsVec;
    winningShardokIdsVec.reserve(10);
    winningShardokIdsVec.insert(
            winningShardokIdsVec.end(),
            std::begin(gameStateView.status().winning_shardok_ids()),
            std::end(gameStateView.status().winning_shardok_ids()));
    for (int i = gameStateView.status().winning_shardok_ids_size(); i < 10; ++i) {
        winningShardokIdsVec.push_back(-1);
    }
    auto winningShardokIdsOffset = fbb.CreateVector(winningShardokIdsVec);

    net::eagle0::shardok::storage::fb::GameStatusBuilder statusBuilder(fbb);
    statusBuilder.add_description(statusDescription);
    statusBuilder.add_state(static_cast<net::eagle0::shardok::storage::fb::GameStatus_::State>(
            gameStateView.status().state()));

    net::eagle0::shardok::storage::fb::EndGameCondition endGameCondition;
    if (gameStateView.status().has_end_game_condition()) {
        switch (gameStateView.status().end_game_condition().condition_case()) {
            case net::eagle0::shardok::common::EndGameCondition::ConditionCase::kVictory:
                endGameCondition.mutate_victory_type(
                        net::eagle0::shardok::storage::fb::VictoryType_VICTORY);
                endGameCondition.mutate_victory_details(
                        static_cast<net::eagle0::shardok::storage::fb::VictoryCondition>(
                                gameStateView.status().end_game_condition().victory()));
                break;

            case net::eagle0::shardok::common::EndGameCondition::ConditionCase::kLoss:
                endGameCondition.mutate_victory_type(
                        net::eagle0::shardok::storage::fb::VictoryType_LOSS);
                endGameCondition.mutate_victory_details(
                        static_cast<net::eagle0::shardok::storage::fb::VictoryCondition>(
                                gameStateView.status().end_game_condition().loss()));
                break;

            case net::eagle0::shardok::common::EndGameCondition::ConditionCase::CONDITION_NOT_SET:
                break;
        }
    }

    statusBuilder.add_end_game_condition(&endGameCondition);
    statusBuilder.add_winning_shardok_ids(winningShardokIdsOffset);
    auto statusOffset = statusBuilder.Finish();

    auto mapOffset = fb::ConvertHexMapProto(fbb, gameStateView.hex_map());

    // units
    bool foundDefender = false;
    UnitId maxId = 0;

    std::vector<UnitId> foundUnits{};
    std::vector<Unit> unitsVec{};
    const int knownUnitCount = gameStateView.units_size() + gameStateView.reserve_units_size() +
                               gameStateView.captured_units_size();
    foundUnits.reserve(knownUnitCount);
    unitsVec.reserve(knownUnitCount + 6);
    for (const auto &uv : gameStateView.units()) {
        unitsVec.push_back(GuessedUnit(
                settings,
                uv,
                pid,
                FromCoordsProto(uv.location()),
                net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT,
                playerAverages));
        if (uv.player_id() == 1) foundDefender = true;
        if (uv.unit_id() > maxId) maxId = uv.unit_id();
        foundUnits.push_back(uv.unit_id());
    }

    for (const auto &uv : gameStateView.reserve_units()) {
        Coords location = FromCoordsProto(uv.location());
        auto status = net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT;

        // If the player is not the defender and there are no known defenders, guess one
        if (!isDefender && defenderExists) {
            if (!foundDefender && uv.player_id() != pid) {
                foundDefender = true;
                status = net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT;
                for (const auto &startPosProto :
                     gameStateView.hex_map().defender_starting_positions().positions()) {
                    Coords startPos = FromCoordsProto(startPosProto);
                    if (!Occupant<UnitViewProto>(gameStateView.units(), startPos).has_value()) {
                        location = startPos;
                        break;
                    }
                }
            }
        }
        unitsVec.push_back(GuessedUnit(settings, uv, pid, location, status, playerAverages));

        if (uv.unit_id() > maxId) maxId = uv.unit_id();
        foundUnits.push_back(uv.unit_id());
    }

    for (const auto &uv : gameStateView.captured_units()) {
        unitsVec.push_back(GuessedUnit(
                settings,
                uv,
                pid,
                FromCoordsProto(uv.location()),
                net::eagle0::shardok::storage::fb::UnitStatus_CAPTURED_UNIT,
                playerAverages));
        if (uv.player_id() == 1) foundDefender = true;
        if (uv.unit_id() > maxId) maxId = uv.unit_id();
        foundUnits.push_back(uv.unit_id());
    }

    for (int i = 0; i < maxId; ++i) {
        auto entry = std::ranges::find(foundUnits, i);
        if (entry == std::end(foundUnits)) {
            net::eagle0::shardok::storage::fb::Unit newUnit;
            newUnit.mutate_unit_id(i);
            newUnit.mutate_has_attached_hero(false);
            newUnit.mutate_commanding_unit_id(-1);
            newUnit.mutate_player_id(0);
            newUnit.mutable_battalion().mutate_type(
                    net::eagle0::shardok::storage::fb::BattalionTypeId_LIGHT_INFANTRY);
            newUnit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_CAPTURED_UNIT);
            newUnit.mutable_location() = Coords(-1, -1);

            unitsVec.push_back(newUnit);
        }
    }

    // If we're the attacker and we don't know of defenders, we're either in SETUP or defenders are
    // all hidden
    if (!isDefender && !foundDefender && defenderExists) {
        Unit newUnit;

        newUnit.mutable_battalion().mutate_size(1000);

        newUnit.mutate_unit_id(maxId + 1);
        newUnit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT);
        newUnit.mutate_has_attached_hero(false);
        newUnit.mutate_commanding_unit_id(-1);

        for (const auto &startPosProto :
             gameStateView.hex_map().defender_starting_positions().positions()) {
            Coords startPos = FromCoordsProto(startPosProto);
            if (!Occupant<UnitViewProto>(gameStateView.units(), startPos).has_value()) {
                newUnit.mutable_location() = startPos;
                newUnit.mutate_hidden(true);
                newUnit.mutate_player_id(1);
                break;
            }
        }

        unitsVec.push_back(newUnit);
    }

    // Add 5 reserved slots for future expansion (e.g., necromancer raise dead)
    UnitId nextReservedId = maxId + 2;  // Start after any guessed defender unit
    for (int i = 0; i < 5; ++i) {
        Unit reservedSlot;
        reservedSlot.mutate_unit_id(nextReservedId + i);
        reservedSlot.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_RESERVED_SLOT);
        // Set other required fields to safe defaults
        reservedSlot.mutate_player_id(-1);
        reservedSlot.mutate_eagle_player_id(-1);
        reservedSlot.mutable_location().mutate_row(-1);
        reservedSlot.mutable_location().mutate_column(-1);
        reservedSlot.mutate_has_attached_hero(false);
        reservedSlot.mutate_commanding_unit_id(-1);
        unitsVec.push_back(reservedSlot);
    }

    // Post-process: guess meteor targets for enemy mages who are casting but whose target is
    // unknown
    for (auto &unit : unitsVec) {
        if (!unit.has_attached_hero()) continue;
        if (unit.player_id() == pid) continue;  // Only guess for enemy mages

        const auto castState = unit.attached_hero().profession_info().meteor_cast_state();
        // Mage has targeted or is ready to cast, but we don't know the target
        if (castState != net::eagle0::shardok::storage::fb::MultiroundMagicState_TARGET &&
            castState != net::eagle0::shardok::storage::fb::MultiroundMagicState_CAST) {
            continue;
        }

        const auto &castTarget = unit.attached_hero().profession_info().cast_target();
        if (castTarget.row() >= 0 && castTarget.column() >= 0) {
            continue;  // Already has a valid target
        }

        // Guess a target: assume mage is targeting us (the viewing player)
        auto guessedTarget = GuessMeteorTarget(
                settings,
                gameStateView.hex_map(),
                unit.location(),
                pid,
                unit.player_id(),
                gameStateView.units());

        if (guessedTarget.has_value()) {
            unit.mutable_attached_hero().mutable_profession_info().mutable_cast_target().mutate_row(
                    guessedTarget->row());
            unit.mutable_attached_hero()
                    .mutable_profession_info()
                    .mutable_cast_target()
                    .mutate_column(guessedTarget->column());
        }
    }

    auto unitsOffset = fbb.CreateVectorOfSortedStructs(&unitsVec);

    std::vector<UnitId> chargeeIdsVec;
    chargeeIdsVec.reserve(6);
    chargeeIdsVec.insert(
            chargeeIdsVec.end(),
            std::begin(gameStateView.possible_chargee_ids()),
            std::end(gameStateView.possible_chargee_ids()));
    for (int i = gameStateView.possible_chargee_ids_size(); i < 6; ++i) {
        chargeeIdsVec.push_back(-1);
    }
    auto chargeeIdsOffset = fbb.CreateVector(chargeeIdsVec);

    net::eagle0::shardok::storage::fb::GameStateBuilder gs(fbb);

    // constants
    gs.add_game_id(gameIdOffset);
    gs.add_month(gameStateView.month());

    // these change, but they're well-known
    gs.add_current_round(gameStateView.current_round());
    gs.add_current_player(gameStateView.current_player());

    gs.add_player_infos(piOffset);

    net::eagle0::shardok::storage::fb::Wind wind(
            static_cast<net::eagle0::shardok::storage::fb::HexMapDirection>(
                    gameStateView.weather().wind().direction()),
            gameStateView.weather().wind().speed_in_mph());

    net::eagle0::shardok::storage::fb::Weather weather(
            static_cast<net::eagle0::shardok::storage::fb::WeatherConditions>(
                    gameStateView.weather().conditions()),
            wind);
    gs.add_weather(&weather);

    gs.add_status(statusOffset);

    // partially known
    gs.add_hex_map(mapOffset);

    gs.add_units(unitsOffset);

    if (gameStateView.has_eligible_charger_id()) {
        gs.add_eligible_charger_id(gameStateView.eligible_charger_id().value());
    } else {
        gs.add_eligible_charger_id(-1);
    }

    gs.add_possible_chargee_ids(chargeeIdsOffset);

    fbb.Finish(gs.Finish());

    GameStateW gsw(fbb);
    // Make sure that controlled unit/commanding unit IDs match up
    for (UnitId unitId = 0; static_cast<unsigned int>(unitId) < gsw->units()->size(); ++unitId) {
        const Unit *unit = gsw->units()->Get(unitId);

        if (unit->has_attached_hero()) {
            UnitId controlledUnitId = unit->attached_hero().control_info().controlled_unit_id();
            if (controlledUnitId != -1) {
                GetMutableUnit(gsw.Get(), controlledUnitId)->mutate_commanding_unit_id(unitId);
            }
        }

        UnitId commandingUnitId = unit->commanding_unit_id();
        if (unit->commanding_unit_id() != -1) {
            GetMutableUnit(gsw.Get(), commandingUnitId)
                    ->mutable_attached_hero()
                    .mutable_control_info()
                    .mutate_controlled_unit_id(unitId);
        }
    }

    return gsw;
}

auto GuessedUnit(
        const SettingsGetter & /*settings*/,
        const UnitViewProto &uv,
        PlayerId toPid,
        Coords location,
        net::eagle0::shardok::storage::fb::UnitStatus status,
        const std::unordered_map<PlayerId, AverageStats> &playerAverages) -> Unit {
    net::eagle0::shardok::storage::fb::Hero attachedHero{};
    if (uv.has_attached_hero()) { attachedHero = GuessedHero(uv.attached_hero()); }

    fb::Unit unit;

    unit.mutate_unit_id(uv.unit_id());
    unit.mutate_player_id(uv.player_id());
    unit.mutate_hidden(uv.hidden());
    if (uv.has_attached_hero()) {
        unit.mutable_attached_hero() = attachedHero;
        unit.mutate_has_attached_hero(true);
    } else {
        unit.mutate_has_attached_hero(false);
    }
    if (uv.has_commanding_unit_id()) {
        unit.mutate_commanding_unit_id(uv.commanding_unit_id().value());
    } else {
        unit.mutate_commanding_unit_id(-1);
    }

    unit.mutable_location() = location;

    // We might have an off-map location for a normal unit if it's hidden. We'll just pretend
    // it's a reserve unit for now. Ideally we'd make a guess as to where it's hidden.
    if (location.row() < 0 && status == net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT) {
        // I *think* hidden units are the only way this can happen, let's check.
        assert(uv.hidden());
        unit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT);
    } else {
        unit.mutate_status(status);
    }

    unit.mutable_battalion() = GuessedBattalion(uv.battalion(), uv.player_id(), playerAverages);

    if (uv.has_remaining_action_points()) {
        unit.mutate_remaining_action_points(uv.remaining_action_points().value());
    } else {
        unit.mutate_remaining_action_points(0);
    }
    if (uv.has_stun_rounds_remaining()) {
        unit.mutate_stun_rounds_remaining(uv.stun_rounds_remaining().value());
    } else {
        unit.mutate_stun_rounds_remaining(0);
    }
    // Volleys: use real value if known, otherwise conservative placeholder
    if (uv.has_volleys_remaining()) {
        unit.mutate_volleys_remaining(uv.volleys_remaining().value());
    } else {
        // Conservative estimate - assume they have volleys if we don't know
        // This prevents AI from underestimating archery threat
        unit.mutate_volleys_remaining(5);
    }

    unit.mutate_has_moved_in_zoc(uv.has_moved_in_zoc());

    for (uint32_t i = 0; i < unit.opponent_knowledge()->size(); ++i) {
        unit.mutable_opponent_knowledge()->Mutate(i, 0);
    }
    unit.mutable_opponent_knowledge()->Mutate(toPid, uv.my_knowledge());

    unit.mutate_can_flee(uv.can_flee());

    // For can_archery and can_start_fire:
    // - If we have battalion stats (armament/training known), trust the proto values
    // - Otherwise, be conservative and assume true (opponent might be able to)
    // - Exception: Units without heroes can't archery or start fires
    const bool haveStats = uv.battalion().has_armament() && uv.battalion().has_training();
    const bool hasHero = uv.has_attached_hero();

    if (haveStats) {
        // We know the unit's capabilities, use the actual values
        unit.mutate_can_archery(uv.can_archery());
        unit.mutate_can_start_fire(uv.can_start_fire());
    } else {
        // We don't know - be conservative (assume they can, so AI doesn't underestimate threat)
        // But units need heroes for archery and fire
        unit.mutate_can_archery(hasHero);
        unit.mutate_can_start_fire(hasHero);
    }

    unit.mutate_fortified(uv.fortified());
    unit.mutate_targeted_unit(-1);

    // Food: use real value if known (when we know battalion stats), otherwise guess high to avoid
    // starvation
    if (uv.has_food_remaining()) {
        unit.mutate_food_remaining(static_cast<float>(uv.food_remaining().value()));
    } else {
        // Conservative estimate - assume they have enough food
        // This prevents AI simulation from incorrectly applying starvation (size=0)
        // Typical daily food cost is ~1-2 per unit, so 100 gives plenty of buffer
        unit.mutate_food_remaining(100.0f);
    }

    if (uv.has_starting_position_index()) {
        unit.mutate_starting_position_index(uv.starting_position_index().value());
    } else {
        unit.mutate_starting_position_index(-1);
    }

    return unit;
}

auto GuessedHero(const HeroViewProto &hv) -> net::eagle0::shardok::storage::fb::Hero {
    net::eagle0::shardok::storage::fb::Hero hero;

    hero.mutate_eagle_hero_id(-1);

    if (hv.has_control_info()) {
        hero.mutable_control_info().mutate_controlled_unit_id(
                hv.control_info().controlled_unit_id());
        hero.mutable_control_info().mutate_controlled_this_round(
                hv.control_info().controlled_this_round());
    } else {
        hero.mutable_control_info().mutate_controlled_unit_id(-1);
        hero.mutable_control_info().mutate_controlled_this_round(false);
    }

    hero.mutate_is_vip(hv.has_is_vip() && hv.is_vip().value());

    if (hv.has_stats()) {
        const auto &stats = hv.stats();
        hero.mutate_strength(stats.strength());
        hero.mutate_strength_xp(kGuessedHeroStat);
        hero.mutate_agility(stats.agility());
        hero.mutate_agility_xp(kGuessedHeroStat);
        hero.mutate_constitution(stats.constitution());
        hero.mutate_constitution_xp(kGuessedHeroStat);
        hero.mutate_charisma(stats.charisma());
        hero.mutate_charisma_xp(kGuessedHeroStat);
        hero.mutate_wisdom(stats.wisdom());
        hero.mutate_wisdom_xp(kGuessedHeroStat);

        hero.mutate_integrity(stats.integrity());
        hero.mutate_ambition(stats.ambition());
        hero.mutate_gregariousness(stats.gregariousness());
        hero.mutate_bravery(stats.bravery());

        hero.mutate_vigor(stats.vigor());
        // The battle-start recovery ceiling is internal state and is not part of HeroView.
        // Use the visible current vigor as a conservative ceiling rather than inventing
        // recovery capacity that the player cannot observe.
        hero.mutate_starting_vigor(stats.vigor());
        hero.mutate_spent_vigor(kGuessedHeroStat);
    } else {
        // Don't believe it's destroyed
        hero.mutate_strength(kGuessedHeroStat);
        hero.mutate_strength_xp(kGuessedHeroStat);
        hero.mutate_agility(kGuessedHeroStat);
        hero.mutate_agility_xp(kGuessedHeroStat);
        hero.mutate_constitution(kGuessedHeroStat);
        hero.mutate_charisma(kGuessedHeroStat);
        hero.mutate_charisma_xp(kGuessedHeroStat);
        hero.mutate_wisdom(kGuessedHeroStat);
        hero.mutate_wisdom_xp(kGuessedHeroStat);

        hero.mutate_integrity(kGuessedHeroStat);
        hero.mutate_ambition(kGuessedHeroStat);
        hero.mutate_gregariousness(kGuessedHeroStat);
        hero.mutate_bravery(kGuessedHeroStat);

        hero.mutate_vigor(kGuessedHeroStat);
        hero.mutate_starting_vigor(kGuessedHeroStat);
        hero.mutate_spent_vigor(kGuessedHeroStat);
    }

    if (hv.has_profession_info()) {
        hero.mutable_profession_info().mutate_profession(
                static_cast<net::eagle0::shardok::storage::fb::Profession>(
                        hv.profession_info().profession()));
        hero.mutable_profession_info().mutate_meteor_cast_state(
                static_cast<net::eagle0::shardok::storage::fb::MultiroundMagicState>(
                        hv.profession_info().meteor_cast_state()));
        if (hv.profession_info().has_cast_target()) {
            hero.mutable_profession_info().mutable_cast_target().mutate_row(
                    hv.profession_info().cast_target().row());
            hero.mutable_profession_info().mutable_cast_target().mutate_column(
                    hv.profession_info().cast_target().column());
        } else {
            hero.mutable_profession_info().mutable_cast_target().mutate_row(-1);
            hero.mutable_profession_info().mutable_cast_target().mutate_column(-1);
        }
    } else {
        hero.mutable_profession_info().mutable_cast_target().mutate_row(-1);
        hero.mutable_profession_info().mutable_cast_target().mutate_column(-1);
        hero.mutable_profession_info().mutate_profession(
                net::eagle0::shardok::storage::fb::Profession_NO_PROFESSION);
        hero.mutable_profession_info().mutate_meteor_cast_state(
                net::eagle0::shardok::storage::fb::MultiroundMagicState_NONE);
    }

    return hero;
}

auto GuessedBattalion(
        const BattalionViewProto &bv,
        PlayerId playerId,
        const std::unordered_map<PlayerId, AverageStats> &playerAverages)
        -> net::eagle0::shardok::storage::fb::Battalion {
    // Three-tier fallback for stats:
    // 1. Use actual stats if present
    // 2. Use average stats for this player if available
    // 3. Use kGuessedStat as default

    int8_t armament = kGuessedBattalionStat;
    int8_t training = kGuessedBattalionStat;

    if (bv.has_armament()) {
        armament = bv.armament().value();
    } else if (playerAverages.contains(playerId)) {
        armament = static_cast<int8_t>(std::round(playerAverages.at(playerId).armament));
    }

    if (bv.has_training()) {
        training = bv.training().value();
    } else if (playerAverages.contains(playerId)) {
        training = static_cast<int8_t>(std::round(playerAverages.at(playerId).training));
    }

    const net::eagle0::shardok::storage::fb::Battalion batt(
            static_cast<net::eagle0::shardok::storage::fb::BattalionTypeId>(bv.type()),
            bv.size(),
            bv.has_morale() ? bv.morale().value() : kGuessedMorale,
            training,
            armament,
            -1,
            kGuessedMorale);

    return batt;
}

auto KnownUnitStats(const GameStateView &gameStateView)
        -> std::unordered_map<PlayerId, AverageStats> {
    struct PlayerStats {
        double armamentSum = 0.0;
        double trainingSum = 0.0;
        int count = 0;
    };

    std::unordered_map<PlayerId, PlayerStats> playerStatsMap;

    // Single pass to collect all stats
    for (const auto &unit : gameStateView.units()) {
        if (unit.player_id() < 0) continue;  // Skip reserve units

        if (unit.battalion().has_armament() && unit.battalion().has_training()) {
            auto &stats = playerStatsMap[unit.player_id()];
            stats.count++;
            stats.armamentSum += unit.battalion().armament().value();
            stats.trainingSum += unit.battalion().training().value();
        }
    }

    // Build result map with averages
    std::unordered_map<PlayerId, AverageStats> result;
    result.reserve(playerStatsMap.size());  // Pre-allocate for efficiency

    for (const auto &[playerId, stats] : playerStatsMap) {
        if (stats.count > 0) {
            result.emplace(
                    playerId,
                    AverageStats{
                            .armament = stats.armamentSum / stats.count,
                            .training = stats.trainingSum / stats.count});
        }
    }

    return result;
}

}  // namespace shardok
