//
// Converts saved Eagle ShardokBattle/GameState protobuf blobs into AI battle configs.
//

#include <fstream>
#include <iostream>
#include <optional>
#include <stdexcept>
#include <string>

#include "AiBattleConfig.hpp"
#include "src/main/protobuf/net/eagle0/eagle/internal/game_state.pb.h"
#include "src/main/protobuf/net/eagle0/eagle/internal/shardok_battle.pb.h"

namespace {

using BattleConfigProto = shardok::ai_battle_simulator::BattleConfigProto;
using net::eagle0::eagle::internal::Battalion;
using net::eagle0::eagle::internal::GameState;
using net::eagle0::eagle::internal::Hero;
using net::eagle0::eagle::internal::ShardokBattle;
using net::eagle0::eagle::internal::ShardokPlayer;
using net::eagle0::shardok::ai_battle_simulator::AIAlgorithmType;
using net::eagle0::shardok::ai_battle_simulator::BattalionConfig;
using net::eagle0::shardok::ai_battle_simulator::HeroConfig;
using net::eagle0::shardok::ai_battle_simulator::PlayerConfig;
using net::eagle0::shardok::ai_battle_simulator::ScoringCalculatorType;
using net::eagle0::shardok::ai_battle_simulator::UnitConfig;

constexpr int NO_PROFESSION = 14;

std::string ReadFile(const std::string &path) {
    std::ifstream input(path, std::ios::binary);
    if (!input.is_open()) { throw std::runtime_error("Failed to open " + path); }
    return std::string(std::istreambuf_iterator<char>(input), std::istreambuf_iterator<char>());
}

void WriteFile(const std::string &path, const std::string &contents) {
    std::ofstream output(path);
    if (!output.is_open()) { throw std::runtime_error("Failed to write " + path); }
    output << contents;
}

std::optional<ScoringCalculatorType> ParseScoring(const std::string &value) {
    if (value == "STANDARD") { return ScoringCalculatorType::STANDARD; }
    if (value == "MCTS_OPTIMIZED") { return ScoringCalculatorType::MCTS_OPTIMIZED; }
    if (value == "EXPERIMENTAL") { return ScoringCalculatorType::EXPERIMENTAL; }
    return std::nullopt;
}

void CopyHero(const Hero &source, HeroConfig *target) {
    target->set_strength(source.strength());
    target->set_agility(source.agility());
    target->set_wisdom(source.wisdom());
    target->set_charisma(source.charisma());
    target->set_constitution(source.constitution());
    target->set_bravery(source.bravery());
    target->set_integrity(source.integrity());
    target->set_ambition(source.ambition());
    target->set_vigor(static_cast<int>(source.vigor()));
}

void CopyBattalion(const Battalion &source, BattalionConfig *target) {
    target->set_size(source.size());
    target->set_training(source.training());
    target->set_armament(source.armament());
    target->set_morale(50.0);
}

int StartingPositionForArmy(const net::eagle0::eagle::internal::MovingArmy &army, bool defender) {
    if (army.has_starting_position_index()) { return army.starting_position_index().value(); }
    return defender ? -1 : 0;
}

void AddPlayerUnits(const GameState &state, const ShardokPlayer &source, PlayerConfig *target) {
    target->set_ai_algorithm(AIAlgorithmType::ITERATIVE_DEEPENING);

    for (const auto &movingArmy : source.army_group().armies()) {
        const int startingPosition = StartingPositionForArmy(movingArmy, source.is_defender());
        for (const auto &combatUnit : movingArmy.army().units()) {
            const int heroId = combatUnit.hero_id();
            const auto heroIt = state.heroes().find(heroId);
            const bool hasHero = heroId > 0 && heroIt != state.heroes().end();

            const Battalion *battalion = nullptr;
            if (combatUnit.has_battalion_id()) {
                const int battalionId = combatUnit.battalion_id().value();
                const auto battalionIt = state.battalions().find(battalionId);
                if (battalionIt != state.battalions().end()) { battalion = &battalionIt->second; }
            }

            const bool hasTroops = battalion != nullptr && battalion->size() > 0;
            if (!hasHero && !hasTroops) { continue; }

            UnitConfig *unit = target->add_units();
            unit->set_starting_position_index(startingPosition);
            unit->set_can_flee(movingArmy.army().has_flee_province_id());

            if (hasHero) {
                unit->set_profession(static_cast<int>(heroIt->second.profession()));
                CopyHero(heroIt->second, unit->mutable_hero());
            } else {
                unit->set_profession(NO_PROFESSION);
            }

            if (battalion != nullptr) {
                unit->set_battalion_type_id(static_cast<int>(battalion->type()));
                CopyBattalion(*battalion, unit->mutable_battalion());
            } else {
                unit->set_battalion_type_id(0);
                unit->mutable_battalion()->set_size(0.0);
                unit->mutable_battalion()->set_training(50.0);
                unit->mutable_battalion()->set_armament(50.0);
                unit->mutable_battalion()->set_morale(50.0);
            }
        }
    }
}

void PrintUsage(const char *program) {
    std::cerr << "Usage: " << program
              << " --battle=<new_battle_proto> --state=<game_state_proto> --output=<config.json>"
              << " [--seed=N] [--max-rounds=N] [--attacker-scoring=STANDARD]"
              << " [--defender-scoring=STANDARD]\n";
}

}  // namespace

