voice-cat/server/src/session_registry.cpp

#include "session_registry.h"

#include <atomic>
#include <chrono>
#include <mutex>
#include <shared_mutex>

#include "conn_session.h"

namespace voicecat::server {

SessionRegistry::SessionRegistry(std::shared_ptr<Database> db) : db_(std::move(db)) {}

void SessionRegistry::load_channels() {
    std::unique_lock lk(mu_);
    channels_.clear();

    auto records = db_->list_channels();
    if (records.empty()) {
        // First run: seed the default channel tree.
        seed_default_channels();
        records = db_->list_channels();
    }

    uint32_t max_id = 2;
    for (auto& rec : records) {
        ChannelEntry entry;
        entry.proto.set_id(rec.id);
        entry.proto.set_parent_id(rec.parent_id);
        entry.proto.set_name(rec.name);
        entry.proto.set_topic(rec.topic);
        entry.proto.set_password_protected(rec.password_protected);
        entry.proto.set_max_users(rec.max_users);
        entry.proto.set_type(rec.type);
        *entry.proto.mutable_audio() = rec.audio;
        entry.proto.set_order(rec.sort_order);
        max_id = std::max(max_id, rec.id);
        channels_[rec.id] = std::move(entry);
    }
    next_channel_id_ = max_id + 1;
}

void SessionRegistry::seed_default_channels() {
    // Lobby: speech profile.
    ChannelRecord lobby;
    lobby.id = 1;
    lobby.name = "Lobby";
    lobby.type = voicecat::v1::CHANNEL_PERMANENT;
    lobby.sort_order = 0;
    lobby.max_users = 20;
    {
        auto& a = lobby.audio;
        a.set_codec(0);
        a.set_mode(voicecat::v1::MODE_MONO);
        a.set_sample_rate(48000);
        a.set_bitrate_bps(24000);
        a.set_frame_ms(20);
        a.set_application(voicecat::v1::OPUS_VOIP);
        a.set_fec(true);
        a.set_expected_packet_loss(10);
        a.set_dtx(true);
        a.set_complexity(5);
    }

    // Music Room: stereo/music profile.
    ChannelRecord music;
    music.id = 2;
    music.name = "Music Room";
    music.type = voicecat::v1::CHANNEL_PERMANENT;
    music.sort_order = 1;
    {
        auto& a = music.audio;
        a.set_codec(0);
        a.set_mode(voicecat::v1::MODE_STEREO);
        a.set_sample_rate(48000);
        a.set_bitrate_bps(128000);
        a.set_frame_ms(20);
        a.set_application(voicecat::v1::OPUS_AUDIO);
        a.set_fec(false);
        a.set_expected_packet_loss(0);
        a.set_dtx(false);
        a.set_complexity(8);
    }

    std::string err;
    db_->create_channel(lobby, "", err);
    db_->create_channel(music, "", err);
}

uint64_t SessionRegistry::register_session(std::weak_ptr<ConnSession> session) {
    std::unique_lock lk(mu_);
    uint64_t id = next_session_id_++;
    sessions_[id] = std::move(session);
    return id;
}

void SessionRegistry::unregister_session(uint64_t session_id) {
    std::unique_lock lk(mu_);
    sessions_.erase(session_id);
    session_permissions_.erase(session_id);

    // Free the per-session UDP/media state too. These maps are keyed by
    // endpoint/token/ssrc (not session id), so scan-and-erase by value. Leaving them
    // behind leaks entries and lets a stale endpoint/token resolve toward a dead
    // session across reconnects (e.g. a wifi-handoff rebind from a new port).
    auto erase_by_value = [session_id](auto& map) {
        for (auto it = map.begin(); it != map.end();) {
            if (it->second == session_id)
                it = map.erase(it);
            else
                ++it;
        }
    };
    erase_by_value(udp_endpoints_);
    erase_by_value(udp_tokens_);
    erase_by_value(ssrc_to_session_);
}

uint32_t SessionRegistry::add_user(uint64_t session_id, const voicecat::v1::User& user) {
    std::unique_lock lk(mu_);
    uint32_t uid = next_user_id_++;
    UserEntry entry;
    entry.proto = user;
    entry.proto.set_id(uid);
    entry.proto.set_channel_id(1);  // start in Lobby
    entry.session_id = session_id;
    users_[uid] = std::move(entry);
    return uid;
}

void SessionRegistry::remove_user(uint32_t user_id) {
    std::unique_lock lk(mu_);
    users_.erase(user_id);
}

bool SessionRegistry::set_user_channel(uint32_t user_id, uint32_t channel_id) {
    std::unique_lock lk(mu_);
    auto ch_it = channels_.find(channel_id);
    if (ch_it == channels_.end()) return false;
    auto user_it = users_.find(user_id);
    if (user_it == users_.end()) return false;
    user_it->second.proto.set_channel_id(channel_id);
    return true;
}

void SessionRegistry::set_user_voice_subscribed(uint32_t user_id, bool subscribed) {
    std::unique_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return;
    it->second.proto.set_voice_subscribed(subscribed);
}

std::vector<voicecat::v1::Channel> SessionRegistry::channel_snapshot() const {
    std::shared_lock lk(mu_);
    std::vector<voicecat::v1::Channel> result;
    result.reserve(channels_.size());
    for (auto& [id, entry] : channels_) result.push_back(entry.proto);
    return result;
}

std::vector<voicecat::v1::User> SessionRegistry::user_snapshot() const {
    std::shared_lock lk(mu_);
    std::vector<voicecat::v1::User> result;
    result.reserve(users_.size());
    for (auto& [id, entry] : users_) result.push_back(entry.proto);
    return result;
}

std::optional<voicecat::v1::User> SessionRegistry::user_snapshot_user(uint32_t user_id) const {
    std::shared_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return std::nullopt;
    return it->second.proto;
}

std::optional<std::string> SessionRegistry::user_nickname(uint32_t user_id) const {
    std::shared_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return std::nullopt;
    return it->second.proto.nickname();
}

std::vector<std::shared_ptr<ConnSession>> SessionRegistry::resolve_text_targets(
    uint64_t sender_session_id, voicecat::v1::TextScope scope, uint32_t target_id) const {
    std::shared_lock lk(mu_);
    std::vector<std::shared_ptr<ConnSession>> targets;

    auto add_session = [&](uint64_t sid) {
        auto sit = sessions_.find(sid);
        if (sit == sessions_.end()) return;
        if (auto sess = sit->second.lock()) targets.push_back(sess);
    };

    if (scope == voicecat::v1::TEXT_CHANNEL) {
        // Fan out to everyone in the target channel, INCLUDING the sender, so the sender's
        // own client renders the message through the same authoritative relay everyone else
        // gets (no optimistic local echo). See docs/protocol.md §6.
        for (auto& [uid, entry] : users_) {
            if (entry.proto.channel_id() != target_id) continue;
            add_session(entry.session_id);
        }
    } else if (scope == voicecat::v1::TEXT_PRIVATE) {
        // target_id is the recipient user_id: deliver to the recipient and echo to the
        // sender (skip the echo if they messaged themselves, to avoid a duplicate).
        auto user_it = users_.find(target_id);
        if (user_it != users_.end()) add_session(user_it->second.session_id);
        if (user_it == users_.end() || user_it->second.session_id != sender_session_id)
            add_session(sender_session_id);
    }
    return targets;
}

void SessionRegistry::broadcast(const voicecat::v1::Envelope& env,
                                 uint64_t exclude_session_id) const {
    std::shared_lock lk(mu_);
    broadcast_unlocked(env, exclude_session_id);
}

void SessionRegistry::broadcast_unlocked(const voicecat::v1::Envelope& env,
                                          uint64_t exclude_session_id) const {
    for (auto& [sid, weak] : sessions_) {
        if (sid == exclude_session_id) continue;
        if (auto sess = weak.lock()) sess->send_envelope(env);
    }
}

std::vector<std::shared_ptr<ConnSession>> SessionRegistry::find_stale_sessions(
    int64_t max_age_ms) const {
    std::shared_lock lk(mu_);
    auto now_ms = std::chrono::duration_cast<std::chrono::milliseconds>(
        std::chrono::steady_clock::now().time_since_epoch()).count();
    std::vector<std::shared_ptr<ConnSession>> stale;
    for (auto& [sid, weak] : sessions_) {
        auto sess = weak.lock();
        if (!sess) continue;
        int64_t seen = sess->last_seen_ms();
        if (seen == 0) continue;  // not yet initialized — skip (shouldn't happen after begin())
        if (now_ms - seen > max_age_ms) stale.push_back(sess);
    }
    return stale;
}

// ── Permissions ───────────────────────────────────────────────────────────────

void SessionRegistry::set_session_permissions(uint64_t session_id,
                                               const voicecat::v1::Permissions& perms) {
    std::unique_lock lk(mu_);
    session_permissions_[session_id] = perms;
}

std::optional<voicecat::v1::Permissions> SessionRegistry::get_session_permissions(
    uint64_t session_id) const {
    std::shared_lock lk(mu_);
    auto it = session_permissions_.find(session_id);
    if (it == session_permissions_.end()) return std::nullopt;
    return it->second;
}

// ── Moderation ────────────────────────────────────────────────────────────────

std::shared_ptr<ConnSession> SessionRegistry::find_session_by_user_id(uint32_t user_id) const {
    std::shared_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return nullptr;
    auto sit = sessions_.find(it->second.session_id);
    if (sit == sessions_.end()) return nullptr;
    return sit->second.lock();
}

namespace {
    voicecat::v1::Envelope make_left_event(uint32_t user_id, const std::string& reason) {
        voicecat::v1::Envelope env;
        auto* ue = env.mutable_user_event();
        ue->set_kind(voicecat::v1::UserEvent::LEFT);
        ue->mutable_user()->set_id(user_id);
        ue->set_left_id(user_id);
        ue->set_reason(reason);
        return env;
    }
}

bool SessionRegistry::kick_user(uint32_t user_id, const std::string& reason) {
    auto target = find_session_by_user_id(user_id);
    if (!target) return false;
    {
        std::shared_lock lk(mu_);
        broadcast(make_left_event(user_id, reason), /*exclude*/ 0);
    }
    // Close outside the registry lock: close() may call back into unregister_session().
    target->send_disconnect_and_close(2, reason);  // code 2 = kicked
    return true;
}

void SessionRegistry::broadcast_left(uint32_t user_id, const std::string& reason) {
    std::shared_lock lk(mu_);
    broadcast_unlocked(make_left_event(user_id, reason), /*exclude*/ 0);
}

bool SessionRegistry::ban_user(uint32_t user_id, const std::string& reason, int64_t expires_at) {
    {
        std::unique_lock lk(mu_);
        auto it = users_.find(user_id);
        if (it != users_.end()) {
            std::string err;
            db_->ban_create("user_id", std::to_string(user_id), reason, expires_at, err);
        }
    }
    return kick_user(user_id, reason);
}

bool SessionRegistry::set_server_mute(uint32_t user_id, bool muted, bool deafened) {
    std::unique_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return false;
    it->second.proto.set_server_muted(muted);
    it->second.proto.set_server_deafened(deafened);

    voicecat::v1::Envelope env;
    auto* ue = env.mutable_user_event();
    ue->set_kind(voicecat::v1::UserEvent::UPDATED);
    *ue->mutable_user() = it->second.proto;
    broadcast_unlocked(env, 0);
    return true;
}

bool SessionRegistry::move_user(uint32_t user_id, uint32_t channel_id) {
    std::unique_lock lk(mu_);
    auto ch_it = channels_.find(channel_id);
    if (ch_it == channels_.end()) return false;
    auto user_it = users_.find(user_id);
    if (user_it == users_.end()) return false;

    user_it->second.proto.set_channel_id(channel_id);

    voicecat::v1::Envelope env;
    auto* ue = env.mutable_user_event();
    ue->set_kind(voicecat::v1::UserEvent::UPDATED);
    *ue->mutable_user() = user_it->second.proto;
    broadcast_unlocked(env, 0);
    return true;
}

// ── Channel CRUD ──────────────────────────────────────────────────────────────

uint32_t SessionRegistry::create_channel(const voicecat::v1::Channel& ch,
                                          const std::string& password,
                                          std::string& error) {
    ChannelRecord rec;
    rec.parent_id = ch.parent_id();
    rec.name = ch.name();
    rec.topic = ch.topic();
    rec.max_users = ch.max_users();
    rec.type = ch.type();
    rec.audio = ch.audio();
    rec.sort_order = ch.order();

    auto result = db_->create_channel(rec, password, error);
    if (!result) return 0;

    std::unique_lock lk(mu_);
    uint32_t id = result->id;
    ChannelEntry entry;
    entry.proto = ch;
    entry.proto.set_id(id);
    entry.proto.set_password_protected(result->password_protected);

    voicecat::v1::Envelope env;
    auto* ce = env.mutable_channel_event();
    ce->set_kind(voicecat::v1::ChannelEvent::CREATED);
    *ce->mutable_channel() = entry.proto;

    channels_[id] = std::move(entry);
    next_channel_id_ = std::max(next_channel_id_, id + 1);

    broadcast_unlocked(env, 0);
    return id;
}

bool SessionRegistry::update_channel(const voicecat::v1::Channel& ch,
                                      const std::string& password,
                                      std::string& error) {
    ChannelRecord rec;
    rec.id = ch.id();
    rec.parent_id = ch.parent_id();
    rec.name = ch.name();
    rec.topic = ch.topic();
    rec.max_users = ch.max_users();
    rec.type = ch.type();
    rec.audio = ch.audio();
    rec.sort_order = ch.order();

    if (!db_->update_channel(rec, password, error)) return false;

    std::unique_lock lk(mu_);
    auto it = channels_.find(ch.id());
    if (it == channels_.end()) {
        error = "channel not found";
        return false;
    }
    it->second.proto = ch;
    it->second.proto.set_password_protected(
        !password.empty() || db_->get_channel(ch.id())->password_protected);

    voicecat::v1::Envelope env;
    auto* ce = env.mutable_channel_event();
    ce->set_kind(voicecat::v1::ChannelEvent::UPDATED);
    *ce->mutable_channel() = it->second.proto;
    broadcast_unlocked(env, 0);
    return true;
}

bool SessionRegistry::delete_channel(uint32_t channel_id, std::string& error) {
    if (channel_id == 1) { error = "cannot delete root channel"; return false; }
    if (!db_->delete_channel(channel_id, error)) return false;

    std::unique_lock lk(mu_);
    channels_.erase(channel_id);

    // Move any users left in the deleted channel to Lobby.
    for (auto& [uid, entry] : users_) {
        if (entry.proto.channel_id() != channel_id) continue;
        entry.proto.set_channel_id(1);
        voicecat::v1::Envelope uev;
        auto* ue = uev.mutable_user_event();
        ue->set_kind(voicecat::v1::UserEvent::UPDATED);
        *ue->mutable_user() = entry.proto;
        broadcast_unlocked(uev, 0);
    }

    voicecat::v1::Envelope env;
    auto* ce = env.mutable_channel_event();
    ce->set_kind(voicecat::v1::ChannelEvent::DELETED);
    ce->set_deleted_id(channel_id);
    broadcast_unlocked(env, 0);
    return true;
}

std::optional<voicecat::v1::Channel> SessionRegistry::get_channel(uint32_t channel_id) const {
    std::shared_lock lk(mu_);
    auto it = channels_.find(channel_id);
    if (it == channels_.end()) return std::nullopt;
    return it->second.proto;
}

bool SessionRegistry::check_channel_password(uint32_t channel_id,
                                              const std::string& password) const {
    return db_->check_channel_password(channel_id, password);
}

// ── UDP / media ────────────────────────────────────────────────────────────────

void SessionRegistry::register_udp_token(const std::array<uint8_t, 16>& token,
                                          uint64_t session_id) {
    std::unique_lock lk(mu_);
    udp_tokens_[token] = session_id;
}

std::shared_ptr<ConnSession> SessionRegistry::find_by_udp_token(
    const std::array<uint8_t, 16>& token) const {
    std::shared_lock lk(mu_);
    auto it = udp_tokens_.find(token);
    if (it == udp_tokens_.end()) return nullptr;
    auto sit = sessions_.find(it->second);
    if (sit == sessions_.end()) return nullptr;
    return sit->second.lock();
}

void SessionRegistry::register_udp_endpoint(asio::ip::udp::endpoint ep,
                                               uint64_t session_id) {
    std::unique_lock lk(mu_);
    udp_endpoints_[ep] = session_id;
}

std::shared_ptr<ConnSession> SessionRegistry::find_by_udp_endpoint(
    const asio::ip::udp::endpoint& ep) const {
    std::shared_lock lk(mu_);
    auto it = udp_endpoints_.find(ep);
    if (it == udp_endpoints_.end()) return nullptr;
    auto sit = sessions_.find(it->second);
    if (sit == sessions_.end()) return nullptr;
    return sit->second.lock();
}

uint32_t SessionRegistry::assign_ssrc(uint64_t session_id) {
    uint32_t ssrc = next_ssrc_.fetch_add(1, std::memory_order_relaxed);
    std::unique_lock lk(mu_);
    ssrc_to_session_[ssrc] = session_id;
    return ssrc;
}

std::vector<std::shared_ptr<ConnSession>> SessionRegistry::find_channel_sessions(
    uint32_t channel_id, uint64_t exclude_session_id) const {
    std::shared_lock lk(mu_);
    std::vector<std::shared_ptr<ConnSession>> result;
    for (auto& [uid, entry] : users_) {
        if (entry.proto.channel_id() != channel_id) continue;
        if (entry.session_id == exclude_session_id) continue;
        auto sit = sessions_.find(entry.session_id);
        if (sit == sessions_.end()) continue;
        if (auto sess = sit->second.lock()) {
            // Only relay voice to sessions that are on the voice plane.
            if (!sess->voice_subscribed()) continue;
            result.push_back(sess);
        }
    }
    return result;
}

std::optional<voicecat::v1::User> SessionRegistry::set_user_stream(
    uint32_t user_id, const voicecat::v1::StreamInfo& info) {
    std::unique_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return std::nullopt;
    auto* streams = it->second.proto.mutable_streams();
    for (int i = 0; i < streams->size(); ++i) {
        if (streams->Get(i).stream_id() == info.stream_id()) {
            *streams->Mutable(i) = info;
            return it->second.proto;
        }
    }
    *streams->Add() = info;
    return it->second.proto;
}

std::optional<voicecat::v1::User> SessionRegistry::clear_user_stream(uint32_t user_id,
                                                                        uint32_t stream_id) {
    std::unique_lock lk(mu_);
    auto it = users_.find(user_id);
    if (it == users_.end()) return std::nullopt;
    auto* streams = it->second.proto.mutable_streams();
    for (int i = 0; i < streams->size(); ++i) {
        if (streams->Get(i).stream_id() == stream_id) {
            streams->erase(streams->begin() + i);
            break;
        }
    }
    return it->second.proto;
}

uint32_t SessionRegistry::user_channel(uint32_t user_id) const {
    std::shared_lock lk(mu_);
    auto it = users_.find(user_id);
    return (it == users_.end()) ? 0 : it->second.proto.channel_id();
}

std::optional<voicecat::v1::AudioConfig> SessionRegistry::channel_audio_config(
    uint32_t channel_id) const {
    std::shared_lock lk(mu_);
    auto it = channels_.find(channel_id);
    if (it == channels_.end()) return std::nullopt;
    return it->second.proto.audio();
}

}  // namespace voicecat::server