//
// AI Battle Simulator Implementation
//

#include "AiBattleSimulator.hpp"

#include <algorithm>
#include <iostream>
#include <sstream>
#include <stdexcept>
#include <utility>

#include "src/main/cpp/net/eagle0/common/FilesystemUtils.hpp"
#include "src/main/cpp/net/eagle0/common/RandomGenerator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/ShardokAIClient.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai_testing_common/AIClientFactory.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai_testing_common/GamePhaseRunner.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai_testing_common/GameSettingsFactory.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/UnitPlacementInfo.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/fb_helpers/GameStateHelpers.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/settings/GameSettings.hpp"
#include "src/main/cpp/net/eagle0/shardok/util/MapLoader.hpp"
#include "src/main/flatbuffer/net/eagle0/shardok/storage/unit.hpp"

#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Wdeprecated-redundant-constexpr-static-def"
#include "src/main/protobuf/net/eagle0/shardok/common/player_info.pb.h"
#pragma clang diagnostic pop

namespace shardok::ai_battle_simulator {

namespace {

constexpr PlayerId ATTACKER_ID = 0;
constexpr PlayerId DEFENDER_ID = 1;

// Default values for unit configuration - matching client-initiated battles
constexpr double DEFAULT_BATTALION_ARMAMENT = 102.0;
constexpr double DEFAULT_BATTALION_TRAINING = 102.0;
constexpr double DEFAULT_BATTALION_MORALE = 76.0;

constexpr int DEFAULT_HERO_STRENGTH = 102;
constexpr int DEFAULT_HERO_AGILITY = 102;
constexpr int DEFAULT_HERO_WISDOM = 102;
constexpr int DEFAULT_HERO_CHARISMA = 102;
constexpr int DEFAULT_HERO_CONSTITUTION = 102;
constexpr int DEFAULT_HERO_BRAVERY = 102;
constexpr int DEFAULT_HERO_INTEGRITY = 0;
constexpr int DEFAULT_HERO_AMBITION = 0;
constexpr int DEFAULT_HERO_VIGOR = 102;

constexpr int PROFESSION_MAGE = 1;
constexpr double MINIMUM_AGILITY_FOR_ARCHERY = 75.0;
constexpr double MINIMUM_AGILITY_FOR_START_FIRE = 70.0;
constexpr double MINIMUM_WISDOM_FOR_START_FIRE = 70.0;

struct SideTotals {
    int attackerUnits = 0;
    int defenderUnits = 0;
    double attackerTroops = 0.0;
    double defenderTroops = 0.0;
};

std::string JsonEscape(const std::string& value) {
    std::ostringstream out;
    for (const char c : value) {
        switch (c) {
            case '"': out << "\\\""; break;
            case '\\': out << "\\\\"; break;
            case '\n': out << "\\n"; break;
            case '\r': out << "\\r"; break;
            case '\t': out << "\\t"; break;
            default: out << c; break;
        }
    }
    return out.str();
}

std::string CompletionReasonLabel(EvaluationCompletionReason reason) {
    switch (reason) {
        case EvaluationCompletionReason::RAN_OUT_OF_COMMANDS: return "ran_out_of_commands";
        case EvaluationCompletionReason::RAN_OUT_OF_TIME: return "ran_out_of_time";
        case EvaluationCompletionReason::NOT_ENOUGH_TIME_TO_CONTINUE:
            return "not_enough_time_to_continue";
    }
}

SideTotals CalculateSideTotals(const GameStateW& state) {
    SideTotals totals;
    const auto* units = state->units();
    if (units == nullptr) { return totals; }

    for (const auto* unit : *units) {
        if (unit == nullptr) { continue; }
        const bool isAlive =
                unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT ||
                unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT ||
                unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_NEVER_ENTERED_UNIT;
        if (!isAlive) { continue; }

        if (unit->player_id() == ATTACKER_ID) {
            totals.attackerUnits++;
            totals.attackerTroops += unit->battalion().size();
        } else if (unit->player_id() == DEFENDER_ID) {
            totals.defenderUnits++;
            totals.defenderTroops += unit->battalion().size();
        }
    }

    return totals;
}

// Convert proto AI algorithm type to ShardokAI enum
AIAlgorithmType ConvertAIAlgorithmType(
        net::eagle0::shardok::ai_battle_simulator::AIAlgorithmType protoType) {
    switch (protoType) {
        case net::eagle0::shardok::ai_battle_simulator::MCTS: return AIAlgorithmType::MCTS;
        case net::eagle0::shardok::ai_battle_simulator::ITERATIVE_DEEPENING:
        default: return AIAlgorithmType::ITERATIVE_DEEPENING;
    }
}

ScoringCalculatorType ConvertScoringCalculatorType(
        net::eagle0::shardok::ai_battle_simulator::ScoringCalculatorType protoType) {
    switch (protoType) {
        case net::eagle0::shardok::ai_battle_simulator::EXPERIMENTAL:
            return ScoringCalculatorType::EXPERIMENTAL;
        case net::eagle0::shardok::ai_battle_simulator::MCTS_OPTIMIZED:
            return ScoringCalculatorType::MCTS_OPTIMIZED;
        case net::eagle0::shardok::ai_battle_simulator::STANDARD:
        case net::eagle0::shardok::ai_battle_simulator::SCORING_CALCULATOR_TYPE_UNSPECIFIED:
        default: return ScoringCalculatorType::STANDARD;
    }
}

mcts::MCTSSimulationPolicy ConvertSimulationPolicy(
        net::eagle0::shardok::ai_battle_simulator::MctsSimulationPolicy protoType) {
    switch (protoType) {
        case net::eagle0::shardok::ai_battle_simulator::RANDOM:
            return mcts::MCTSSimulationPolicy::RANDOM;
        case net::eagle0::shardok::ai_battle_simulator::FILTERED_RANDOM:
            return mcts::MCTSSimulationPolicy::FILTERED_RANDOM;
        case net::eagle0::shardok::ai_battle_simulator::BEST_IMMEDIATE:
            return mcts::MCTSSimulationPolicy::BEST_IMMEDIATE;
        case net::eagle0::shardok::ai_battle_simulator::WEIGHTED_BEST_IMMEDIATE:
            return mcts::MCTSSimulationPolicy::WEIGHTED_BEST_IMMEDIATE;
        case net::eagle0::shardok::ai_battle_simulator::WEIGHTED_HEURISTIC:
        case net::eagle0::shardok::ai_battle_simulator::MCTS_SIMULATION_POLICY_UNSPECIFIED:
        default: return mcts::MCTSSimulationPolicy::WEIGHTED_HEURISTIC;
    }
}

mcts::MCTSBackpropagationPolicy ConvertBackpropagationPolicy(
        net::eagle0::shardok::ai_battle_simulator::MctsBackpropagationPolicy protoType) {
    switch (protoType) {
        case net::eagle0::shardok::ai_battle_simulator::AVERAGING:
            return mcts::MCTSBackpropagationPolicy::AVERAGING;
        case net::eagle0::shardok::ai_battle_simulator::MINIMAX:
        case net::eagle0::shardok::ai_battle_simulator::MCTS_BACKPROPAGATION_POLICY_UNSPECIFIED:
        default: return mcts::MCTSBackpropagationPolicy::MINIMAX;
    }
}

mcts::MCTSConfig BuildMCTSConfig(
        const net::eagle0::shardok::ai_battle_simulator::PlayerConfig& playerConfig,
        int battleSeed,
        PlayerId playerId) {
    mcts::MCTSConfig config{};
    config.backpropagationPolicy = mcts::MCTSBackpropagationPolicy::MINIMAX;
    config.maxPlayerFlips = 1;
    config.maxSimulationFlips = 1;
    config.simulationPolicy = mcts::MCTSSimulationPolicy::WEIGHTED_HEURISTIC;
    if (battleSeed != 0) { config.useMultithreading = false; }
    if (battleSeed != 0) {
        config.randomSeed = static_cast<uint32_t>(battleSeed + (1009 * playerId));
    }

    if (!playerConfig.has_mcts_config()) { return config; }

    const auto& protoConfig = playerConfig.mcts_config();
    if (protoConfig.exploration_constant() > 0.0) {
        config.explorationConstant = protoConfig.exploration_constant();
    }
    if (protoConfig.max_simulation_depth() > 0) {
        config.maxSimulationDepth = protoConfig.max_simulation_depth();
    }
    if (protoConfig.max_tree_depth() > 0) { config.maxTreeDepth = protoConfig.max_tree_depth(); }
    config.useMultithreading = protoConfig.use_multithreading();
    if (protoConfig.num_threads() > 0) { config.numThreads = protoConfig.num_threads(); }
    config.simulationPolicy = ConvertSimulationPolicy(protoConfig.simulation_policy());
    config.backpropagationPolicy =
            ConvertBackpropagationPolicy(protoConfig.backpropagation_policy());
    if (protoConfig.max_player_flips() >= 0) {
        config.maxPlayerFlips = protoConfig.max_player_flips();
    }
    if (protoConfig.max_simulation_flips() >= 0) {
        config.maxSimulationFlips = protoConfig.max_simulation_flips();
    }
    if (protoConfig.random_seed() != 0) {
        config.randomSeed = static_cast<uint32_t>(protoConfig.random_seed());
    }
    return config;
}

bool HasAttachedHero(int profession) { return profession > 0; }

// Get value with default if not set (proto3 uses 0 as default, we check for that)
template<typename T>
T GetOrDefault(T value, T defaultValue) {
    return (value == 0) ? defaultValue : value;
}

bool HeroCanArchery(const net::eagle0::shardok::ai_battle_simulator::UnitConfig& unitConfig) {
    if (!HasAttachedHero(unitConfig.profession())) { return false; }
    return GetOrDefault(unitConfig.hero().agility(), DEFAULT_HERO_AGILITY) >=
           MINIMUM_AGILITY_FOR_ARCHERY;
}

bool HeroCanStartFire(const net::eagle0::shardok::ai_battle_simulator::UnitConfig& unitConfig) {
    if (!HasAttachedHero(unitConfig.profession())) { return false; }
    return unitConfig.profession() == PROFESSION_MAGE ||
           (GetOrDefault(unitConfig.hero().agility(), DEFAULT_HERO_AGILITY) >=
                    MINIMUM_AGILITY_FOR_START_FIRE &&
            GetOrDefault(unitConfig.hero().wisdom(), DEFAULT_HERO_WISDOM) >=
                    MINIMUM_WISDOM_FOR_START_FIRE);
}

double GetBattalionSizeOrDefault(
        const net::eagle0::shardok::ai_battle_simulator::BattalionConfig& battalion,
        double defaultValue) {
    return battalion.has_size() ? battalion.size() : defaultValue;
}

bool GetCanFleeOrDefault(const net::eagle0::shardok::ai_battle_simulator::UnitConfig& unitConfig) {
    return unitConfig.has_can_flee() ? unitConfig.can_flee() : true;
}

double GetBattalionArmamentOrDefault(
        const net::eagle0::shardok::ai_battle_simulator::BattalionConfig& battalion) {
    return battalion.has_armament() ? battalion.armament() : DEFAULT_BATTALION_ARMAMENT;
}

double GetBattalionTrainingOrDefault(
        const net::eagle0::shardok::ai_battle_simulator::BattalionConfig& battalion) {
    return battalion.has_training() ? battalion.training() : DEFAULT_BATTALION_TRAINING;
}

double GetBattalionMoraleOrDefault(
        const net::eagle0::shardok::ai_battle_simulator::BattalionConfig& battalion) {
    return battalion.has_morale() ? battalion.morale() : DEFAULT_BATTALION_MORALE;
}

}  // namespace

AiBattleSimulator::AiBattleSimulator(
        const BattleConfigProto& config,
        GameSettingsSPtr gameSettings,
        std::optional<GameStateW> initialGameState,
        BattleTraceOptions traceOptions)
    : config_(config),
      gameSettings_(std::move(gameSettings)),
      initialGameState_(std::move(initialGameState)),
      traceOptions_(std::move(traceOptions)) {
    if (!gameSettings_) {
        // Initialize default game settings
        gameSettings_ = ai_testing_common::GameSettingsFactory::CreateDefault();
    }
    if (config_.random_seed() != 0) {
        auto setter = gameSettings_->GetSetter();
        setter.SetRandomGenerator(std::make_shared<StdLibraryGenerator>(
                static_cast<std::mt19937_64::result_type>(config_.random_seed())));
    }
    if (config_.max_rounds() > 0) {
        auto setter = gameSettings_->GetSetter();
        setter.SetInt("maxRounds", config_.max_rounds());
    }
}

BattleResult AiBattleSimulator::RunBattle() {
    std::cout << "Starting AI vs AI battle simulation...\n";
    std::cout << "  Map: " << config_.map_name() << "\n";
    std::cout << "  Month: " << config_.month() << "\n";
    std::cout << "  Max rounds: " << config_.max_rounds() << "\n";
    std::cout << "  Random seed: " << config_.random_seed() << "\n";
    std::cout << "  Attacker AI: "
              << net::eagle0::shardok::ai_battle_simulator::AIAlgorithmType_Name(
                         config_.attacker().ai_algorithm())
              << "\n";
    std::cout << "  Defender AI: "
              << net::eagle0::shardok::ai_battle_simulator::AIAlgorithmType_Name(
                         config_.defender().ai_algorithm())
              << "\n";

    // Create initial game state
    auto gameState = initialGameState_.has_value() ? *initialGameState_ : CreateInitialGameState();

    // Get hex map from game state
    const auto* hexMap = gameState->hex_map();

    // Create AI clients for both players
    auto attackerAI = CreateAIClient(ATTACKER_ID, hexMap);
    auto defenderAI = CreateAIClient(DEFENDER_ID, hexMap);

    // Create the engine once for the entire simulation
    ShardokEngine engine(gameSettings_, gameState);

    int setupCommands = 0;
    if (!initialGameState_.has_value()) {
        // Run setup phase
        std::cout << "\nStarting setup phase...\n";
        auto setupResult = config_.use_configured_placements()
                                   ? RunConfiguredSetupPhase(engine)
                                   : RunSetupPhase(engine, *attackerAI, *defenderAI);
        if (setupResult.endReason != BattleResult::EndReason::DRAW) {
            // Game ended during setup (shouldn't happen normally)
            return setupResult;
        }
        setupCommands = setupResult.totalCommands;
    } else {
        std::cout << "\nStarting from saved Shardok state; skipping setup phase.\n";
    }

    // Run battle phase
    std::cout << "\nStarting battle phase...\n";
    return RunBattlePhase(engine, *attackerAI, *defenderAI, setupCommands);
}

GameStateW AiBattleSimulator::CreateInitialGameState() const {
    // Load map
    auto hexMapProto = LoadMap(config_.map_name());

    // Attacker info
    net::eagle0::shardok::common::PlayerInfo attackerInfo;
    attackerInfo.set_player_id(ATTACKER_ID);
    attackerInfo.set_is_defender(false);
    attackerInfo.set_starting_food(1000);
    attackerInfo.add_victory_conditions(
            net::eagle0::shardok::common::VICTORY_CONDITION_LAST_PLAYER_STANDING);
    attackerInfo.add_victory_conditions(
            net::eagle0::shardok::common::VICTORY_CONDITION_HOLDS_CRITICAL_TILES);

    // Defender info
    net::eagle0::shardok::common::PlayerInfo defenderInfo;
    defenderInfo.set_player_id(DEFENDER_ID);
    defenderInfo.set_is_defender(true);
    defenderInfo.set_starting_food(1000);
    defenderInfo.add_victory_conditions(
            net::eagle0::shardok::common::VICTORY_CONDITION_LAST_PLAYER_STANDING);
    defenderInfo.add_victory_conditions(
            net::eagle0::shardok::common::VICTORY_CONDITION_WIN_AFTER_MAX_ROUNDS);

    std::vector<net::eagle0::shardok::common::PlayerInfo> playerInfoProtos{
            std::move(attackerInfo),
            std::move(defenderInfo)};

    // Create units from config
    std::vector<net::eagle0::shardok::storage::fb::Unit> units;
    units.reserve(config_.attacker().units_size() + config_.defender().units_size());

    // Attacker units
    const int attackerUnitCount = config_.attacker().units_size();
    for (int i = 0; i < attackerUnitCount; ++i) {
        const auto& unitConfig = config_.attacker().units(i);

        net::eagle0::shardok::storage::fb::Unit unit{};
        unit.mutate_player_id(ATTACKER_ID);
        unit.mutate_unit_id(i);
        unit.mutate_eagle_player_id(ATTACKER_ID);
        unit.mutable_location() = net::eagle0::shardok::storage::fb::Coords(-1, -1);
        unit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT);
        unit.mutate_remaining_action_points(12);
        unit.mutate_hidden(false);
        unit.mutate_fortified(false);
        unit.mutate_can_flee(GetCanFleeOrDefault(unitConfig));
        unit.mutate_can_start_fire(false);
        unit.mutate_can_archery(false);
        unit.mutate_stun_rounds_remaining(0);
        unit.mutate_commanding_unit_id(-1);
        unit.mutate_targeted_unit(-1);
        unit.mutate_starting_position_index(unitConfig.starting_position_index());
        unit.mutate_has_moved_in_zoc(false);
        unit.mutate_volleys_remaining(0);
        unit.mutate_food_remaining(1000.0f);

        // Battalion
        net::eagle0::shardok::storage::fb::Battalion battalion;
        const auto battalionTypeId =
                static_cast<net::eagle0::shardok::storage::fb::BattalionTypeId>(
                        unitConfig.battalion_type_id());
        battalion.mutate_type(battalionTypeId);

        // Get battalion type to determine default capacity
        const auto& battalionType = gameSettings_->GetGetter().GetBattalionType(battalionTypeId);
        const double defaultSize = battalionType->capacity;
        unit.mutate_can_archery(battalionType->alwaysArcheryCapable || HeroCanArchery(unitConfig));
        unit.mutate_can_start_fire(HeroCanStartFire(unitConfig));

        battalion.mutate_size(GetBattalionSizeOrDefault(unitConfig.battalion(), defaultSize));
        battalion.mutate_armament(GetBattalionArmamentOrDefault(unitConfig.battalion()));
        battalion.mutate_training(GetBattalionTrainingOrDefault(unitConfig.battalion()));
        battalion.mutate_morale(GetBattalionMoraleOrDefault(unitConfig.battalion()));
        unit.mutable_battalion() = battalion;

        // Hero
        if (HasAttachedHero(unitConfig.profession())) {
            unit.mutate_has_attached_hero(true);

            net::eagle0::shardok::storage::fb::Hero hero;
            hero.mutate_strength(GetOrDefault(unitConfig.hero().strength(), DEFAULT_HERO_STRENGTH));
            hero.mutate_strength_xp(0);
            hero.mutate_agility(GetOrDefault(unitConfig.hero().agility(), DEFAULT_HERO_AGILITY));
            hero.mutate_agility_xp(0);
            hero.mutate_wisdom(GetOrDefault(unitConfig.hero().wisdom(), DEFAULT_HERO_WISDOM));
            hero.mutate_wisdom_xp(0);
            hero.mutate_charisma(GetOrDefault(unitConfig.hero().charisma(), DEFAULT_HERO_CHARISMA));
            hero.mutate_charisma_xp(0);
            hero.mutate_constitution(
                    GetOrDefault(unitConfig.hero().constitution(), DEFAULT_HERO_CONSTITUTION));
            hero.mutate_constitution_xp(0);
            const int vigor = GetOrDefault(unitConfig.hero().vigor(), DEFAULT_HERO_VIGOR);
            hero.mutate_vigor(vigor);
            hero.mutate_starting_vigor(vigor);
            hero.mutate_spent_vigor(0);
            hero.mutate_bravery(GetOrDefault(unitConfig.hero().bravery(), DEFAULT_HERO_BRAVERY));
            hero.mutate_integrity(
                    GetOrDefault(unitConfig.hero().integrity(), DEFAULT_HERO_INTEGRITY));
            hero.mutate_ambition(GetOrDefault(unitConfig.hero().ambition(), DEFAULT_HERO_AMBITION));
            hero.mutate_eagle_hero_id(i + 1);
            hero.mutate_is_vip(false);

            hero.mutable_profession_info().mutate_profession(
                    static_cast<net::eagle0::shardok::storage::fb::Profession>(
                            unitConfig.profession()));
            hero.mutable_profession_info().mutate_meteor_cast_state(
                    net::eagle0::shardok::storage::fb::MultiroundMagicState_NONE);

            hero.mutable_control_info().mutate_controlled_unit_id(-1);
            hero.mutable_control_info().mutate_controlled_this_round(false);

            unit.mutable_attached_hero() = hero;
        } else {
            unit.mutate_has_attached_hero(false);
        }
        if (battalion.size() <= 0 && !unit.has_attached_hero()) {
            unit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_RESERVED_SLOT);
        }

