#include "session_registry.h"

#ifdef VOICECAT_HAS_NET

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

#include "conn_session.h"

namespace voicecat::server {

void SessionRegistry::init_default_channels() {
    std::unique_lock lk(mu_);
    ChannelEntry lobby;
    lobby.proto.set_id(1);
    lobby.proto.set_name("Lobby");
    lobby.proto.set_type(voicecat::v1::CHANNEL_PERMANENT);
    lobby.proto.set_order(0);
    channels_[1] = std::move(lobby);
}

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);
}

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;
}

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::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;

    if (scope == voicecat::v1::TEXT_CHANNEL) {
        // Find channel_id of the target, then all users in that channel
        for (auto& [uid, entry] : users_) {
            if (entry.proto.channel_id() != target_id) continue;
            if (entry.session_id == sender_session_id) continue;
            auto sit = sessions_.find(entry.session_id);
            if (sit == sessions_.end()) continue;
            if (auto sess = sit->second.lock()) targets.push_back(sess);
        }
    } else if (scope == voicecat::v1::TEXT_PRIVATE) {
        // target_id is user_id
        auto user_it = users_.find(target_id);
        if (user_it != users_.end()) {
            auto sit = sessions_.find(user_it->second.session_id);
            if (sit != sessions_.end()) {
                if (auto sess = sit->second.lock()) targets.push_back(sess);
            }
        }
    }
    return targets;
}

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

// ── M2: 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()) result.push_back(sess);
    }
    return result;
}

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();
}

}  // namespace voicecat::server

#endif  // VOICECAT_HAS_NET