//
//  ShardokGameController.cpp
//  eagle0_engine
//
//  Created by Dan Crosby on 7/9/17.
//  Copyright © 2017 none. All rights reserved.
//

#include "ShardokGameController.hpp"

#include <execinfo.h>

#include <algorithm>
#include <iostream>
#include <iterator>
#include <ranges>
#include <sstream>
#include <string>
#include <thread>

#include "src/main/cpp/net/eagle0/common/ShardokLatencyTrace.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/ShardokAIClient.hpp"
#include "src/main/cpp/net/eagle0/shardok/ai/StartingPositionRecommender.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/protobuf/net/eagle0/shardok/api/game_state_view.pb.h"

namespace shardok {

auto PlayerBattleProgressFrom(const net::eagle0::shardok::storage::fb::GameState &gameState)
        -> std::vector<PlayerBattleProgress> {
    std::vector<PlayerBattleProgress> progress;
    progress.reserve(gameState.player_infos()->size());
    for (const auto *player : *gameState.player_infos()) {
        int32_t troopCount = 0;
        bool isRetreating = false;
        for (const auto *unit : *gameState.units()) {
            if (unit->player_id() != player->player_id()) continue;
            switch (unit->status()) {
                case net::eagle0::shardok::storage::fb::UnitStatus_NORMAL_UNIT:
                case net::eagle0::shardok::storage::fb::UnitStatus_RESERVE_UNIT:
                case net::eagle0::shardok::storage::fb::UnitStatus_NEVER_ENTERED_UNIT:
                    troopCount += unit->battalion().size();
                    break;
                case net::eagle0::shardok::storage::fb::UnitStatus_FLED_UNIT:
                    isRetreating = true;
                    break;
                default: break;
            }
        }
        progress.push_back(PlayerBattleProgress{player->player_id(), troopCount, isRetreating});
    }
    return progress;
}
using net::eagle0::common::ScopedShardokLatencyTrace;

using net::eagle0::shardok::common::GameStatus;

namespace {

[[nodiscard]] auto DescribeCommandForLog(const CommandSPtr &command) -> std::string {
    if (!command) { return "<null command>"; }

    std::ostringstream out;
    out << "type=" << command->GetCommandType()
        << " player=" << static_cast<int>(command->GetPlayerId())
        << " actor=" << command->GetActorUnitId() << " target=("
        << static_cast<int>(command->GetTargetRow()) << ','
        << static_cast<int>(command->GetTargetColumn()) << ')';

    if (command->HasOdds()) { out << " odds_percentile=" << command->GetOddsPercentile(); }

    try {
        out << " proto={" << command->GetCommandProto().ShortDebugString() << '}';
    } catch (const std::exception &e) { out << " proto_error={" << e.what() << '}'; }

    return out.str();
}

[[nodiscard]] auto DescribeCommandsForLog(const CommandList &commands) -> std::string {
    std::ostringstream out;
    out << "available_command_count=" << commands.size();
    for (size_t i = 0; i < commands.size(); ++i) {
        out << "\n  [" << i << "] " << DescribeCommandForLog(commands[i]);
    }
    return out.str();
}

[[nodiscard]] auto DescribeCurrentStateForLog(const ShardokEngine &engine) -> std::string {
    const auto *state = engine.GetCurrentGameState().Get();
    if (state == nullptr) { return "state=<null>"; }

    const auto vectorSize = [](const auto *vector) -> size_t {
        return vector == nullptr ? 0 : vector->size();
    };

    std::ostringstream out;
    out << "state={round=" << static_cast<int>(state->current_round())
        << " current_player=" << static_cast<int>(state->current_player())
        << " history_count=" << engine.GetUnfilteredHistoryCount();

    if (const auto *status = state->status(); status != nullptr) {
        out << " status=" << static_cast<int>(status->state()) << " end_condition=";
        if (const auto *endCondition = status->end_game_condition(); endCondition != nullptr) {
            out << "{type=" << static_cast<int>(endCondition->victory_type())
                << " victory_details=" << static_cast<int>(endCondition->victory_details())
                << " draw_details=" << static_cast<int>(endCondition->draw_details()) << '}';
        } else {
            out << "<null>";
        }
    } else {
        out << " status=<null>";
    }

    out << " units=" << vectorSize(state->units())
        << " player_infos=" << vectorSize(state->player_infos());

    if (const auto *hexMap = state->hex_map(); hexMap != nullptr) {
        out << " map=" << hexMap->row_count() << 'x' << hexMap->column_count();
    } else {
        out << " map=<null>";
    }

    out << '}';
    return out.str();
}

}  // namespace

class InvalidTokenException : public std::exception {};

class MismatchedPlayerIdException : public std::exception {
private:
    PlayerId shardokPlayerId;
    int eagleFactionId;