        // Initialize opponent knowledge
        unit.mutable_opponent_knowledge()->Mutate(0, 0);
        unit.mutable_opponent_knowledge()->Mutate(1, 0);

        units.push_back(std::move(unit));
    }

    // Defender units
    const int defenderUnitIdStart = attackerUnitCount;
    const int defenderUnitCount = config_.defender().units_size();
    for (int i = 0; i < defenderUnitCount; ++i) {
        const auto& unitConfig = config_.defender().units(i);

        net::eagle0::shardok::storage::fb::Unit unit{};
        unit.mutate_player_id(DEFENDER_ID);
        unit.mutate_unit_id(defenderUnitIdStart + i);
        unit.mutate_eagle_player_id(DEFENDER_ID);
        unit.mutable_location() = net::eagle0::shardok::storage::fb::Coords(-1, -1);
        unit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT);
        unit.mutate_remaining_action_points(0);  // Player 1 starts with 0 AP (not their turn yet)
        unit.mutate_hidden(false);
        unit.mutate_fortified(false);
        unit.mutate_can_flee(GetCanFleeOrDefault(unitConfig));
        unit.mutate_can_start_fire(false);
        unit.mutate_can_archery(false);
        unit.mutate_stun_rounds_remaining(0);
        unit.mutate_commanding_unit_id(-1);
        unit.mutate_targeted_unit(-1);
        unit.mutate_starting_position_index(unitConfig.starting_position_index());
        unit.mutate_has_moved_in_zoc(false);
        unit.mutate_volleys_remaining(0);
        unit.mutate_food_remaining(1000.0f);

        // Battalion
        net::eagle0::shardok::storage::fb::Battalion battalion;
        const auto battalionTypeId =
                static_cast<net::eagle0::shardok::storage::fb::BattalionTypeId>(
                        unitConfig.battalion_type_id());
        battalion.mutate_type(battalionTypeId);

        // Get battalion type to determine default capacity
        const auto& battalionType = gameSettings_->GetGetter().GetBattalionType(battalionTypeId);
        const double defaultSize = battalionType->capacity;
        unit.mutate_can_archery(battalionType->alwaysArcheryCapable || HeroCanArchery(unitConfig));
        unit.mutate_can_start_fire(HeroCanStartFire(unitConfig));

        battalion.mutate_size(GetBattalionSizeOrDefault(unitConfig.battalion(), defaultSize));
        battalion.mutate_armament(GetBattalionArmamentOrDefault(unitConfig.battalion()));
        battalion.mutate_training(GetBattalionTrainingOrDefault(unitConfig.battalion()));
        battalion.mutate_morale(GetBattalionMoraleOrDefault(unitConfig.battalion()));
        unit.mutable_battalion() = battalion;

        // Hero
        if (HasAttachedHero(unitConfig.profession())) {
            unit.mutate_has_attached_hero(true);

            net::eagle0::shardok::storage::fb::Hero hero;
            hero.mutate_strength(GetOrDefault(unitConfig.hero().strength(), DEFAULT_HERO_STRENGTH));
            hero.mutate_strength_xp(0);
            hero.mutate_agility(GetOrDefault(unitConfig.hero().agility(), DEFAULT_HERO_AGILITY));
            hero.mutate_agility_xp(0);
            hero.mutate_wisdom(GetOrDefault(unitConfig.hero().wisdom(), DEFAULT_HERO_WISDOM));
            hero.mutate_wisdom_xp(0);
            hero.mutate_charisma(GetOrDefault(unitConfig.hero().charisma(), DEFAULT_HERO_CHARISMA));
            hero.mutate_charisma_xp(0);
            hero.mutate_constitution(
                    GetOrDefault(unitConfig.hero().constitution(), DEFAULT_HERO_CONSTITUTION));
            hero.mutate_constitution_xp(0);
            const int vigor = GetOrDefault(unitConfig.hero().vigor(), DEFAULT_HERO_VIGOR);
            hero.mutate_vigor(vigor);
            hero.mutate_starting_vigor(vigor);
            hero.mutate_spent_vigor(0);
            hero.mutate_bravery(GetOrDefault(unitConfig.hero().bravery(), DEFAULT_HERO_BRAVERY));
            hero.mutate_integrity(
                    GetOrDefault(unitConfig.hero().integrity(), DEFAULT_HERO_INTEGRITY));
            hero.mutate_ambition(GetOrDefault(unitConfig.hero().ambition(), DEFAULT_HERO_AMBITION));
            hero.mutate_eagle_hero_id(defenderUnitIdStart + i + 1);
            hero.mutate_is_vip(false);

            hero.mutable_profession_info().mutate_profession(
                    static_cast<net::eagle0::shardok::storage::fb::Profession>(
                            unitConfig.profession()));
            hero.mutable_profession_info().mutate_meteor_cast_state(
                    net::eagle0::shardok::storage::fb::MultiroundMagicState_NONE);

            hero.mutable_control_info().mutate_controlled_unit_id(-1);
            hero.mutable_control_info().mutate_controlled_this_round(false);

            unit.mutable_attached_hero() = hero;
        } else {
            unit.mutate_has_attached_hero(false);
        }
        if (battalion.size() <= 0 && !unit.has_attached_hero()) {
            unit.mutate_status(net::eagle0::shardok::storage::fb::UnitStatus_RESERVED_SLOT);
        }

        // Initialize opponent knowledge
        unit.mutable_opponent_knowledge()->Mutate(0, 0);
        unit.mutable_opponent_knowledge()->Mutate(1, 0);

        units.push_back(std::move(unit));
    }

    // Create game state
    const bool isWinter = (config_.month() >= 10 || config_.month() <= 2);
    return shardok::fb::SetupInitialGameState(
            "ai_battle_sim",
            hexMapProto,
            playerInfoProtos,
            units,
            config_.month(),
            isWinter,
            gameSettings_->GetGetter());
}

