WSCommModifyService/WebSocketSession.h

#pragma once

#include <boost/beast/core.hpp>
#include <boost/beast/http.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/asio/ip/tcp.hpp>
#include <boost/asio/strand.hpp>
#include <boost/asio/dispatch.hpp>
#include <cstdlib>
#include <functional>
#include <iostream>
#include <memory>
#include <string>
#include <thread>
#include <mutex>
#include <queue>
#include <atomic>

namespace beast = boost::beast;
namespace http = beast::http;
namespace websocket = beast::websocket;
namespace net = boost::asio;
using tcp = boost::asio::ip::tcp;

// 前向声明
class WebSocketHandler;

class WebSocketSession : public std::enable_shared_from_this<WebSocketSession> {
public:
    using MessageHandler = std::function<void(std::shared_ptr<WebSocketSession>, const std::string&)>;

private:
    websocket::stream<beast::tcp_stream> ws_;
    beast::flat_buffer buffer_;
    http::request<http::dynamic_body> req_;
    
    std::mutex send_mutex_;
    std::queue<std::shared_ptr<std::string>> send_queue_;
    std::shared_ptr<std::string> current_message_; // 添加这个成员变量
    bool is_sending_ = false;
    std::atomic<bool> closed_{false};
    
    MessageHandler message_handler_;
    
    // 添加对WebSocketHandler的引用
    std::shared_ptr<WebSocketHandler> websocket_handler_;

public:
    explicit WebSocketSession(tcp::socket&& socket, http::request<http::dynamic_body> req);
    ~WebSocketSession();
    
    void run();
    void safe_send(std::shared_ptr<std::string> message);
    void close();
    void set_message_handler(MessageHandler handler);
    
    // 添加设置WebSocketHandler的方法
    void set_websocket_handler(std::shared_ptr<WebSocketHandler> handler) { websocket_handler_ = handler; }

private:
    void on_accept(beast::error_code ec);
    void do_read();
    void on_read(beast::error_code ec, std::size_t bytes_transferred);
    void do_write();
    void on_write(beast::error_code ec, std::size_t bytes_transferred);
};