//
//  EagleInterfaceGrpcServer.cpp
//  shardok_engine
//
//  Created by Dan Crosby on 1/28/19.
//  Copyright © 2019 none. All rights reserved.
//

#include "EagleInterfaceGrpcServer.hpp"

#include <algorithm>
#include <iostream>
#include <iterator>
#include <memory>
#include <ranges>
#include <string>
#include <thread>
#include <utility>

#include "src/main/cpp/net/eagle0/common/FilesystemUtils.hpp"
#include "src/main/cpp/net/eagle0/common/ShardokLatencyTrace.hpp"
#include "src/main/cpp/net/eagle0/common/UnitConversions.hpp"
#include "src/main/cpp/net/eagle0/shardok/controller/ShardokGameController.hpp"
#include "src/main/cpp/net/eagle0/shardok/library/ShardokEngine.hpp"
#include "src/main/cpp/net/eagle0/shardok/server/GameOverResponsePopulator.hpp"
#include "src/main/cpp/net/eagle0/shardok/util/MapLoader.hpp"
#include "src/main/protobuf/net/eagle0/common/victory_condition.pb.h"

namespace shardok {
using grpc::Status;
using grpc::StatusCode;
using net::eagle0::common::GameSetupInfo;
using net::eagle0::common::ScopedShardokLatencyTrace;
using net::eagle0::common::ShardokLatencyFields;
using net::eagle0::common::VictoryCondition;
using net::eagle0::shardok::api::PlacementCommand;
using net::eagle0::shardok::common::EndGameCondition;

class NewGameException : public std::exception {
private:
    Status status;

public:
    explicit NewGameException(const StatusCode s, const string &m) : status(s, m) {}