std::unique_ptr<ShardokAIClient> AiBattleSimulator::CreateAIClient(
        PlayerId playerId,
        const HexMap* hexMap) const {
    const auto& playerConfig = (playerId == ATTACKER_ID) ? config_.attacker() : config_.defender();
    const bool isDefender = (playerId == DEFENDER_ID);

    AIAlgorithmType algorithmType = ConvertAIAlgorithmType(playerConfig.ai_algorithm());
    const ScoringCalculatorType scoringType =
            ConvertScoringCalculatorType(playerConfig.scoring_calculator());
    const auto mctsConfig = BuildMCTSConfig(playerConfig, config_.random_seed(), playerId);

    return ai_testing_common::AIClientFactory::Create(
            playerId,
            isDefender,
            hexMap,
            gameSettings_->GetGetter(),
            algorithmType,
            scoringType,
            true,
            mctsConfig);
}

BattleResult AiBattleSimulator::RunSetupPhase(
        ShardokEngine& engine,
        ShardokAIClient& attackerAI,
        ShardokAIClient& defenderAI) {
    auto getAI = [&](PlayerId playerId) -> ShardokAIClient& {
        return (playerId == ATTACKER_ID) ? attackerAI : defenderAI;
    };

    ai_testing_common::PhaseResult phaseResult;

    while (engine.GetCurrentGameState()->status()->state() ==
           net::eagle0::shardok::storage::fb::GameStatus_::State_SET_UP) {
        auto currentState = engine.GetCurrentGameState();
        PlayerId currentPlayer = currentState->current_player();

        auto availableCommands = engine.GetAvailableCommandProtos(currentPlayer, false);
        if (availableCommands.empty()) { break; }

        ShardokAIClient& activeAI = getAI(currentPlayer);
        const SideTotals preTotals = CalculateSideTotals(currentState);
        auto choiceResults = activeAI.ChooseCommandIndex(engine);
        const auto command = availableCommands.at(choiceResults.chosenIndex);

        const auto randomGenerator = gameSettings_->GetGetter().GetRandomGenerator();
        engine.PostCommand(currentPlayer, choiceResults.chosenIndex, randomGenerator);
        phaseResult.commandsExecuted++;
        const auto postState = engine.GetCurrentGameState();
        const SideTotals postTotals = CalculateSideTotals(postState);

        if (traceOptions_.output != nullptr) {
            *traceOptions_.output
                    << "{"
                    << "\"run_label\":\"" << JsonEscape(traceOptions_.runLabel) << "\","
                    << "\"config\":\"" << JsonEscape(traceOptions_.configPath) << "\","
                    << "\"state_file\":\"" << JsonEscape(traceOptions_.statePath) << "\","
                    << "\"phase\":\"setup\","
                    << "\"sequence\":" << phaseResult.commandsExecuted << ","
                    << "\"round\":" << static_cast<int>(currentState->current_round()) << ","
                    << "\"player\":" << static_cast<int>(currentPlayer) << ","
                    << "\"side\":\"" << (currentPlayer == ATTACKER_ID ? "attacker" : "defender")
                    << "\","
                    << "\"command_type\":\""
                    << net::eagle0::shardok::common::CommandType_Name(command.type()) << "\","
                    << "\"actor_unit\":" << (command.has_actor() ? command.actor().value() : -1)
                    << ","
                    << "\"target_row\":" << command.target().row() << ","
                    << "\"target_column\":" << command.target().column() << ","
                    << "\"chosen_index\":" << choiceResults.chosenIndex << ","
                    << "\"available_command_count\":" << choiceResults.availableCommandCount << ","
                    << "\"depth_achieved\":" << choiceResults.depthAchieved << ","
                    << "\"commands_evaluated\":" << choiceResults.commandCountEvaluated << ","
                    << "\"completion_reason\":\""
                    << CompletionReasonLabel(choiceResults.completionReason) << "\","
                    << "\"forced_commands\":0,"
                    << "\"pre_attacker_units\":" << preTotals.attackerUnits << ","
                    << "\"pre_defender_units\":" << preTotals.defenderUnits << ","
                    << "\"pre_attacker_troops\":" << preTotals.attackerTroops << ","
                    << "\"pre_defender_troops\":" << preTotals.defenderTroops << ","
                    << "\"post_attacker_units\":" << postTotals.attackerUnits << ","
                    << "\"post_defender_units\":" << postTotals.defenderUnits << ","
                    << "\"post_attacker_troops\":" << postTotals.attackerTroops << ","
                    << "\"post_defender_troops\":" << postTotals.defenderTroops << "}\n";
        }

        if (engine.GameIsOver()) {
            phaseResult.gameEnded = true;
            break;
        }
    }

    std::cout << "Setup phase complete. Commands executed: " << phaseResult.commandsExecuted
              << "\n";

    // Check if game ended unexpectedly
    if (phaseResult.gameEnded) {
        return CreateResultFromGameState(
                engine.GetCurrentGameState(),
                0,
                phaseResult.commandsExecuted);
    }

    // Return a "not finished" result
    BattleResult result;
    result.winner = -1;
    result.totalRounds = 0;
    result.totalCommands = phaseResult.commandsExecuted;
    result.endReason = BattleResult::EndReason::DRAW;  // Temporary placeholder
    result.description = "Setup phase completed";
    return result;
}

