#include "WebsocketServer.h"
#include <spdlog/spdlog.h>

WebsocketServer::WebsocketServer(const tcp::endpoint& endpoint)
    : worker_pool_(std::thread::hardware_concurrency()),  // 传递线程池大小
      boss_context_(1), // Boss 线程使用单线程 io_context
      boss_work_guard_(asio::make_work_guard(boss_context_)),
      acceptor_(boss_context_, endpoint),
      session_strand_(boss_context_.get_executor()) // 确保 sessions_ 操作是线程安全的
{
    // 设置地址重用和监听队列
    acceptor_.set_option(asio::socket_base::reuse_address(true));
    acceptor_.listen(kBacklogSize);
}

WebsocketServer::~WebsocketServer() {
    stop();
}

void WebsocketServer::start() {
    // 启动Boss线程（1个）
    boss_threads_.emplace_back([this] {
#ifdef _WIN32
        // 可设置线程亲和性，绑定到核心0（示例）
        SetThreadAffinityMask(GetCurrentThread(), 1ULL << 0);
#endif
        boss_context_.run();
    });

    // 启动Worker线程（N-1个）
    worker_pool_.start();

    // 开始接受连接
    asio::post(boss_context_, [this] { doAccept(); });

    spdlog::info("[Server] Started on {} with {} threads.",
        acceptor_.local_endpoint().address().to_string(),
        worker_pool_.size() + 1);
}

void WebsocketServer::stop() {
    // 关闭 Acceptor
    if (acceptor_.is_open()) {
        acceptor_.close();
    }

    // 关闭所有会话
    asio::post(session_strand_, [this] {
        for (auto& session : sessions_) {
            session->close();
        }
        sessions_.clear();
    });

    // 停止线程池
    boss_context_.stop();
    worker_pool_.stop();

    for (auto& t : boss_threads_) {
        if (t.joinable()) t.join();
    }
    spdlog::info("[Server] Stopped.");
}

void WebsocketServer::broadcast(const Message& msg) const
{
    // 此处采用消息的二进制形式发送，或直接调用 session->send(msg) 进行序列化后发送
    asio::post(session_strand_, [this, msg] {
        for (auto& session : sessions_) {
            session->send(msg);
        }
    });
}

size_t WebsocketServer::connectionCount() const {
    return session_count_.load();  // 使用原子变量，减少锁竞争
}

bool WebsocketServer::isRunning() const {
    return acceptor_.is_open();
}

void WebsocketServer::doAccept() {
    acceptor_.async_accept(
        [this](boost::system::error_code ec, tcp::socket socket) {
            if (!ec) {
                spdlog::info("[Server] New connection accepted");
                // 设置 Socket 选项
                setupSocketOptions(socket);
                // 将新连接的 socket 移交给 worker_context_
                worker_pool_.post([this, s = std::move(socket)]() mutable {
                    spdlog::info("[Worker] New session lambda executed");

                    auto session = std::make_shared<WebsocketSession>(std::move(s));

                    // 绑定会话关闭回调
                    session->setOnClose([this](const std::shared_ptr<WebsocketSession>& s) {
                        removeSession(s);
                    });

                    addSession(session);

                    // 构造消息处理 Pipeline（类似 Netty 的 ChannelPipeline）
                    auto pipeline = std::make_shared<ChannelPipeline>();
                    pipeline->addHandler("loggingHandler", std::make_shared<LoggingHandler>());
                    session->setPipeline(pipeline);

                    // 启动握手（如果需要严格验证 URI，可参考前面的示例扩展 handshake 逻辑）
                    session->start();
                });

                // 继续接受新连接
                doAccept();
            } else if (ec != asio::error::operation_aborted) {
                spdlog::error("[Server] Accept error: {}", ec.message());

                // 防止高频错误导致CPU占用过高，使用 `steady_timer` 延迟重新监听
                auto timer = std::make_shared<asio::steady_timer>(boss_context_, std::chrono::milliseconds(100));
                timer->async_wait([this, timer](const boost::system::error_code& ec) {
                    if (!ec) {
                        doAccept();
                    }
                });
            }
        }
    );
}

void WebsocketServer::setupSocketOptions(tcp::socket& socket) {
    // 设置 TCP 参数
    socket.set_option(asio::ip::tcp::no_delay(true));      // 关闭 Nagle 算法
    socket.set_option(asio::socket_base::keep_alive(true));// 保持连接
    socket.set_option(asio::socket_base::receive_buffer_size(kRecvBufferSize));
    socket.set_option(asio::socket_base::send_buffer_size(kSendBufferSize));
}

void WebsocketServer::addSession(const std::shared_ptr<WebsocketSession>& session) {
    asio::post(session_strand_, [this, session] {
        sessions_.insert(session);
        ++session_count_;
    });
}

void WebsocketServer::removeSession(const std::shared_ptr<WebsocketSession>& session) {
    asio::post(session_strand_, [this, session] {
        sessions_.erase(session);
        --session_count_;
    });
}