//
// Self-Play Generator - Generate ML training data from AI vs AI games
//

#include <filesystem>
#include <iomanip>
#include <iostream>
#include <memory>
#include <string>

#include "SelfPlayBattleRunner.hpp"
#include "StateRecorder.hpp"
#include "src/main/cpp/net/eagle0/common/FilesystemUtils.hpp"
#include "src/main/cpp/net/eagle0/common/TsvParser.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/ShardokAIClient.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai_battle_simulator/AiBattleConfig.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai_battle_simulator/AiBattleSimulator.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/action_point_distances/FixedActionPointDistances.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/BattalionTypeRegistrar.hpp"
#include "src/main/cpp/net/eagle0/shardok/util/MapLoader.hpp"

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

namespace {

using shardok::AIAlgorithmType;
using shardok::GameSettingsSPtr;
using shardok::GameStateW;
using shardok::PlayerId;
using shardok::ShardokAIClient;
using shardok::ShardokEngine;
using shardok::ai_battle_simulator::AiBattleConfigLoader;
using shardok::ai_battle_simulator::AiBattleSimulator;
using shardok::ai_battle_simulator::BattleResult;
using shardok::ml::StateRecorder;

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

struct GeneratorConfig {
    int numGames = 100;
    std::string outputDir = "/tmp/shardok_training_data";
    std::string configPath;  // If empty, use default config
    AIAlgorithmType aiAlgorithm = AIAlgorithmType::ITERATIVE_DEEPENING;
    bool verbose = false;
};

void PrintUsage(const char* programName) {
    std::cout << "Self-Play Generator - Generate ML Training Data\n"
              << "\n"
              << "Usage:\n"
              << "  " << programName << " [OPTIONS]\n"
              << "\n"
              << "Options:\n"
              << "  --num_games=N          Number of games to generate (default: 100)\n"
              << "  --output_dir=PATH      Output directory for data files (default: "
                 "/tmp/shardok_training_data)\n"
              << "  --config=PATH          Battle config JSON file (default: use built-in "
                 "config)\n"
              << "  --ai=TYPE              AI algorithm: iterative_deepening or mcts (default: "
                 "iterative_deepening)\n"
              << "  --verbose              Print detailed game information\n"
              << "  --help                 Show this help message\n"
              << "\n"
              << "Examples:\n"
              << "  # Generate 1000 games with default settings\n"
              << "  " << programName << " --num_games=1000\n"
              << "\n"
              << "  # Generate games using MCTS AI\n"
              << "  " << programName << " --num_games=500 --ai=mcts\n"
              << "\n"
              << "  # Use custom battle config\n"
              << "  " << programName << " --config=my_battle.json --num_games=100\n"
              << "\n";
}

GeneratorConfig ParseCommandLine(int argc, char* argv[]) {
    GeneratorConfig config;

    for (int i = 1; i < argc; ++i) {
        std::string arg(argv[i]);

        if (arg == "--help" || arg == "-h") {
            PrintUsage(argv[0]);
            std::exit(0);
        } else if (arg.starts_with("--num_games=")) {
            config.numGames = std::stoi(arg.substr(12));
        } else if (arg.starts_with("--output_dir=")) {
            config.outputDir = arg.substr(13);
        } else if (arg.starts_with("--config=")) {
            config.configPath = arg.substr(9);
        } else if (arg.starts_with("--ai=")) {
            std::string aiType = arg.substr(5);
            if (aiType == "mcts") {
                config.aiAlgorithm = AIAlgorithmType::MCTS;
            } else if (aiType == "iterative_deepening") {
                config.aiAlgorithm = AIAlgorithmType::ITERATIVE_DEEPENING;
            } else {
                std::cerr << "Error: Unknown AI type: " << aiType << "\n";
                std::exit(1);
            }
        } else if (arg == "--verbose") {
            config.verbose = true;
        } else {
            std::cerr << "Error: Unknown argument: " << arg << "\n";
            PrintUsage(argv[0]);
            std::exit(1);
        }
    }

    return config;
}

/**
 * @brief Run a complete self-play game with state recording.
 */
struct GameResult {
    PlayerId winner;
    int totalRounds;
    int totalCommands;
    int statesRecorded;
};

GameResult RunCompleteGameWithRecording(
        const net::eagle0::shardok::ai_battle_simulator::BattleConfig& battleConfig,
        const GameSettingsSPtr& gameSettings,
        AIAlgorithmType aiAlgorithm,
        StateRecorder& recorder,
        bool verbose) {
    // Use SelfPlayBattleRunner to run the game with full state recording
    shardok::ml::SelfPlayBattleRunner runner(battleConfig, gameSettings);
    auto result = runner.RunBattleWithRecording(recorder, aiAlgorithm, verbose);

    // Convert to GameResult
    GameResult gameResult;
    gameResult.winner = result.winner;
    gameResult.totalRounds = result.totalRounds;
    gameResult.totalCommands = result.totalCommands;
    gameResult.statesRecorded = result.statesRecorded;

    return gameResult;
}

void GenerateTrainingData(const GeneratorConfig& config) {
    std::cout << "Self-Play Training Data Generator\n";
    std::cout << "==================================\n";
    std::cout << "Output directory: " << config.outputDir << "\n";
    std::cout << "Number of games: " << config.numGames << "\n";
    std::cout << "AI algorithm: "
              << (config.aiAlgorithm == AIAlgorithmType::MCTS ? "MCTS" : "Iterative Deepening")
              << "\n\n";

    // Create output directory
    std::filesystem::create_directories(config.outputDir);

    // Load battle config
    // Note: We don't create a single config upfront because we want RANDOMIZED configs per game
    if (!config.configPath.empty()) {
        // If user provides a custom config, validate it exists
        auto testConfig = AiBattleConfigLoader::LoadFromJsonFile(config.configPath);
        if (!testConfig) {
            std::cerr << "Error: Failed to load config from " << config.configPath << "\n";
            std::exit(1);
        }
        std::cout << "Using custom config from: " << config.configPath << "\n";
    } else {
        std::cout << "Using randomized battle configurations for ML training\n";
    }

    // Initialize game settings once (shared across all games)
    auto gameSettings = std::make_shared<shardok::GameSettings>();
    auto setter = gameSettings->GetSetter();
    shardok::BattalionTypeRegistrar::RegisterBattalionTypes(setter);

    const std::string settingsPath =
            FilesystemUtils::StaticShardokFilesDirectory() + "settings.tsv";
    const std::string settingsTsv = std::string(byte_vector::FromPath(settingsPath));
    TsvParser parser;
    const auto valuesAndTypes = parser.ParseColumnEntryTsv(settingsTsv);
    setter.SetFromTypesAndValues(valuesAndTypes[1], valuesAndTypes[0]);

    std::cout << "\nGenerating games...\n";

    int attackerWins = 0;
    int defenderWins = 0;
    int draws = 0;
    size_t totalStatesRecorded = 0;

    for (int gameNum = 0; gameNum < config.numGames; ++gameNum) {
        StateRecorder recorder;

        if (config.verbose || gameNum % 10 == 0) {
            std::cout << "Game " << (gameNum + 1) << "/" << config.numGames << "...\n";
        }

        try {
            // Create a new randomized config for each game (or load from file)
            std::unique_ptr<net::eagle0::shardok::ai_battle_simulator::BattleConfig> battleConfig;
            if (config.configPath.empty()) {
                battleConfig = AiBattleConfigLoader::CreateRandomizedBattleConfig();
            } else {
                battleConfig = AiBattleConfigLoader::LoadFromJsonFile(config.configPath);
            }

            // Override AI algorithm if specified
            battleConfig->mutable_attacker()->set_ai_algorithm(
                    config.aiAlgorithm == AIAlgorithmType::MCTS
                            ? net::eagle0::shardok::ai_battle_simulator::MCTS
                            : net::eagle0::shardok::ai_battle_simulator::ITERATIVE_DEEPENING);
            battleConfig->mutable_defender()->set_ai_algorithm(
                    config.aiAlgorithm == AIAlgorithmType::MCTS
                            ? net::eagle0::shardok::ai_battle_simulator::MCTS
                            : net::eagle0::shardok::ai_battle_simulator::ITERATIVE_DEEPENING);

            if (config.verbose) {
                std::cout << "  Config: Month=" << battleConfig->month()
                          << ", Attacker=" << battleConfig->attacker().units_size() << " units"
                          << ", Defender=" << battleConfig->defender().units_size() << " units\n";
            }

            auto result = RunCompleteGameWithRecording(
                    *battleConfig,
                    gameSettings,
                    config.aiAlgorithm,
                    recorder,
                    config.verbose);

            // Update statistics
            if (result.winner == ATTACKER_ID) {
                attackerWins++;
            } else if (result.winner == DEFENDER_ID) {
                defenderWins++;
            } else {
                draws++;
            }

            totalStatesRecorded += result.statesRecorded;

            if (config.verbose) {
                std::cout << "  Winner: "
                          << (result.winner == ATTACKER_ID   ? "Attacker"
                              : result.winner == DEFENDER_ID ? "Defender"
                                                             : "Draw")
                          << "\n";
                std::cout << "  Rounds: " << result.totalRounds
                          << ", Commands: " << result.totalCommands
                          << ", States: " << result.statesRecorded << "\n";
            }

            // Write data file (when we have states recorded)
            if (result.statesRecorded > 0) {
                recorder.SetGameOutcome(result.winner);

                std::ostringstream filename;
                filename << config.outputDir << "/game_" << std::setw(6) << std::setfill('0')
                         << gameNum << ".bin";
                recorder.WriteToFile(filename.str());
            }

        } catch (const std::exception& e) {
            std::cerr << "Error in game " << gameNum << ": " << e.what() << "\n";
            continue;
        }
    }

    // Print summary
    std::cout << "\n";
    std::cout << "==================================\n";
    std::cout << "Training Data Generation Complete\n";
    std::cout << "==================================\n";
    std::cout << "Games generated: " << config.numGames << "\n";
    std::cout << "Total states recorded: " << totalStatesRecorded << "\n";
    if (config.numGames > 0) {
        std::cout << "Average states per game: "
                  << (totalStatesRecorded / static_cast<double>(config.numGames)) << "\n";
    }
    std::cout << "\n";
    std::cout << "Outcomes:\n";
    std::cout << "  Attacker wins: " << attackerWins << " ("
              << (100.0 * attackerWins / config.numGames) << "%)\n";
    std::cout << "  Defender wins: " << defenderWins << " ("
              << (100.0 * defenderWins / config.numGames) << "%)\n";
    std::cout << "  Draws: " << draws << " (" << (100.0 * draws / config.numGames) << "%)\n";
    std::cout << "\n";

    if (totalStatesRecorded == 0) {
        std::cout << "WARNING: No states were recorded. Check verbose output for errors.\n";
    } else {
        std::cout << "Data files written to: " << config.outputDir << "\n";
    }
}

}  // namespace

int main(int argc, char* argv[]) {
    // Set exec path for FilesystemUtils
    FilesystemUtils::SetExecPath(argv[0]);

    // Set cache directory for ActionPointDistances
    shardok::FixedActionPointDistances::SetCacheDirectory(
            FilesystemUtils::CacheFilesDirectory() + "apdCache/");

    try {
        auto config = ParseCommandLine(argc, argv);
        GenerateTrainingData(config);
        return 0;

    } catch (const std::exception& e) {
        std::cerr << "Fatal error: " << e.what() << "\n";
        return 1;
    }
}