BattleResult AiBattleSimulator::RunConfiguredSetupPhase(ShardokEngine& engine) {
    auto placePlayer = [&](PlayerId playerId,
                           const net::eagle0::shardok::ai_battle_simulator::PlayerConfig&
                                   playerConfig,
                           int unitIdOffset) {
        std::vector<UnitPlacementInfo> placements;
        const int unitCount = playerConfig.units_size();
        placements.reserve(unitCount);
        for (int i = 0; i < unitCount; ++i) {
            const auto& unitConfig = playerConfig.units(i);
            if (!unitConfig.has_fixed_position()) {
                throw std::invalid_argument(
                        "use_configured_placements requires has_fixed_position for every unit");
            }
            placements.emplace_back(
                    static_cast<UnitId>(unitIdOffset + i),
                    Coords(unitConfig.row(), unitConfig.column()));
        }

        const auto randomGenerator = gameSettings_->GetGetter().GetRandomGenerator();
        engine.PostPlacementCommands(playerId, placements, randomGenerator);
        engine.PostFinishedPlacementCommand(playerId, randomGenerator);
        if (config_.apply_forced_commands()) {
            engine.PostWhileCurrentPlayerHasOnlyOneOption(randomGenerator);
        }
    };

    placePlayer(ATTACKER_ID, config_.attacker(), 0);
    if (!engine.GameIsOver()) {
        placePlayer(DEFENDER_ID, config_.defender(), config_.attacker().units_size());
    }

    std::cout << "Configured setup complete.\n";

    if (engine.GameIsOver()) {
        return CreateResultFromGameState(
                engine.GetCurrentGameState(),
                0,
                static_cast<int>(engine.GetUnfilteredHistoryCount()));
    }

    BattleResult result;
    result.winner = -1;
    result.totalRounds = 0;
    result.totalCommands = static_cast<int>(engine.GetUnfilteredHistoryCount());
    result.endReason = BattleResult::EndReason::DRAW;
    result.description = "Configured setup completed";
    return result;
}

