#include <atomic>
#include <csignal>
#include <memory>
#include <thread>

#include "rdp_worker/common/logger.hpp"
#include "rdp_worker/config/config.hpp"
#include "rdp_worker/coordination/control_plane.hpp"
#include "rdp_worker/dataplane/direct_wss_server.hpp"
#include "rdp_worker/runtime/session_manager.hpp"

namespace {

std::atomic<bool> g_stop{false};

void HandleSignal(int) {
    g_stop.store(true);
}

}  // namespace

int main() {
    std::signal(SIGINT, HandleSignal);
    std::signal(SIGTERM, HandleSignal);

    auto logger = std::make_shared<rdp_worker::common::Logger>("rdp-worker");
    const auto config = rdp_worker::config::LoadFromEnv();
    auto control_plane = std::make_shared<rdp_worker::coordination::ControlPlane>(config, logger);
    control_plane->Connect();
    control_plane->ReleaseOwnedLeasesOnStartup();
    control_plane->RegisterWorker();

    auto session_manager = std::make_shared<rdp_worker::runtime::SessionManager>(control_plane, logger);
    rdp_worker::dataplane::DirectWssServer direct_wss_server(config, session_manager, logger);
    direct_wss_server.Start();

    std::thread heartbeat_thread([&]() {
        while (!g_stop.load()) {
            control_plane->SendHeartbeat();
            std::this_thread::sleep_for(config.worker_heartbeat_interval);
        }
    });

    while (!g_stop.load()) {
        if (auto assignment = control_plane->PollAssignment(config.assignment_poll_interval); assignment.has_value()) {
            session_manager->ApplyAssignment(*assignment);
        }
    }

    direct_wss_server.Stop();
    session_manager->StopAll();
    if (heartbeat_thread.joinable()) {
        heartbeat_thread.join();
    }
    return 0;
}