//
//  ShardokAIClient.cpp
//  eagle0
//
//  Created by Dan Crosby on 10/18/17.
//  Copyright © 2017 none. All rights reserved.
//

#include "ShardokAIClient.hpp"

#define DEBUG_FLEE_DECISIONS

// Enable to dump game state and debug tree to /tmp for debugging
// #define ENABLE_MCTS_DEBUG_DUMP

#ifdef ENABLE_MCTS_DEBUG_DUMP
#include <chrono>
#include <fstream>
#include <iomanip>
#include <sstream>
#endif

#include <set>
#include <unordered_map>

#include "AIAttackerStrategySelector.hpp"
#include "AIConfig.hpp"
#include "AIDefenderStrategySelector.hpp"
#include "AIFleeDecisionCalculator.hpp"
#include "AIScoreUtilities.hpp"
#include "AITimeBudget.hpp"
#include "IterativeDeepeningAI.hpp"
#include "mcts/ShardokMCTSAI.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/MCTSOptimizedAIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/NormalizedAIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/score/StandardAIScoreCalculator.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/util/HexMapUtils.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/view_filters/GameStateGuesser.hpp"
#include "src/main/protobuf/net/eagle0/shardok/api/action_result_view.pb.h"

namespace shardok {

using net::eagle0::shardok::api::ActionResultView;
using net::eagle0::shardok::api::GameStateView;

static constexpr bool kPerformanceLogging = true;

void ApplyUpdate(GameStateView & /*currentView*/, const ActionResultView & /*update*/) {}

auto RoundsRemaining(const GameSettingsSPtr &settings, const GameStateView &gsv) -> int {
    const int maxRounds = settings->GetGetter().Backing().max_rounds();

    return maxRounds - gsv.current_round();
}

ShardokAIClient::ShardokAIClient(
        const PlayerId playerId,
        const bool isDefender,
        const bool isAllAiBattle,
        const HexMap *hexMap,
        const SettingsGetter &settings,
        const AIAlgorithmType aiAlgorithmType,
        const ScoringCalculatorType scoringCalculatorType,
        const mcts::MCTSConfig &mctsConfig)
    : playerId(playerId),
      isDefender(isDefender),
      isAllAiBattle(isAllAiBattle),
      aiAlgorithmType(aiAlgorithmType),
      scoringCalculatorType(scoringCalculatorType),
      alCache(std::make_unique<AttackLocationsCache>(hexMap, settings)),
      waterCrossingCommandChooser(playerId, apdCache),
      mctsConfig(mctsConfig) {
    // Pre-generate the most common cache entries for better performance
    const auto mapId = ActionPointDistancesCache::GetMapId(hexMap);

    // Pre-fetch for all battalion types, both with and without brave water
    using BattalionTypeId = net::eagle0::shardok::storage::fb::BattalionTypeId;

    for (int typeId = BattalionTypeId::BattalionTypeId_MIN;
         typeId <= BattalionTypeId::BattalionTypeId_MAX;
         typeId++) {
        const auto battalionTypeId = static_cast<BattalionTypeId>(typeId);
        const auto battalionType = settings.GetBattalionType(battalionTypeId);

        // Pre-fetch without brave water (braveWaterActionPointCost = -1)
        apdCache->GetRaw(hexMap, mapId, battalionType, false, -1);

        // Pre-fetch with brave water (includeBravingWater = true, braveWaterActionPointCost = 0)
        apdCache->GetRaw(hexMap, mapId, battalionType, true, 0);
    }

    // Consolidate all the pre-fetched entries into the persistent cache
    apdCache->ConsolidateThreadLocalCache_Racy();
}

void CheckCommand(const CommandSPtr &realCommand, const CommandSPtr &guessedCommand) {
    // Verify that the AI's guessed state produces the same available commands as the real state.
    // We only compare fields that uniquely identify a command - metadata fields like action_points,
    // will_unhide, next_round_target_info are not part of command identity.

    if (realCommand->GetCommandType() != guessedCommand->GetCommandType()) {
        throw ShardokInternalErrorException("Command type mismatch between real and guessed state");
    }

    if (realCommand->GetPlayerId() != guessedCommand->GetPlayerId()) {
        throw ShardokInternalErrorException("Player ID mismatch between real and guessed state");
    }

    if (realCommand->GetActorUnitId() != guessedCommand->GetActorUnitId()) {
        throw ShardokInternalErrorException("Actor unit mismatch between real and guessed state");
    }

    if (realCommand->GetTargetRow() != guessedCommand->GetTargetRow() ||
        realCommand->GetTargetColumn() != guessedCommand->GetTargetColumn()) {
        throw ShardokInternalErrorException(
                "Target coordinates mismatch between real and guessed state");
    }

    // For commands with odds (like FLEE), verify the odds match
    if (realCommand->HasOdds() != guessedCommand->HasOdds()) {
        throw ShardokInternalErrorException(
                "Odds presence mismatch between real and guessed state");
    }

    if (realCommand->HasOdds() && guessedCommand->HasOdds()) {
        if (realCommand->GetOddsPercentile() != guessedCommand->GetOddsPercentile()) {
            throw ShardokInternalErrorException(
                    "Odds percentile mismatch between real and guessed state");
        }
    }
}

auto ShardokAIClient::StandardChooseCommandIndex(
        const GameSettingsSPtr &settings,
        const GameStateW &guessedState,
        const CommandListSPtr &realAvailableCommands) const -> CommandChoiceResults {
    const auto settingsGetter = settings->GetGetter();
    const auto guessedEngine = ShardokEngine(settings, guessedState);
    const auto guessedCommands = guessedEngine.GetAvailableCommandsForAIPlayer(playerId);
    const auto commandCount = guessedCommands->size();

    // Calculate time budget based on game situation using new dynamic per-command settings
    const auto timeBudget =
            CalculateTimeBudget(playerId, settings, guessedState, commandCount, isAllAiBattle);

    // Configure MCTS based on proximity to enemy
    // When far from enemy: use AVERAGING with maxPlayerFlips=0 (single-player lookahead)
    //   - AVERAGING naturally penalizes longer paths through variance
    //   - No opponent nodes, so no one-bad-child problem
    // When close to enemy: use MINIMAX with maxPlayerFlips=1 (adversarial lookahead)
    //   - MINIMAX correctly models opponent choosing best response
    //   - Explores through one opponent turn for tactical accuracy
    auto adjustedMCTSConfig = mctsConfig;

    // For fair evaluation: simulate leaves to opponent's turn start (maxSimulationFlips=1)
    // This ensures all leaves are scored at the same game phase:
    // - Leaves at playerFlips=0 (still my turn): simulate through END_TURN to playerFlips=1
    // - Leaves at playerFlips=1 (opponent's turn): evaluate immediately
    // Result: consistent comparison of "what happens after I end my turn"
    // adjustedMCTSConfig.maxSimulatfixionFlips = 1;

    // adjustedMCTSConfig.maxPlayerFlips = 0;
    // if (timeBudget.isCloseToEnemy) {
    //     adjustedMCTSConfig.maxPlayerFlips = 1;
    //     adjustedMCTSConfig.backpropagationPolicy = mcts::MCTSBackpropagationPolicy::MINIMAX;
    //     if constexpr (kPerformanceLogging) {
    //         printf("MCTS Config: Close to enemy - using maxPlayerFlips=1, MINIMAX backprop\n");
    //     }
    // } else {
    //     adjustedMCTSConfig.maxPlayerFlips = 0;
    //     adjustedMCTSConfig.backpropagationPolicy = mcts::MCTSBackpropagationPolicy::AVERAGING;
    //     if constexpr (kPerformanceLogging) {
    //         printf("MCTS Config: Far from enemy - using maxPlayerFlips=0, AVERAGING backprop\n");
    //     }
    // }

    // Verify that the AI's guessed state produces the same available commands as reality
    if (commandCount != realAvailableCommands->size()) {
        fprintf(stderr,
                "AI command count mismatch: guessed=%zu real=%zu player=%d round=%d\n",
                commandCount,
                realAvailableCommands->size(),
                playerId,
                guessedState->current_round());

        const auto maxDump = std::max(commandCount, realAvailableCommands->size());
        for (size_t i = 0; i < maxDump; i++) {
            const auto *realName = i < realAvailableCommands->size()
                                           ? net::eagle0::shardok::common::CommandType_Name(
                                                     (*realAvailableCommands)[i]->GetCommandType())
                                                     .c_str()
                                           : "(none)";
            const auto *guessedName = i < commandCount
                                              ? net::eagle0::shardok::common::CommandType_Name(
                                                        (*guessedCommands)[i]->GetCommandType())
                                                        .c_str()
                                              : "(none)";
            const int realActor = i < realAvailableCommands->size()
                                          ? (*realAvailableCommands)[i]->GetActorUnitId()
                                          : -1;
            const int guessedActor =
                    i < commandCount ? (*guessedCommands)[i]->GetActorUnitId() : -1;
            const int realRow = i < realAvailableCommands->size()
                                        ? (*realAvailableCommands)[i]->GetTargetRow()
                                        : -1;
            const int realCol = i < realAvailableCommands->size()
                                        ? (*realAvailableCommands)[i]->GetTargetColumn()
                                        : -1;
            const int guessedRow = i < commandCount ? (*guessedCommands)[i]->GetTargetRow() : -1;
            const int guessedCol = i < commandCount ? (*guessedCommands)[i]->GetTargetColumn() : -1;
            fprintf(stderr,
                    "  [%zu] real: %s unit=%d (%d,%d)  guessed: %s unit=%d (%d,%d)\n",
                    i,
                    realName,
                    realActor,
                    realRow,
                    realCol,
                    guessedName,
                    guessedActor,
                    guessedRow,
                    guessedCol);
        }

        // Find positions in real commands that are missing from guessed commands
        // Group by unit to find the missing position per unit
        std::unordered_map<int, std::set<std::pair<int, int>>> realPositions;
        std::unordered_map<int, std::set<std::pair<int, int>>> guessedPositions;
        for (size_t i = 0; i < realAvailableCommands->size(); i++) {
            const auto &cmd = (*realAvailableCommands)[i];
            realPositions[cmd->GetActorUnitId()].emplace(
                    cmd->GetTargetRow(),
                    cmd->GetTargetColumn());
        }
        for (size_t i = 0; i < commandCount; i++) {
            const auto &cmd = (*guessedCommands)[i];
            guessedPositions[cmd->GetActorUnitId()].emplace(
                    cmd->GetTargetRow(),
                    cmd->GetTargetColumn());
        }
        for (const auto &[unitId, positions] : realPositions) {
            for (const auto &[row, col] : positions) {
                if (!guessedPositions[unitId].contains({row, col})) {
                    fprintf(stderr,
                            "  Missing from guessed: unit=%d pos=(%d,%d)",
                            unitId,
                            row,
                            col);
                    const Coords missingCoords(row, col);
                    const auto *occupant = guessedState.GetOccupant(missingCoords);
                    if (occupant) {
                        fprintf(stderr,
                                " -> guessed occupant: unit_id=%d player=%d status=%d "
                                "loc=(%d,%d) hidden=%d battalion_size=%d\n",
                                occupant->unit_id(),
                                occupant->player_id(),
                                static_cast<int>(occupant->status()),
                                occupant->location().row(),
                                occupant->location().column(),
                                occupant->hidden(),
                                occupant->battalion().size());
                    } else {
                        fprintf(stderr, " -> no occupant in guessed state\n");
                    }
                }
            }
        }

        // Dump all guessed state units for full picture
        fprintf(stderr, "  Guessed state units (%u total):\n", guessedState->units()->size());
        for (size_t i = 0; i < guessedState->units()->size(); i++) {
            const auto *u = guessedState->units()->Get(static_cast<unsigned int>(i));
            fprintf(stderr,
                    "    unit_id=%d player=%d status=%d loc=(%d,%d) hidden=%d "
                    "starting_pos_idx=%d\n",
                    u->unit_id(),
                    u->player_id(),
                    static_cast<int>(u->status()),
                    u->location().row(),
                    u->location().column(),
                    u->hidden(),
                    u->starting_position_index());
        }

        throw ShardokInternalErrorException(
                "AI command count mismatch: guessed=" + std::to_string(commandCount) +
                " real=" + std::to_string(realAvailableCommands->size()));
    }
    for (size_t i = 0; i < commandCount; i++) {
        CheckCommand((*realAvailableCommands)[i], (*guessedCommands)[i]);
    }

    // Extract values directly from settings for strategy selection
    const auto maxRounds = settingsGetter.Backing().max_rounds();
    const auto braveWaterCost = settingsGetter.Backing().brave_water_action_point_cost();
    const auto battalionTypeGetter = [&settingsGetter](BattalionTypeId typeId) {
        return settingsGetter.GetBattalionType(typeId);
    };

    // Create scorer for actual scoring during search - type selected at construction
    std::unique_ptr<AIScoreCalculator> scorer;
    switch (scoringCalculatorType) {
        case ScoringCalculatorType::NORMALIZED:
            scorer = MakeNormalizedAIScoreCalculator(settingsGetter, apdCache, alCache);
            break;
        case ScoringCalculatorType::MCTS_OPTIMIZED:
            scorer = MakeMCTSOptimizedAIScoreCalculator(settingsGetter, apdCache, alCache);
            break;
        case ScoringCalculatorType::STANDARD:
        default: scorer = MakeStandardAIScoreCalculator(settingsGetter, apdCache, alCache); break;
    }

    // Determine strategy once for consistent scoring throughout iterative deepening
    const auto castleCoords = AllCastleCoords(guessedState->hex_map());
    const AIStrategy strategy = isDefender ? AIDefenderStrategySelector::BestDefenderStrategy(
                                                     guessedState,
                                                     castleCoords,
                                                     maxRounds,
                                                     apdCache,
                                                     battalionTypeGetter)
                                           : AIAttackerStrategySelector::BestAttackerStrategy(
                                                     playerId,
                                                     guessedState,
                                                     castleCoords,
                                                     maxRounds,
                                                     apdCache,
                                                     alCache,
                                                     battalionTypeGetter,
                                                     braveWaterCost,
                                                     waterCrossingCommandChooser,
                                                     realAvailableCommands);

    // AI implementation chosen at runtime via constructor parameter
    IterativeDeepeningAI::SearchResult search_result;

    if (aiAlgorithmType == AIAlgorithmType::MCTS) {
#ifdef ENABLE_MCTS_DEBUG_DUMP
        // Set unique debug dump path for each action using timestamp
        const auto now = std::chrono::system_clock::now();
        const auto nowTime = std::chrono::system_clock::to_time_t(now);
        const auto nowMs =
                std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()) %
                1000;

        std::ostringstream pathStream;
        pathStream << "/tmp/shardok_debug_"
                   << std::put_time(std::localtime(&nowTime), "%Y%m%d_%H%M%S") << "_"
                   << std::setfill('0') << std::setw(3) << nowMs.count() << "_p"
                   << static_cast<int>(playerId) << ".txt";
        adjustedMCTSConfig.debugDumpPath = pathStream.str();

        // Also dump the game state to a file for reproduction
        std::ostringstream statePathStream;
        statePathStream << "/tmp/shardok_state_"
                        << std::put_time(std::localtime(&nowTime), "%Y%m%d_%H%M%S") << "_"
                        << std::setfill('0') << std::setw(3) << nowMs.count() << "_p"
                        << static_cast<int>(playerId) << ".bin";
        const std::string statePath = statePathStream.str();

        // Write the flatbuffer game state to file using SaveTo method
        if (guessedState.SaveTo(statePath)) {
            printf("Game state dumped to: %s\n", statePath.c_str());
        } else {
            printf("Failed to dump game state to: %s\n", statePath.c_str());
        }
#endif  // ENABLE_MCTS_DEBUG_DUMP

        // Using Monte Carlo Tree Search AI (with abstraction layer)
        ShardokMCTSAI ai(
                playerId,
                isDefender,
                strategy,
                castleCoords,
                *scorer,
                apdCache,
                alCache,
                adjustedMCTSConfig);
        search_result = ai.Search(settings, guessedState, timeBudget);
    } else {
        // Using Iterative Deepening AI (default)
        IterativeDeepeningAI ai(
                playerId,
                isDefender,
                strategy,
                castleCoords,
                *scorer,
                apdCache,
                battalionTypeGetter);
        search_result =
                ai.IterativeSearch(settings, guessedState, realAvailableCommands, timeBudget);
    }