    auto GetStatus() const -> Status { return status; }
};

auto EagleInterfaceImpl::ControllerForGame(const GameId &gameId, const GameSetupInfo &setupInfo)
        -> std::shared_ptr<ShardokGameController> {
    ScopedShardokLatencyTrace trace(
            "controller_for_game",
            gameId,
            {{"known_result_count", std::to_string(setupInfo.known_result_count())},
             {"has_current_game_state", setupInfo.current_game_state().empty() ? "false" : "true"},
             {"has_new_game_request", setupInfo.has_new_game_request() ? "true" : "false"}});
    auto controller = gamesManager->GetController(gameId);

    // If game exists but Eagle is requesting a fresh start with no prior history
    // (happens after admin rewind), replace the old game
    if (controller != nullptr && setupInfo.has_new_game_request() &&
        setupInfo.known_result_count() == 0) {
        std::cout << "Replacing existing game " << gameId << " with fresh game (rewind detected)"
                  << '\n';
        gamesManager->RemoveController(gameId);
        controller = nullptr;
    }

    if (controller == nullptr) {
        if (setupInfo.current_game_state().empty()) {
            StartGame(setupInfo.new_game_request());
        } else {
            gamesManager->UnlockedCreateFromGameStateBytes(
                    byte_vector(setupInfo.current_game_state()),
                    setupInfo.known_result_count(),
                    "");
        }

        controller = gamesManager->GetController(gameId);
    }

    if (controller == nullptr) { throw ShardokClientErrorException("Not found"); }

    return controller;
}

EagleInterfaceImpl::EagleInterfaceImpl(
        std::shared_ptr<ShardokGamesManager> manager,
        std::string authToken)
    : gamesManager(std::move(manager)),
      tokenValidator_(std::move(authToken)) {
    if (tokenValidator_.IsEnabled()) {
        std::cout << "Token authentication enabled" << '\n';
    } else {
        std::cout << "Token authentication disabled (no token configured)" << '\n';
    }
}

auto ConvertPlayerInfo(
        const google::protobuf::RepeatedPtrField<net::eagle0::common::PlayerSetupInfo> &allPis,
        PlayerId shardokPid) -> PlayerInfoWithUnits {
    const net::eagle0::common::PlayerSetupInfo &pi = allPis[shardokPid];

    const bool isAi =
            pi.type() ==
            net::eagle0::common::PlayerSetupInfo::PlayerType::PlayerSetupInfo_PlayerType_AI;

    const int32_t eagleFactionId = pi.eagle_faction_id();

    vector<PlayerId> allPids{};
    const int playerInfoCount = allPis.size();
    allPids.reserve(playerInfoCount);
    for (int i = 0; i < playerInfoCount; ++i) { allPids.push_back(static_cast<PlayerId>(i)); }

    vector<Unit> shardokUnits;
    shardokUnits.reserve(pi.units().size());
    std::ranges::transform(
            pi.units(),
            std::back_inserter(shardokUnits),
            [shardokPid, &allPids](const net::eagle0::common::CommonUnit &unit) {
                return ConvertUnit(unit, shardokPid, allPids);
            });

    vector<net::eagle0::shardok::common::VictoryCondition> victoryConditions;
    victoryConditions.reserve(pi.victory_conditions().size());
    std::ranges::transform(
            pi.victory_conditions(),
            std::back_inserter(victoryConditions),
            [](const int32_t vc) { return net::eagle0::shardok::common::VictoryCondition(vc); });

    vector<PlayerId> alliedPids{};
    alliedPids.reserve(pi.allies_size());
    for (const auto &ally : pi.allies()) {
        const int32_t alliedFid = ally.eagle_faction_id();
        for (int i = 0; i < playerInfoCount; ++i) {
            if (allPis[i].eagle_faction_id() == alliedFid) {
                alliedPids.push_back(static_cast<PlayerId>(i));
                break;
            }
        }
    }

    return PlayerInfoWithUnits(
            shardokPid,
            isAi,
            eagleFactionId,
            pi.defender(),
            pi.food(),
            victoryConditions,
            alliedPids,
            shardokUnits,
            pi.cannot_become_outlaw());
}

void EagleInterfaceImpl::StartGame(const NewGameRequest &request) {
    ScopedShardokLatencyTrace trace(
            "start_game",
            request.game_id(),
            {{"eagle_game_id", std::to_string(request.eagle_game_id())},
             {"player_count", std::to_string(request.user_infos_size())},
             {"custom_map_bytes", std::to_string(request.custom_map_bytes().size())}});
    // Get the custom map, or the map name if there isn't one
    const string &mapName = request.hex_map_name();
    const string &customMapBytes = request.custom_map_bytes();
    if (mapName.empty() && customMapBytes.empty()) {
        throw NewGameException(
                StatusCode::INVALID_ARGUMENT,
                "Must include map path unless a custom map is present");
    }

    const int month = request.month();
    if (month == 0) {
        throw NewGameException(StatusCode::INVALID_ARGUMENT, "Must include battle month");
    }

    // Also grab a collection of all the hero IDs for logging
    vector<HeroId> heroIds{};
    for (const auto &info : request.user_infos()) {
        for (const auto &unit : info.units()) { heroIds.push_back(unit.hero().eagle_hero_id()); }
    }
    std::ranges::sort(heroIds);

    // Convert each interface player struct into the Shardok storage equivalents
    vector<PlayerInfoWithUnits> playersInfo{};
    const int userInfoCount = request.user_infos_size();
    playersInfo.reserve(userInfoCount);
    for (int i = 0; i < userInfoCount; ++i) {
        playersInfo.push_back(ConvertPlayerInfo(request.user_infos(), static_cast<PlayerId>(i)));
    }

    // Validate that exactly one player has WIN_AFTER_MAX_ROUNDS
    int winAfterMaxRoundsCount = 0;
    for (const auto &psi : playersInfo) {
        for (const auto &vc : psi.playerInfo.victory_conditions()) {
            if (vc == net::eagle0::shardok::common::VICTORY_CONDITION_WIN_AFTER_MAX_ROUNDS) {
                winAfterMaxRoundsCount++;
            }
        }
    }
    if (winAfterMaxRoundsCount != 1) {
        throw NewGameException(
                StatusCode::INVALID_ARGUMENT,
                "Exactly one player must have WIN_AFTER_MAX_ROUNDS victory condition, found " +
                        std::to_string(winAfterMaxRoundsCount));
    }

    // Grab the nonempty gameId
    const string &gameId = request.game_id();
    if (gameId.empty()) {
        throw NewGameException(StatusCode::INVALID_ARGUMENT, "Must include game ID");
    }

    string unitsString;
    for (const auto &hid : heroIds) { unitsString += std::to_string(hid) + " "; }

    // Build the watcher → allied-Shardok-pids mapping. A non-participant
    // faction "watches" via any participant that lists it as an ally.
    std::vector<std::pair<int32_t, std::vector<PlayerId>>> watcherAllies;
    watcherAllies.reserve(request.watcher_eagle_faction_ids_size());
    for (const int32_t watcherFid : request.watcher_eagle_faction_ids()) {
        std::vector<PlayerId> allyPids;
        allyPids.reserve(userInfoCount);
        for (int i = 0; i < userInfoCount; ++i) {
            for (const auto &ally : request.user_infos(i).allies()) {
                if (ally.eagle_faction_id() == watcherFid) {
                    allyPids.push_back(static_cast<PlayerId>(i));
                    break;
                }
            }
        }
        if (!allyPids.empty()) { watcherAllies.emplace_back(watcherFid, std::move(allyPids)); }
    }

    std::cout << "Starting game " << gameId << " with units " << unitsString << '\n';
    gamesManager->UnlockedCreateSpecifiedGame(
            gameId,
            playersInfo,
            mapName,
            customMapBytes,
            month,
            request.tutorial_battle_config(),
            request.SerializeAsString(),
            watcherAllies);
}

auto EagleInterfaceImpl::PostCommand(
        ServerContext *context,
        const PostCommandRequest *request,
        GameStatusResponse *response) -> Status {
    ScopedShardokLatencyTrace totalTrace(
            "grpc_post_command_total",
            request->game_id(),
            {{"player_id", std::to_string(request->player_id())},
             {"eagle_faction_id", std::to_string(request->eagle_faction_id())},
             {"index", std::to_string(request->index())},
             {"known_result_count",
              std::to_string(request->game_setup_info().known_result_count())},
             {"request_bytes", std::to_string(request->ByteSizeLong())}});
    // Validate auth token
    if (auto authStatus = tokenValidator_.ValidateOrStatus(context)) { return *authStatus; }

    std::shared_ptr<ShardokGameController> controller;
    try {
        controller = ControllerForGame(request->game_id(), request->game_setup_info());
    } catch (const NewGameException &e) { return e.GetStatus(); }
    try {
        {
            ScopedShardokLatencyTrace trace(
                    "controller_post_command",
                    request->game_id(),
                    {{"player_id", std::to_string(request->player_id())},
                     {"index", std::to_string(request->index())}});
            controller->PostCommand(
                    static_cast<PlayerId>(request->player_id()),
                    request->eagle_faction_id(),
                    request->token(),
                    request->index(),
                    request->has_roll() ? std::optional<int32_t>(request->roll().value())
                                        : std::nullopt);
        }

        {
            ScopedShardokLatencyTrace trace(
                    "populate_post_command_response",
                    request->game_id(),
                    {{"starting_action_id",
                      std::to_string(request->game_setup_info().known_result_count())}});
            PopulateGameStatusResponse(
                    controller,
                    request->game_setup_info().known_result_count(),
                    response);
        }
    } catch (const ShardokClientErrorException &e) {
        std::cerr << "Bad request from client: " << e.what() << '\n';
        return Status(StatusCode::INVALID_ARGUMENT, e.what());
    } catch (const ShardokInternalErrorException &e) {
        std::cerr << "Internal error: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    } catch (const std::exception &e) {
        std::cerr << "Unknown exception: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    }

    return Status::OK;
}

auto EagleInterfaceImpl::PostPlacementCommands(
        ServerContext *context,
        const PlacementCommandsRequest *request,
        GameStatusResponse *response) -> Status {
    ScopedShardokLatencyTrace totalTrace(
            "grpc_post_placement_commands_total",
            request->game_id(),
            {{"player_id", std::to_string(request->player_id())},
             {"eagle_faction_id", std::to_string(request->eagle_faction_id())},
             {"placement_count", std::to_string(request->placement_commands_size())},
             {"known_result_count",
              std::to_string(request->game_setup_info().known_result_count())},
             {"request_bytes", std::to_string(request->ByteSizeLong())}});
    // Validate auth token
    if (auto authStatus = tokenValidator_.ValidateOrStatus(context)) { return *authStatus; }

    std::shared_ptr<ShardokGameController> controller;
    try {
        controller = ControllerForGame(request->game_id(), request->game_setup_info());
    } catch (const NewGameException &e) { return e.GetStatus(); }
    try {
        vector<UnitPlacementInfo> upis;
        upis.reserve(request->placement_commands().size());
        std::ranges::transform(
                request->placement_commands(),
                std::back_inserter(upis),
                [](const net::eagle0::shardok::api::PlacementCommand &cmd) {
                    return UnitPlacementInfo(
                            static_cast<UnitId>(cmd.unit_id()),
                            Coords(cmd.row(), cmd.column()));
                });
        {
            ScopedShardokLatencyTrace trace(
                    "controller_post_placement_commands",
                    request->game_id(),
                    {{"player_id", std::to_string(request->player_id())},
                     {"placement_count", std::to_string(request->placement_commands_size())}});
            controller->PostPlacementCommands(
                    static_cast<PlayerId>(request->player_id()),
                    request->eagle_faction_id(),
                    request->token(),
                    upis);
        }

        {
            ScopedShardokLatencyTrace trace(
                    "populate_post_placement_response",
                    request->game_id(),
                    {{"starting_action_id",
                      std::to_string(request->game_setup_info().known_result_count())}});
            PopulateGameStatusResponse(
                    controller,
                    request->game_setup_info().known_result_count(),
                    response);
        }
    } catch (const ShardokClientErrorException &e) {
        std::cerr << "Bad request from client: " << e.what() << '\n';
        return Status(StatusCode::INVALID_ARGUMENT, e.what());
    } catch (const ShardokInternalErrorException &e) {
        std::cerr << "Internal error: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    } catch (const std::exception &e) {
        std::cerr << "Unknown exception: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    }

    return Status::OK;
}

void EagleInterfaceImpl::PopulateGameStatusResponse(
        std::shared_ptr<ShardokGameController> &controller,
        const int64_t startingActionId,
        GameStatusResponse *response) {
    ScopedShardokLatencyTrace trace(
            "populate_game_status_response",
            controller->GetGameId(),
            {{"starting_action_id", std::to_string(startingActionId)}});
    response->set_game_id(controller->GetGameId());

    auto results = controller->GetUpdates(startingActionId);

    if (results.mainResults.empty()) {
        // check for game over
        if (controller->GameIsOver()) {
            PopulateGameOverResponse(
                    controller->GetGameId(),
                    controller->GameOverStatus(),
                    controller->ResolvedPlayerInfos(),
                    controller->EndGameUnits(),
                    response->mutable_game_over_response());
            // player resolution infos
            //            response->mutable_game_over_response()->mutable_user_infos() ...
        }
    } else {
        response->mutable_game_update_response()->mutable_update_responses()->Add(
                std::begin(results.mainResults),
                std::end(results.mainResults));
        for (const auto &op : results.filteredResults) {
            auto *newFilteredResponse =
                    response->mutable_game_update_response()->add_filtered_update_responses();
            newFilteredResponse->set_eagle_faction_id(op.eagleFactionId);
            newFilteredResponse->mutable_action_result_views()->Add(
                    std::begin(op.resultViews),
                    std::end(op.resultViews));

            if (op.availableCommands) {
                *newFilteredResponse->mutable_available_commands() = *op.availableCommands;
            }
        }
        response->mutable_game_update_response()->set_total_action_result_count(
                results.newUnfilteredCount);

        *response->mutable_game_update_response()->mutable_current_game_state() =
                string(controller->GetCurrentGameStateBytes());

        response->mutable_game_update_response()->set_current_round(results.currentRound);
        for (const auto &[playerId, troopCount, isRetreating] : results.playerBattleProgress) {
            auto *ptc = response->mutable_game_update_response()->add_player_troop_counts();
            ptc->set_player_id(playerId);
            ptc->set_troop_count(troopCount);
            ptc->set_is_retreating(isRetreating);
        }
    }
}

auto EagleInterfaceImpl::GetHexMap(
        ServerContext *context,
        const HexMapRequest *request,
        HexMapResponse *response) -> Status {
    // Validate auth token
    if (auto authStatus = tokenValidator_.ValidateOrStatus(context)) { return *authStatus; }

    // TODO: return appropriate status code if a bad map name is sent
    *response->mutable_map() = LoadMap(request->map_name());

    return Status::OK;
}

auto EagleInterfaceImpl::GetHexMapNames(
        ServerContext *context,
        const HexMapNamesRequest * /*request*/,
        HexMapNamesResponse *response) -> Status {
    // Validate auth token
    if (auto authStatus = tokenValidator_.ValidateOrStatus(context)) { return *authStatus; }

    for (const string &mapName : GetMapNames()) { response->add_map_names(mapName); }

    return Status::OK;
}

auto EagleInterfaceImpl::SubscribeToGame(
        ServerContext *context,
        const GameSubscriptionRequest *request,
        grpc::ServerWriter<GameStatusResponse> *writer) -> Status {
    ScopedShardokLatencyTrace subscribeSetupTrace(
            "grpc_subscribe_to_game_setup",
            request->game_id(),
            {{"known_result_count",
              std::to_string(request->game_setup_info().known_result_count())},
             {"request_bytes", std::to_string(request->ByteSizeLong())}});
    // Validate auth token
    if (auto authStatus = tokenValidator_.ValidateOrStatus(context)) { return *authStatus; }

    std::shared_ptr<ShardokGameController> controller;
    try {
        controller = ControllerForGame(request->game_id(), request->game_setup_info());
    } catch (const NewGameException &e) {
        return e.GetStatus();
    } catch (const ShardokClientErrorException &e) {
        std::cerr << "Bad request from client: " << e.what() << '\n';
        return Status(StatusCode::INVALID_ARGUMENT, e.what());
    } catch (const ShardokInternalErrorException &e) {
        std::cerr << "Internal error: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    } catch (const std::exception &e) {
        std::cerr << "Unknown exception in SubscribeToGame: " << e.what() << '\n';
        return Status(StatusCode::INTERNAL, e.what());
    }

    if (!controller) { return Status(StatusCode::NOT_FOUND, "Game not found"); }

    // Send initial state
    GameStatusResponse initialResponse;
    {
        ScopedShardokLatencyTrace trace(
                "subscribe_initial_populate",
                controller->GetGameId(),
                {{"starting_action_id",
                  std::to_string(request->game_setup_info().known_result_count())}});
        PopulateGameStatusResponse(
                controller,
                request->game_setup_info().known_result_count(),
                &initialResponse);
    }
    {
        ScopedShardokLatencyTrace trace(
                "subscribe_initial_write",
                controller->GetGameId(),
                {{"response_bytes", std::to_string(initialResponse.ByteSizeLong())}});
        if (!writer->Write(initialResponse)) {
            return Status::OK;  // Client disconnected
        }
    }

    // If game was already over, we're done
    if (initialResponse.has_game_over_response()) { return Status::OK; }

    // Track the actual count after initial response to avoid duplicate sends
    const int64_t countAfterInitialResponse =
            initialResponse.has_game_update_response()
                    ? initialResponse.game_update_response().total_action_result_count()
                    : request->game_setup_info().known_result_count();

    // Create a subscriber that writes to the gRPC stream
    class GrpcStreamSubscriber : public StreamSubscriber {
    private:
        grpc::ServerWriter<GameStatusResponse> *writer_;
        ServerContext *context_;
        std::atomic<bool> active_{true};
        std::string gameId_;

    public:
        GrpcStreamSubscriber(
                grpc::ServerWriter<GameStatusResponse> *w,
                ServerContext *ctx,
                std::string gameId)
            : writer_(w),
              context_(ctx),
              gameId_(std::move(gameId)) {}

        void OnUpdate(
                const vector<ActionResult> &mainResults,
                const vector<OnePlayerUpdates> &filteredResults,
                int32_t newUnfilteredCount,
                const byte_vector &currentGameState,
                int32_t currentRound,
                const vector<PlayerBattleProgress> &playerBattleProgress) override {
            if (!active_) return;

            GameStatusResponse response;
            response.set_game_id(gameId_);

            response.mutable_game_update_response()->mutable_update_responses()->Add(
                    std::begin(mainResults),
                    std::end(mainResults));

            // Add filtered results for each player with their faction IDs
            for (const auto &playerUpdate : filteredResults) {
                auto *filtered =
                        response.mutable_game_update_response()->add_filtered_update_responses();
                filtered->set_eagle_faction_id(playerUpdate.eagleFactionId);
                filtered->mutable_action_result_views()->Add(
                        std::begin(playerUpdate.resultViews),
                        std::end(playerUpdate.resultViews));
                if (playerUpdate.availableCommands) {
                    *filtered->mutable_available_commands() = *playerUpdate.availableCommands;
                }
            }

            response.mutable_game_update_response()->set_total_action_result_count(
                    newUnfilteredCount);
            *response.mutable_game_update_response()->mutable_current_game_state() =
                    std::string(currentGameState.begin(), currentGameState.end());

            response.mutable_game_update_response()->set_current_round(currentRound);
            for (const auto &[playerId, troopCount, isRetreating] : playerBattleProgress) {
                auto *ptc = response.mutable_game_update_response()->add_player_troop_counts();
                ptc->set_player_id(playerId);
                ptc->set_troop_count(troopCount);
                ptc->set_is_retreating(isRetreating);
            }

            ScopedShardokLatencyTrace trace(
                    "subscribe_update_write",
                    gameId_,
                    {{"main_result_count", std::to_string(mainResults.size())},
                     {"filtered_result_count", std::to_string(filteredResults.size())},
                     {"new_unfiltered_count", std::to_string(newUnfilteredCount)},
                     {"response_bytes", std::to_string(response.ByteSizeLong())}});
            if (!writer_->Write(response)) { active_ = false; }
        }

        void OnGameOver(const GameOverInfo &info) override {
            if (!active_) return;

            GameStatusResponse response;
            response.set_game_id(gameId_);

            PopulateGameOverResponse(
                    gameId_,
                    info.gameStatus,
                    info.playerInfos,
                    info.endGameUnits,
                    response.mutable_game_over_response());

            ScopedShardokLatencyTrace trace(
                    "subscribe_game_over_write",
                    gameId_,
                    {{"response_bytes", std::to_string(response.ByteSizeLong())}});
            writer_->Write(response);
            active_ = false;
        }

        void OnBattleReset() override {
            if (!active_) return;

            GameStatusResponse response;
            response.set_game_id(gameId_);
            response.mutable_battle_reset_response();

            ScopedShardokLatencyTrace trace(
                    "subscribe_battle_reset_write",
                    gameId_,
                    {{"response_bytes", std::to_string(response.ByteSizeLong())}});
            if (!writer_->Write(response)) { active_ = false; }
        }

        [[nodiscard]] auto IsActive() const -> bool override {
            return active_ && !context_->IsCancelled();
        }
    };

    auto subscriber =
            std::make_shared<GrpcStreamSubscriber>(writer, context, controller->GetGameId());

    controller->RegisterSubscriber(subscriber);

    // Wait for updates and push them until game ends or subscriber disconnects.
    // For tutorial battles, a BATTLE_RESET result means the defender lost and we
    // should restart the game from the original request.
    int64_t actionCount = countAfterInitialResponse;
    WaitResult result = WaitResult::GAME_OVER;

    while (true) {
        result = controller->WaitForUpdatesAndPush(subscriber, actionCount);

        if (result != WaitResult::BATTLE_RESET) break;

        std::cout << "SubscribeToGame: Tutorial battle reset, restarting game "
                  << controller->GetGameId() << '\n';

        // Save the serialized request before removing the controller
        const std::string savedRequest = controller->SerializedRequest();
        const GameId gameId = controller->GetGameId();

        controller->UnregisterSubscriber(subscriber.get());

        gamesManager->RemoveController(gameId);

        // Deserialize and restart the game from the original request
        NewGameRequest newGameRequest;
        if (!newGameRequest.ParseFromString(savedRequest)) {
            return Status(StatusCode::INTERNAL, "Failed to parse saved tutorial battle request");
        }
        StartGame(newGameRequest);

        controller = gamesManager->GetController(gameId);
        if (!controller) {
            std::cout << "SubscribeToGame: Failed to restart game after reset" << '\n';
            return Status(StatusCode::INTERNAL, "Failed to restart tutorial battle");
        }

        // Send initial state from the new game
        GameStatusResponse resetInitialResponse;
        PopulateGameStatusResponse(controller, 0, &resetInitialResponse);
        if (!writer->Write(resetInitialResponse)) break;

        // Reset action count based on the new game's state
        actionCount =
                resetInitialResponse.has_game_update_response()
                        ? resetInitialResponse.game_update_response().total_action_result_count()
                        : 0;

        // Re-register subscriber with the new controller
        controller->RegisterSubscriber(subscriber);
    }

    controller->UnregisterSubscriber(subscriber.get());

    if (controller->HasAIThreadError()) {
        std::cerr << "SubscribeToGame: AI thread error: " << controller->GetAIThreadError() << '\n';
        return Status(StatusCode::INTERNAL, controller->GetAIThreadError());
    }

    if (result == WaitResult::GAME_OVER) {
        std::cout << "SubscribeToGame: Game ended normally" << '\n';
    } else {
        std::cout << "SubscribeToGame: Subscriber disconnected" << '\n';
    }

    return Status::OK;
}

}  // namespace shardok

#ifndef NDEBUG
template class std::vector<net::eagle0::shardok::storage::ResolvedUnit>;
#endif
