#include "rdp_worker/runtime/session_manager.hpp"

#include "rdp_worker/runtime/direct_bind_policy.hpp"

namespace rdp_worker::runtime {

SessionManager::SessionManager(std::shared_ptr<coordination::ControlPlane> control_plane,
                               std::shared_ptr<common::Logger> logger)
    : control_plane_(std::move(control_plane)),
      logger_(std::move(logger)) {}

void SessionManager::ApplyAssignment(const Assignment& assignment) {
    std::lock_guard<std::mutex> lock(mutex_);
    const auto iterator = sessions_.find(assignment.session_id);
    if (iterator != sessions_.end()) {
        iterator->second->ApplyAssignment(assignment);
        logger_->Info("updated assignment for existing session " + assignment.session_id);
        return;
    }

    auto runtime = std::make_shared<SessionRuntime>(assignment, control_plane_, logger_);
    runtime->Start();
    sessions_.emplace(assignment.session_id, runtime);
    logger_->Info("started new runtime for session " + assignment.session_id);
}

void SessionManager::StopAll() {
    std::lock_guard<std::mutex> lock(mutex_);
    for (auto& [_, runtime] : sessions_) {
        runtime->Stop(true, "worker_shutdown");
    }
    sessions_.clear();
}

bool SessionManager::BindDirectDataPlaneAttachment(const dataplane::DataPlaneTokenClaims& claims, std::string& reason) {
    return BindDirectDataPlaneRuntime(claims, reason) != nullptr;
}

std::shared_ptr<SessionRuntime> SessionManager::BindDirectDataPlaneRuntime(const dataplane::DataPlaneTokenClaims& claims, std::string& reason) {
    std::lock_guard<std::mutex> lock(mutex_);
    const auto iterator = sessions_.find(claims.session_id);
    if (iterator == sessions_.end()) {
        reason = "missing_runtime";
        return nullptr;
    }

    const Assignment snapshot = iterator->second->Snapshot();
    const auto validation = ValidateDirectDataPlaneBind(snapshot, claims);
    if (!validation.ok) {
        reason = validation.reason;
        return nullptr;
    }

    reason.clear();
    logger_->Info("event=data_plane_bind_success session=" + claims.session_id +
                  " attachment=" + claims.attachment_id +
                  " user=" + claims.user_id +
                  " organization=" + claims.organization_id +
                  " resource=" + claims.resource_id);
    return iterator->second;
}

}  // namespace rdp_worker::runtime