int main(int argc, char *argv[]) {
    try {
        std::string battlePath;
        std::string statePath;
        std::string outputPath;
        int seed = 1;
        int maxRounds = 31;
        ScoringCalculatorType attackerScoring = ScoringCalculatorType::STANDARD;
        ScoringCalculatorType defenderScoring = ScoringCalculatorType::STANDARD;

        for (int i = 1; i < argc; ++i) {
            const std::string arg = argv[i];
            if (arg.starts_with("--battle=")) {
                battlePath = arg.substr(9);
            } else if (arg.starts_with("--state=")) {
                statePath = arg.substr(8);
            } else if (arg.starts_with("--output=")) {
                outputPath = arg.substr(9);
            } else if (arg.starts_with("--seed=")) {
                seed = std::stoi(arg.substr(7));
            } else if (arg.starts_with("--max-rounds=")) {
                maxRounds = std::stoi(arg.substr(13));
            } else if (arg.starts_with("--attacker-scoring=")) {
                const auto scoring = ParseScoring(arg.substr(19));
                if (!scoring.has_value()) { throw std::runtime_error("Unknown attacker scoring"); }
                attackerScoring = *scoring;
            } else if (arg.starts_with("--defender-scoring=")) {
                const auto scoring = ParseScoring(arg.substr(19));
                if (!scoring.has_value()) { throw std::runtime_error("Unknown defender scoring"); }
                defenderScoring = *scoring;
            } else {
                PrintUsage(argv[0]);
                return 1;
            }
        }

        if (battlePath.empty() || statePath.empty() || outputPath.empty()) {
            PrintUsage(argv[0]);
            return 1;
        }

        ShardokBattle battle;
        if (!battle.ParseFromString(ReadFile(battlePath))) {
            throw std::runtime_error("Failed to parse battle proto");
        }

        GameState state;
        if (!state.ParseFromString(ReadFile(statePath))) {
            throw std::runtime_error("Failed to parse game state proto");
        }

        BattleConfigProto config;
        config.set_map_name(battle.hex_map_name());
        config.set_month(state.current_date().month());
        config.set_max_rounds(maxRounds);
        config.set_random_seed(seed);
        config.set_use_configured_placements(false);
        config.set_apply_forced_commands(true);
        config.mutable_attacker()->set_scoring_calculator(attackerScoring);
        config.mutable_defender()->set_scoring_calculator(defenderScoring);

        for (const auto &player : battle.players()) {
            AddPlayerUnits(
                    state,
                    player,
                    player.is_defender() ? config.mutable_defender() : config.mutable_attacker());
        }

        WriteFile(
                outputPath,
                shardok::ai_battle_simulator::AiBattleConfigLoader::ToJsonString(config));

        std::cerr << battle.shardok_game_id()
                  << ",attacker_units=" << config.attacker().units_size()
                  << ",defender_units=" << config.defender().units_size() << "\n";
        return 0;
    } catch (const std::exception &e) {
        std::cerr << "Error: " << e.what() << "\n";
        return 1;
    }
}