BattleResult AiBattleSimulator::RunBattlePhase(
        ShardokEngine& engine,
        ShardokAIClient& attackerAI,
        ShardokAIClient& defenderAI,
        int setupCommands) {
    int totalCommands = setupCommands;
    int attackerCommands = 0;
    int defenderCommands = 0;
    int attackerDepthTotal = 0;
    int defenderDepthTotal = 0;
    int currentRound = std::max<int>(1, engine.GetCurrentGameState()->current_round());

    const int harnessMaxRounds = config_.max_rounds() > 0
                                         ? config_.max_rounds()
                                         : gameSettings_->GetGetter().Backing().max_rounds();

    while (!engine.GameIsOver() && currentRound <= harnessMaxRounds) {
        auto currentState = engine.GetCurrentGameState();
        PlayerId currentPlayer = currentState->current_player();

        auto availableCommands = engine.GetAvailableCommandProtos(currentPlayer, false);

        if (availableCommands.empty()) {
            std::cout << "No commands available for player " << static_cast<int>(currentPlayer)
                      << "\n";
            break;
        }

        // Choose which AI to use
        ShardokAIClient& activeAI = (currentPlayer == ATTACKER_ID) ? attackerAI : defenderAI;

        // Get AI decision
        const SideTotals preTotals = CalculateSideTotals(currentState);
        auto choiceResults = activeAI.ChooseCommandIndex(engine);
        const auto command = availableCommands.at(choiceResults.chosenIndex);

        // Apply command
        const auto randomGenerator = gameSettings_->GetGetter().GetRandomGenerator();
        engine.PostCommand(currentPlayer, choiceResults.chosenIndex, randomGenerator);
        totalCommands++;
        int forcedCommands = 0;
        if (currentPlayer == ATTACKER_ID) {
            attackerCommands++;
            attackerDepthTotal += choiceResults.depthAchieved;
        } else if (currentPlayer == DEFENDER_ID) {
            defenderCommands++;
            defenderDepthTotal += choiceResults.depthAchieved;
        }

        if (config_.apply_forced_commands()) {
            forcedCommands = engine.PostWhileCurrentPlayerHasOnlyOneOption(randomGenerator);
            totalCommands += forcedCommands;
        }

        // Check round progression
        auto newState = engine.GetCurrentGameState();
        const SideTotals postTotals = CalculateSideTotals(newState);
        if (traceOptions_.output != nullptr) {
            const int sequence = totalCommands - forcedCommands;
            *traceOptions_.output
                    << "{"
                    << "\"run_label\":\"" << JsonEscape(traceOptions_.runLabel) << "\","
                    << "\"config\":\"" << JsonEscape(traceOptions_.configPath) << "\","
                    << "\"state_file\":\"" << JsonEscape(traceOptions_.statePath) << "\","
                    << "\"phase\":\"battle\","
                    << "\"sequence\":" << sequence << ","
                    << "\"round\":" << static_cast<int>(currentState->current_round()) << ","
                    << "\"player\":" << static_cast<int>(currentPlayer) << ","
                    << "\"side\":\"" << (currentPlayer == ATTACKER_ID ? "attacker" : "defender")
                    << "\","
                    << "\"command_type\":\""
                    << net::eagle0::shardok::common::CommandType_Name(command.type()) << "\","
                    << "\"actor_unit\":" << (command.has_actor() ? command.actor().value() : -1)
                    << ","
                    << "\"target_row\":" << command.target().row() << ","
                    << "\"target_column\":" << command.target().column() << ","
                    << "\"chosen_index\":" << choiceResults.chosenIndex << ","
                    << "\"available_command_count\":" << choiceResults.availableCommandCount << ","
                    << "\"depth_achieved\":" << choiceResults.depthAchieved << ","
                    << "\"commands_evaluated\":" << choiceResults.commandCountEvaluated << ","
                    << "\"completion_reason\":\""
                    << CompletionReasonLabel(choiceResults.completionReason) << "\","
                    << "\"forced_commands\":" << forcedCommands << ","
                    << "\"pre_attacker_units\":" << preTotals.attackerUnits << ","
                    << "\"pre_defender_units\":" << preTotals.defenderUnits << ","
                    << "\"pre_attacker_troops\":" << preTotals.attackerTroops << ","
                    << "\"pre_defender_troops\":" << preTotals.defenderTroops << ","
                    << "\"post_attacker_units\":" << postTotals.attackerUnits << ","
                    << "\"post_defender_units\":" << postTotals.defenderUnits << ","
                    << "\"post_attacker_troops\":" << postTotals.attackerTroops << ","
                    << "\"post_defender_troops\":" << postTotals.defenderTroops << "}\n";
        }
        if (newState->current_round() > currentRound) {
            std::cout << "Round " << currentRound
                      << " completed. Commands this round: " << totalCommands << "\n";
            currentRound = newState->current_round();
        }
    }

    std::cout << "Battle phase complete. Total rounds: " << currentRound
              << ", Total commands: " << totalCommands << "\n";

    auto result =
            CreateResultFromGameState(engine.GetCurrentGameState(), currentRound, totalCommands);
    if (!engine.GameIsOver() && currentRound > harnessMaxRounds) {
        result.winner = DEFENDER_ID;
        result.totalRounds = harnessMaxRounds;
        result.endReason = BattleResult::EndReason::MAX_ROUNDS_REACHED;
        result.description = "Defender won by surviving maximum rounds";
    }
    result.attackerCommands = attackerCommands;
    result.defenderCommands = defenderCommands;
    result.attackerAverageDepth =
            attackerCommands == 0 ? 0.0
                                  : static_cast<double>(attackerDepthTotal) / attackerCommands;
    result.defenderAverageDepth =
            defenderCommands == 0 ? 0.0
                                  : static_cast<double>(defenderDepthTotal) / defenderCommands;
    return result;
}