    CommandChoiceResults result{};
    result.chosenIndex = search_result.bestCommandIndex;
    result.availableCommandCount = search_result.availableCommandCount;
    result.depthAchieved = search_result.depthAchieved;
    result.commandCountEvaluated = search_result.commandCountEvaluated;
    result.completionReason = search_result.completionReason;

    if constexpr (kPerformanceLogging) {
        if (result.commandCountEvaluated < result.availableCommandCount) {
            printf("ID AI: Depth %d - evaluated %lu/%zu commands\n",
                   result.depthAchieved,
                   result.commandCountEvaluated,
                   result.availableCommandCount);
        }
        const auto chosenCommandType =
                (*realAvailableCommands)[result.chosenIndex]->GetCommandType();
        printf("ID AI: Search complete - achieved depth %d for best command %zu (%s)\n",
               result.depthAchieved,
               result.chosenIndex,
               net::eagle0::shardok::common::CommandType_Name(chosenCommandType).c_str());

        fflush(stdout);
    }

    return result;
}

auto ShardokAIClient::LateRoundAttackerChooseCommandIndex(
        const GameSettingsSPtr &settings,
        const GameStateW &guessedState,
        const CommandListSPtr &realAvailableCommands) const -> CommandChoiceResults {
    if (const auto dismissCommand = std::ranges::find_if(
                *realAvailableCommands,
                [](const CommandSPtr &cmd) {
                    return cmd->GetCommandType() ==
                           net::eagle0::shardok::common::DISMISS_UNIT_COMMAND;
                });
        dismissCommand == realAvailableCommands->end()) {
        return StandardChooseCommandIndex(settings, guessedState, realAvailableCommands);
    } else {
        CommandChoiceResults results{};
        results.chosenIndex =
                static_cast<size_t>(std::distance(realAvailableCommands->begin(), dismissCommand));
        results.availableCommandCount = realAvailableCommands->size();
        results.depthAchieved = 1;          // Simple heuristic choice
        results.commandCountEvaluated = 1;  // Only evaluated one command type
        results.completionReason =
                EvaluationCompletionReason::RAN_OUT_OF_COMMANDS;  // Heuristic choice
        return results;
    }
}