    std::string message;

public:
    MismatchedPlayerIdException(const PlayerId shardokPid, const int eagleFid)
        : shardokPlayerId(shardokPid),
          eagleFactionId(eagleFid) {
        message = "Mismatched playerId " + std::to_string(shardokPlayerId) +
                  " with eagle factionId " + std::to_string(eagleFactionId);
    }

    [[nodiscard]] auto what() const noexcept -> const char * override { return message.c_str(); }
};

class MissingPlayerIdException : public std::exception {
private:
    PlayerId shardokPlayerId;

    std::string message;

public:
    MissingPlayerIdException(const PlayerId shardokPid) : shardokPlayerId(shardokPid) {
        message = "Missing playerId " + std::to_string(shardokPlayerId);
    }

    [[nodiscard]] auto what() const noexcept -> const char * override { return message.c_str(); }
};

std::vector<std::shared_ptr<ShardokAIClient>> ShardokGameController::MakeAIClients(
        const std::unique_ptr<ShardokEngine> &e) {
    std::vector<std::shared_ptr<ShardokAIClient>> aic;

    mcts::MCTSConfig mctsConfig;
    mctsConfig.backpropagationPolicy = mcts::MCTSBackpropagationPolicy::MINIMAX;
    mctsConfig.maxPlayerFlips = 1;
    mctsConfig.maxSimulationFlips = 1;
    mctsConfig.simulationPolicy = mcts::MCTSSimulationPolicy::WEIGHTED_HEURISTIC;

    // Check if all players are AI - if so, use faster time budgets
    const auto &playerInfos = e->GetPlayerInfos();
    aic.reserve(playerInfos.size());
    const bool isAllAiBattle =
            std::ranges::all_of(playerInfos, [](const auto &pi) { return pi.is_ai(); });

    for (const auto &pi : playerInfos) {
        if (pi.is_ai()) {
            auto newClient = std::make_shared<ShardokAIClient>(
                    pi.player_id(),
                    pi.is_defender(),
                    isAllAiBattle,
                    e->GetCurrentGameState()->hex_map(),
                    e->GetGameSettings()->GetGetter(),
                    AIAlgorithmType::ITERATIVE_DEEPENING,
                    ScoringCalculatorType::STANDARD,
                    mctsConfig);

            // MCTS config is dynamically adjusted in ShardokAIClient based on proximity:
            // - Far from enemy: AVERAGING with maxPlayerFlips=0 (natural wandering prevention)
            // - Close to enemy: MINIMAX with maxPlayerFlips=1 (adversarial lookahead)

            aic.push_back(newClient);
        }
    }

    return aic;
}

void ShardokGameController::LockedNotifyClients() const { updateCondition.notify_all(); }

auto ShardokGameController::LockedAIClientForPid(PlayerId pid) const
        -> std::shared_ptr<ShardokAIClient> {
    const auto it = std::ranges::find_if(aiClients, [pid](const auto &client) {
        return client->GetPlayerId() == pid;
    });
    return (it != aiClients.end()) ? *it : nullptr;
}

void ShardokGameController::DoAIThread() {
    std::cout << "Entering AI thread." << '\n';
    if (aiClients.empty()) { std::cout << "No AI players, exiting AI thread." << '\n'; }

    while (aiThreadKeepGoing) {
        // Phase 1: Gather data for AI decision (brief lock)
        std::shared_ptr<ShardokAIClient> aiClient;
        PlayerId playerId;
        GameSettingsSPtr settings;
        net::eagle0::shardok::api::GameStateView gsv;
        CommandListSPtr availableCommands;
        size_t expectedHistoryCount;

        {
            std::unique_lock lk(masterLock);

            if (engine->GameIsOver()) {
                aiThreadKeepGoing = false;
                continue;
            }

            playerId = engine->GetCurrentPlayerId();
            aiClient = LockedAIClientForPid(playerId);

            if (!aiClient) {
                // Not an AI player's turn - wait for signal
                aiCondition.wait(lk);
                aiThreadKeepGoing = aiThreadKeepGoing && !engine->GameIsOver();
                continue;
            }

            // Get copies of everything the AI needs
            settings = engine->GetGameSettings();
            gsv = engine->GetGameStateView(playerId);
            availableCommands = engine->GetAvailableCommandsForAIPlayer(playerId);
            expectedHistoryCount = engine->GetUnfilteredHistoryCount();
        }
        // Lock released - polls can now get through

        if (availableCommands->empty()) {
            std::cout << "no commands for player " << playerId << '\n';
            continue;
        }

        std::string aiPhase = "choosing command";
        int64_t chosenIndex = -1;
        try {
            // Phase 2: AI thinks (NO LOCK - this is the slow part)
            auto results = [&] {
                ScopedShardokLatencyTrace trace(
                        "ai_choose_command",
                        cachedGameId,
                        {{"player_id", std::to_string(playerId)},
                         {"round", std::to_string(gsv.current_round())},
                         {"available_command_count", std::to_string(availableCommands->size())}});
                return aiClient->ChooseCommandIndex(settings, gsv, availableCommands, std::nullopt);
            }();
            chosenIndex = results.chosenIndex;

            // Phase 3: Post the command (brief lock)
            {
                std::unique_lock lk(masterLock);

                // Verify state hasn't changed while we were thinking
                if (engine->GetUnfilteredHistoryCount() != expectedHistoryCount) {
                    // State changed (e.g., human posted command) - re-evaluate
                    std::cout << "AI: State changed while thinking, re-evaluating" << '\n';
                    continue;
                }

                if (engine->GameIsOver()) {
                    aiThreadKeepGoing = false;
                    continue;
                }

                {
                    aiPhase = "posting chosen command";
                    ScopedShardokLatencyTrace trace(
                            "ai_engine_post_command",
                            cachedGameId,
                            {{"player_id", std::to_string(playerId)},
                             {"round", std::to_string(gsv.current_round())},
                             {"chosen_index", std::to_string(results.chosenIndex)}});
                    engine->PostCommand(playerId, results.chosenIndex);
                }
                LockedNotifyClients();
                {
                    aiPhase = "posting forced follow-up commands";
                    ScopedShardokLatencyTrace trace(
                            "ai_engine_post_forced_commands",
                            cachedGameId,
                            {{"player_id", std::to_string(playerId)},
                             {"round", std::to_string(gsv.current_round())}});
                    engine->PostWhileCurrentPlayerHasOnlyOneOption(nullptr);
                }
                LockedNotifyClients();
                aiThreadKeepGoing = !engine->GameIsOver();
            }
        } catch (const std::exception &e) {
            std::ostringstream error;
            error << "AI error in game " << cachedGameId << ", player " << playerId << ", round "
                  << gsv.current_round() << ": " << e.what();

            std::ostringstream context;
            context << "AI failure context: phase=" << aiPhase
                    << ", expected_history_count=" << expectedHistoryCount
                    << ", current_history_count=" << engine->GetUnfilteredHistoryCount()
                    << ", chosen_index=" << chosenIndex
                    << ", current_player=" << static_cast<int>(engine->GetCurrentPlayerId()) << ", "
                    << DescribeCurrentStateForLog(*engine);

            const std::string commandsContext = availableCommands
                                                        ? DescribeCommandsForLog(*availableCommands)
                                                        : "available_command_count=<cleared>";

            std::cerr << "AI thread exception in game " << cachedGameId << ", player " << playerId
                      << ", round " << gsv.current_round() << ": " << e.what() << '\n';
            std::cerr << context.str() << '\n';
            std::cerr << commandsContext << '\n';

            void *backtraceArray[20];
            const int backtraceSize = backtrace(backtraceArray, 20);
            backtrace_symbols_fd(backtraceArray, backtraceSize, STDERR_FILENO);

            aiThreadErrorMessage = error.str() + "\n" + context.str() + "\n" + commandsContext;
            aiThreadFailed.store(true);
            updateCondition.notify_all();
            break;
        }
    }
    std::cout << "Exiting AI thread." << '\n';
}

ShardokGameController::~ShardokGameController() {
    {
        std::unique_lock lk(masterLock);
        aiThreadKeepGoing = false;
        aiCondition.notify_one();
    }
    aiThread.join();
}

void CheckFactionId(
        const std::unique_ptr<ShardokEngine> &engine,
        const PlayerId shardokPlayerId,
        const int eagleFactionId) {
    bool found = false;
    for (const auto &pi : *engine->GetCurrentGameState()->player_infos()) {
        if (pi->player_id() != shardokPlayerId) continue;
        found = true;
        if (pi->eagle_faction_id() != eagleFactionId) {
            throw MismatchedPlayerIdException(shardokPlayerId, eagleFactionId);
        }
    }

    if (!found) throw MissingPlayerIdException(shardokPlayerId);
}

void ShardokGameController::PostCommand(
        const PlayerId shardokPlayerId,
        const int eagleFactionId,
        const int64_t token,
        const int64_t commandIndex,
        const std::optional<int32_t> roll) {
    ScopedShardokLatencyTrace totalTrace(
            "controller_post_command_total",
            cachedGameId,
            {{"player_id", std::to_string(shardokPlayerId)},
             {"eagle_faction_id", std::to_string(eagleFactionId)},
             {"token", std::to_string(token)},
             {"command_index", std::to_string(commandIndex)}});
    std::scoped_lock<std::mutex> guard(masterLock);

    CheckFactionId(engine, shardokPlayerId, eagleFactionId);

    const auto expectedToken = static_cast<int64_t>(engine->GetUnfilteredHistoryCount());
    if (token < expectedToken) {
        std::cout << "Double token in postCommand" << '\n';
        // The client is missing some updates; probably it's a double-submit
        // Let's just ignore and see how it goes
    } else if (token > expectedToken) {
        // This shouldn't be possible
        throw InvalidTokenException();
    } else {
        {
            ScopedShardokLatencyTrace trace(
                    "engine_post_command",
                    cachedGameId,
                    {{"player_id", std::to_string(shardokPlayerId)},
                     {"command_index", std::to_string(commandIndex)}});
            engine->PostCommand(shardokPlayerId, commandIndex, nullptr, roll);
        }
        LockedNotifyClients();
        {
            ScopedShardokLatencyTrace trace(
                    "engine_post_forced_commands",
                    cachedGameId,
                    {{"player_id", std::to_string(shardokPlayerId)}});
            engine->PostWhileCurrentPlayerHasOnlyOneOption(nullptr);
        }
    }

    LockedNotifyClients();
    aiCondition.notify_one();
}

void ShardokGameController::PostPlacementCommands(
        const PlayerId shardokPlayerId,
        const int eagleFactionId,
        const int64_t token,
        const std::vector<UnitPlacementInfo> &infos) {
    ScopedShardokLatencyTrace totalTrace(
            "controller_post_placement_commands_total",
            cachedGameId,
            {{"player_id", std::to_string(shardokPlayerId)},
             {"eagle_faction_id", std::to_string(eagleFactionId)},
             {"token", std::to_string(token)},
             {"placement_count", std::to_string(infos.size())}});
    {
        std::scoped_lock<std::mutex> guard(masterLock);

        CheckFactionId(engine, shardokPlayerId, eagleFactionId);

        const auto expectedToken = static_cast<int64_t>(engine->GetUnfilteredHistoryCount());
        if (token < expectedToken) {
            std::cout << "Double token in postPlacementCommands" << '\n';
            // The client is missing some updates; probably it's a double-submit
            // Let's just ignore and see how it goes
        } else if (token > expectedToken) {
            // This shouldn't be possible
            throw InvalidTokenException();
        } else {
            // The token matches, execute the commands
            {
                ScopedShardokLatencyTrace trace(
                        "engine_post_placement_commands",
                        cachedGameId,
                        {{"player_id", std::to_string(shardokPlayerId)},
                         {"placement_count", std::to_string(infos.size())}});
                engine->PostPlacementCommands(shardokPlayerId, infos);
            }
            {
                ScopedShardokLatencyTrace trace(
                        "engine_post_finished_placement",
                        cachedGameId,
                        {{"player_id", std::to_string(shardokPlayerId)}});
                engine->PostFinishedPlacementCommand(shardokPlayerId);
            }
            {
                ScopedShardokLatencyTrace trace(
                        "engine_post_placement_forced_commands",
                        cachedGameId,
                        {{"player_id", std::to_string(shardokPlayerId)}});
                engine->PostWhileCurrentPlayerHasOnlyOneOption(nullptr);
            }
        }

        LockedNotifyClients();
        aiCondition.notify_one();
    }
}

auto ShardokGameController::GetCurrentGameStateBytes() -> byte_vector {
    std::scoped_lock<std::mutex> guard(masterLock);
    return engine->GetCurrentGameStateBytes();
}

auto ShardokGameController::GetUpdates(const int64_t startingActionId) -> AllUpdates {
    ScopedShardokLatencyTrace totalTrace(
            "controller_get_updates_total",
            cachedGameId,
            {{"starting_action_id", std::to_string(startingActionId)}});
    AllUpdates updates{};
    std::unique_ptr<ShardokEngine> recommendationEngine;
    std::shared_ptr<AvailableCommands> recommendationCommands;
    PlayerId recommendationPlayerId = -1;
    std::pair<int64_t, PlayerId> recommendationCacheKey;

    incomingRegistrations++;
    std::vector<net::eagle0::shardok::storage::ShardokActionWithResultingState> awrs;
    {
        std::unique_lock<std::mutex> guard(masterLock);

        updates.newUnfilteredCount = engine->GetUnfilteredHistoryCount();

        {
            ScopedShardokLatencyTrace trace(
                    "engine_get_game_history",
                    cachedGameId,
                    {{"starting_action_id", std::to_string(startingActionId)}});
            awrs = engine->GetGameHistory(startingActionId);
        }
        incomingRegistrations--;

        // Note: Previously this had a 5-second wait for AI players to support long-polling.
        // With streaming (WaitForUpdatesAndPush), the caller already waits for updates,
        // so this wait is no longer needed. GetGameStatus is deprecated in favor of streaming.

        updates.mainResults.reserve(awrs.size());
        std::ranges::transform(
                awrs,
                std::back_inserter(updates.mainResults),
                [](const ShardokActionWithResultingState &a) { return a.action_result(); });

        const auto playerInfos = engine->GetPlayerInfos();
        updates.filteredResults.reserve(playerInfos.size() + 1);
        for (const auto &pi : playerInfos) {
            if (pi.is_ai()) continue;

            auto fid = pi.eagle_faction_id();

            std::vector<net::eagle0::shardok::api::ActionResultView> actionResultViews;
            {
                ScopedShardokLatencyTrace trace(
                        "engine_filter_new_results",
                        cachedGameId,
                        {{"player_id", std::to_string(pi.player_id())},
                         {"eagle_faction_id", std::to_string(fid)},
                         {"starting_action_id", std::to_string(startingActionId)}});
                actionResultViews = engine->FilterNewResults(pi.player_id(), startingActionId);
            }

            std::shared_ptr<AvailableCommands> acs = nullptr;

            const auto availableCommands = [&] {
                ScopedShardokLatencyTrace trace(
                        "engine_get_available_command_protos",
                        cachedGameId,
                        {{"player_id", std::to_string(pi.player_id())}});
                return engine->GetAvailableCommandProtos(pi.player_id(), true);
            }();
            const auto previewCommands = [&] {
                ScopedShardokLatencyTrace trace(
                        "engine_get_preview_commands",
                        cachedGameId,
                        {{"player_id", std::to_string(pi.player_id())}});
                return engine->GetPreviewCommands(pi.player_id());
            }();

            if (!availableCommands.empty() || !previewCommands.empty()) {
                acs = std::make_shared<AvailableCommands>();
                acs->mutable_current_command()->Reserve(availableCommands.size());
                *acs->mutable_current_command() = {
                        std::begin(availableCommands),
                        std::end(availableCommands)};
                acs->mutable_preview_command()->Reserve(previewCommands.size());
                *acs->mutable_preview_command() = {
                        std::begin(previewCommands),
                        std::end(previewCommands)};

                if (engine->GetGameStatus()->state() ==
                            net::eagle0::shardok::storage::fb::GameStatus_::State_SET_UP &&
                    engine->GetCurrentPlayerId() == pi.player_id() && !availableCommands.empty()) {
                    recommendationEngine = std::make_unique<ShardokEngine>(*engine, false);
                    recommendationCommands = acs;
                    recommendationPlayerId = pi.player_id();
                    recommendationCacheKey = {updates.newUnfilteredCount, recommendationPlayerId};
                }
            }

            updates.filteredResults.emplace_back(fid, actionResultViews, acs);
        }

        // Per-watcher filtered views. A watcher is a non-participant Eagle
        // faction allied to one or more participants — they get a view that
        // treats their allied participants' units as fully visible (hero
        // stats, hidden-from-opponents actions like meteor casts, etc.).
        // Participants whose allies list named this watcher are passed in as
        // alliedPids so UnitFilter / ActionResultFilter reveal those units to
        // the watcher.
        for (const auto &[watcherFid, allyPids] : watcherAllies) {
            // Watchers can't post commands, so availableCommands is null.
            updates.filteredResults.emplace_back(
                    watcherFid,
                    [&] {
                        ScopedShardokLatencyTrace trace(
                                "engine_filter_new_results_for_watcher",
                                cachedGameId,
                                {{"eagle_faction_id", std::to_string(watcherFid)},
                                 {"starting_action_id", std::to_string(startingActionId)}});
                        return engine->FilterNewResultsForWatcher(-1, allyPids, startingActionId);
                    }(),
                    nullptr);
        }

        updates.filteredResults.emplace_back(
                -1,
                [&] {
                    ScopedShardokLatencyTrace trace(
                            "engine_filter_new_results_all",
                            cachedGameId,
                            {{"starting_action_id", std::to_string(startingActionId)}});
                    return engine->FilterNewResults(-1, startingActionId);
                }(),
                nullptr);

        auto gameStateBytes = [&] {
            ScopedShardokLatencyTrace trace("engine_get_current_game_state_bytes", cachedGameId);
            return engine->GetCurrentGameStateBytes();
        }();
        updates.currentGameState.swap(gameStateBytes);

        // Extract battle progress data from the FlatBuffer game state
        const auto *gs = engine->GetCurrentGameState().Get();
        updates.currentRound = gs->current_round();
        updates.playerBattleProgress = PlayerBattleProgressFrom(*gs);
    }

    if (recommendationEngine) {
        std::scoped_lock<std::mutex> recommendationGuard(recommendationLock);
        auto cached = recommendedStartingPositionsCache.find(recommendationCacheKey);
        if (cached == recommendedStartingPositionsCache.end()) {
            cached = recommendedStartingPositionsCache
                             .emplace(
                                     recommendationCacheKey,
                                     RecommendStartingPositions(
                                             *recommendationEngine,
                                             recommendationPlayerId))
                             .first;
        }
        recommendationCommands->mutable_recommended_starting_positions()->Reserve(
                cached->second.size());
        *recommendationCommands->mutable_recommended_starting_positions() = {
                cached->second.begin(),
                cached->second.end()};
    }

    return updates;
}

auto ShardokGameController::GameIsOver() -> bool {
    std::scoped_lock<std::mutex> guard(masterLock);

    const auto &gameStatus = engine->GetGameStatus();
    return ShardokEngine::GameIsOver(gameStatus);
}

auto ShardokGameController::GameOverStatus() -> net::eagle0::shardok::common::GameStatus {
    std::scoped_lock<std::mutex> guard(masterLock);

    return fb::ToProto(engine->GetGameStatus());
}

auto ShardokGameController::EndGameUnits()
        -> std::vector<net::eagle0::shardok::storage::ResolvedUnit> {
    std::scoped_lock<std::mutex> guard(masterLock);
    return engine->EndGameUnits();
}

auto ShardokGameController::ResolvedPlayerInfos()
        -> std::vector<net::eagle0::shardok::common::PlayerInfo> {
    std::scoped_lock<std::mutex> guard(masterLock);
    return engine->GetPlayerInfos();
}

void ShardokGameController::RegisterSubscriber(std::shared_ptr<StreamSubscriber> subscriber) {
    std::scoped_lock<std::mutex> guard(subscriberLock);
    subscribers.push_back(subscriber);
}

void ShardokGameController::UnregisterSubscriber(const StreamSubscriber *subscriber) {
    std::scoped_lock<std::mutex> guard(subscriberLock);
    subscribers.erase(
            std::remove_if(
                    subscribers.begin(),
                    subscribers.end(),
                    [subscriber](const std::weak_ptr<StreamSubscriber> &weakSub) {
                        auto sub = weakSub.lock();
                        return !sub || sub.get() == subscriber;
                    }),
            subscribers.end());
}

auto ShardokGameController::WaitForUpdatesAndPush(
        std::shared_ptr<StreamSubscriber> subscriber,
        int64_t startingActionId) -> WaitResult {
    int64_t lastPushedActionId = startingActionId;

    while (subscriber->IsActive()) {
        bool gameOver = false;
        bool tutorialReset = false;
        GameOverInfo gameOverInfo{};

        {
            std::unique_lock<std::mutex> guard(masterLock);

            // Wait for updates, game over, or AI thread failure
            updateCondition.wait(guard, [this, lastPushedActionId] {
                return engine->GetUnfilteredHistoryCount() >
                               static_cast<size_t>(lastPushedActionId) ||
                       engine->GameIsOver() || aiThreadFailed.load();
            });

            if (!subscriber->IsActive()) { return WaitResult::DISCONNECTED; }

            if (aiThreadFailed.load()) { return WaitResult::DISCONNECTED; }

            gameOver = engine->GameIsOver();

            if (gameOver) {
                // Check if this is a tutorial battle where the defender lost
                tutorialReset = engine->IsTutorialBattleEnabled() && engine->DidDefenderLose();

                if (!tutorialReset) {
                    gameOverInfo.gameStatus = fb::ToProto(engine->GetGameStatus());
                    gameOverInfo.playerInfos = engine->GetPlayerInfos();
                    gameOverInfo.endGameUnits = engine->EndGameUnits();
                }
            }
        }
        // Lock released - GetUpdates will acquire its own lock

        // IMPORTANT: Send any pending updates FIRST, including the final actions
        // that caused the game to end. Previously, we returned early on gameOver
        // without sending these final updates, causing the client to never see
        // the last few battle actions.
        AllUpdates updates = GetUpdates(lastPushedActionId);
        lastPushedActionId = updates.newUnfilteredCount;

        if (!updates.mainResults.empty()) {
            subscriber->OnUpdate(
                    updates.mainResults,
                    updates.filteredResults,
                    updates.newUnfilteredCount,
                    updates.currentGameState,
                    updates.currentRound,
                    updates.playerBattleProgress);
        }

        // Tutorial battle reset: defender lost, signal reset instead of game over
        if (tutorialReset) {
            subscriber->OnBattleReset();
            return WaitResult::BATTLE_RESET;
        }

        // Now send gameOver notification after all updates have been sent
        if (gameOver) {
            subscriber->OnGameOver(gameOverInfo);
            return WaitResult::GAME_OVER;
        }
    }

    return WaitResult::DISCONNECTED;  // Subscriber disconnected
}

}  // namespace shardok