bool AiBattleSimulator::IsGameOver(const GameStateW& state) const {
    const auto stateValue = state->status()->state();
    return stateValue == net::eagle0::shardok::storage::fb::GameStatus_::State_VICTORY ||
           stateValue == net::eagle0::shardok::storage::fb::GameStatus_::State_DRAW;
}

BattleResult AiBattleSimulator::CreateResultFromGameState(
        const GameStateW& state,
        int totalRounds,
        int totalCommands) const {
    BattleResult result;
    result.totalRounds = totalRounds;
    result.totalCommands = totalCommands;

    const auto* status = state->status();
    const auto stateValue = status->state();

    // Determine winner and end reason
    if (stateValue == net::eagle0::shardok::storage::fb::GameStatus_::State_VICTORY) {
        // Extract winner from winning_shardok_ids
        const auto* winningIds = status->winning_shardok_ids();
        if (winningIds && winningIds->size() > 0) {
            result.winner = winningIds->Get(0);

            if (result.winner == ATTACKER_ID) {
                if (status->end_game_condition()->victory_details() ==
                    net::eagle0::shardok::storage::fb::
                            VictoryCondition_VICTORY_CONDITION_HOLDS_CRITICAL_TILES) {
                    result.endReason = BattleResult::EndReason::CRITICAL_TILES_CONTROLLED;
                    result.description = "Attacker won by controlling critical tiles";
                } else {
                    result.endReason = BattleResult::EndReason::LAST_PLAYER_STANDING;
                    result.description = "Attacker won by eliminating all defenders";
                }
            } else if (result.winner == DEFENDER_ID) {
                if (status->end_game_condition()->victory_details() ==
                    net::eagle0::shardok::storage::fb::
                            VictoryCondition_VICTORY_CONDITION_WIN_AFTER_MAX_ROUNDS) {
                    result.endReason = BattleResult::EndReason::MAX_ROUNDS_REACHED;
                    result.description = "Defender won by surviving maximum rounds";
                } else {
                    result.endReason = BattleResult::EndReason::LAST_PLAYER_STANDING;
                    result.description = "Defender won by eliminating all attackers";
                }
            } else {
                result.winner = -1;
                result.endReason = BattleResult::EndReason::DRAW;
                result.description = "Game ended with unknown winner";
            }
        } else {
            result.winner = -1;
            result.endReason = BattleResult::EndReason::DRAW;
            result.description = "Game ended with no winner";
        }
    } else if (stateValue == net::eagle0::shardok::storage::fb::GameStatus_::State_DRAW) {
        result.winner = -1;
        result.endReason = BattleResult::EndReason::DRAW;
        result.description = "Game ended in a draw";
    } else {
        // If the harness stops before the engine posts its max-round status transition, Shardok's
        // battle rule still awards the battle to the defender.
        result.winner = DEFENDER_ID;
        result.endReason = BattleResult::EndReason::MAX_ROUNDS_REACHED;
        result.description = "Defender won by surviving maximum rounds";
    }

    // Collect surviving units (for winner, or all units if draw)
    const auto* units = state->units();
    if (units) {
        result.survivingUnits.reserve(units->size());
        for (const auto* unit : *units) {
            if (!unit) { continue; }

            // Only include units that are still alive
            // Based on GameStateFilter.cpp logic
            const bool isAlive =
                    unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT ||
                    unit->status() == net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT ||
                    unit->status() ==
                            net::eagle0::shardok::storage::fb::UnitStatus_NEVER_ENTERED_UNIT;

            if (!isAlive) { continue; }

            if (unit->player_id() == ATTACKER_ID) {
                result.attackerSurvivingUnits++;
                result.attackerSurvivingTroops += unit->battalion().size();
            } else if (unit->player_id() == DEFENDER_ID) {
                result.defenderSurvivingUnits++;
                result.defenderSurvivingTroops += unit->battalion().size();
            }

            // For victories, only include winner's units; for draws, include all
            if (result.winner != -1 && unit->player_id() != result.winner) { continue; }

            SurvivingUnit survivor;
            survivor.unitId = unit->unit_id();
            survivor.battalionSize = unit->battalion().size();
            survivor.profession = unit->has_attached_hero()
                                          ? unit->attached_hero().profession_info().profession()
                                          : 0;

            // Get battalion type name
            const auto& battalionType =
                    gameSettings_->GetGetter().GetBattalionType(unit->battalion().type());
            survivor.battalionTypeName = battalionType->name;

            result.survivingUnits.push_back(survivor);
        }
    }

    return result;
}

}  // namespace shardok::ai_battle_simulator