auto ShardokAIClient::FinalRoundAttackerChooseCommandIndex(
        const GameSettingsSPtr &settings,
        const GameStateW &guessedState,
        const CommandListSPtr &realAvailableCommands) const -> CommandChoiceResults {
    const auto fleeCommand =
            std::ranges::find_if(*realAvailableCommands, [](const CommandSPtr &cmd) {
                return cmd->GetCommandType() == net::eagle0::shardok::common::FLEE_COMMAND;
            });

    if (fleeCommand == realAvailableCommands->end()) {
        return LateRoundAttackerChooseCommandIndex(settings, guessedState, realAvailableCommands);
    }

    // Extract values directly from settings for flee decision evaluation
    const auto settingsGetter = settings->GetGetter();
    const auto maxRounds = settingsGetter.Backing().max_rounds();
    const auto minimumFleeOddsThreshold = settingsGetter.Backing().ai_minimum_flee_odds_threshold();
    const auto desperateFleeThreshold = settingsGetter.Backing().ai_desperate_flee_threshold();

    // Use the flee decision calculator
    const auto fleeDecision = AIFleeDecisionCalculator::EvaluateFleeVsFight(
            playerId,
            guessedState,
            realAvailableCommands,
            fleeCommand,
            maxRounds,
            minimumFleeOddsThreshold,
            desperateFleeThreshold,
#ifdef DEBUG_FLEE_DECISIONS
            true  // Enable debug logging
#else
            false
#endif
    );

    if (fleeDecision.shouldFlee) {
        CommandChoiceResults results{};
        results.chosenIndex = fleeDecision.commandIndex;
        results.availableCommandCount = realAvailableCommands->size();
        results.depthAchieved = 1;          // Heuristic choice
        results.commandCountEvaluated = 1;  // Only evaluated one command type
        results.completionReason = EvaluationCompletionReason::RAN_OUT_OF_COMMANDS;
        return results;
    } else {
        // Fight instead of flee
        return StandardChooseCommandIndex(settings, guessedState, realAvailableCommands);
    }
}

auto ShardokAIClient::ChooseCommandIndex(
        const GameSettingsSPtr &settings,
        const GameStateView &gsv,
        const CommandListSPtr &realAvailableCommands) const -> CommandChoiceResults {
    static int typeChosenCount[net::eagle0::shardok::common::CommandType_MAX + 1];
    static int totalChoices = 0;

    CommandChoiceResults results{};

    const auto guessedState = GameStateGuesser::GuessedState(playerId, settings->GetGetter(), gsv);

    if (const int roundsRemaining = RoundsRemaining(settings, gsv);
        !isDefender && roundsRemaining <= 1) {
        results =
                FinalRoundAttackerChooseCommandIndex(settings, guessedState, realAvailableCommands);
    } else if (!isDefender && roundsRemaining <= 3) {
        results =
                LateRoundAttackerChooseCommandIndex(settings, guessedState, realAvailableCommands);
    } else {
        results = StandardChooseCommandIndex(settings, guessedState, realAvailableCommands);
    }

    const auto chosenType = (*realAvailableCommands)[results.chosenIndex]->GetCommandType();
    typeChosenCount[static_cast<int>(chosenType)]++;
    totalChoices++;

    if (totalChoices % 1000 == 0) {
        printf("TYPES CHOSEN:\n");
        vector<std::pair<int, CommandType>> choices{};
        for (int i = 0; i <= net::eagle0::shardok::common::CommandType_MAX; i++) {
            if (typeChosenCount[i] != 0) {
                choices.emplace_back(typeChosenCount[i], static_cast<CommandType>(i));
            }
        }

        std::ranges::sort(choices);
        std::ranges::reverse(choices);
        for (const auto &[index, choice] : choices) {
            printf("%5d %s\n", index, CommandType_Name(choice).c_str());
        }
        printf("\n\n");
    }

    return results;
}

auto ShardokAIClient::ChooseCommandIndex(const ShardokEngine &engine) const
        -> CommandChoiceResults {
    if (const auto &availableCommands = engine.GetAvailableCommandsForAIPlayer(playerId);
        availableCommands->empty()) {
        printf("no commands for player %d\n", playerId);
        throw ShardokInternalErrorException(
                "Asked to choose a command, but there are none available");
    } else {
        const auto &settings = engine.GetGameSettings();
        const auto &gsv = engine.GetGameStateView(GetPlayerId());

        const auto results = ChooseCommandIndex(settings, gsv, availableCommands);
        apdCache->ConsolidateThreadLocalCache_Racy();
        return results;
    }
}

}  // namespace shardok
