#pragma once

#include <boost/beast/core.hpp>
#include <boost/beast/websocket.hpp>
#include <boost/beast/http.hpp>
#include <boost/asio.hpp>
#include <memory>
#include <functional>
#include <spdlog/spdlog.h>

#include "Message.h"
#include "ChannelPipeline.h"
#include "LoggingHandler.h"

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

class WebsocketSession : public std::enable_shared_from_this<WebsocketSession> {
public:
    // 构造函数
    explicit WebsocketSession(tcp::socket socket)
        : handshake_verified_(false),
        ws_(std::move(socket)),
        strand_(ws_.get_executor())
    {
    }

    // 启动握手和读取循环，严格验证请求 URI 是否为 "/ws"
    void start() {
        // 读取 HTTP 请求头
        http::async_read(ws_.next_layer(), buffer_, req_,
            [self = shared_from_this()](beast::error_code ec, std::size_t) {
                if (!ec) {
                    // 验证 URI 是否为 "/ws"
                    if (self->req_.target() == "/ws") {
                        self->handshake_verified_ = true;
                        // 使用 HTTP 请求信息完成 WebSocket 握手
                        self->ws_.async_accept(self->req_,
                            [self](beast::error_code ec) {
                                if (!ec) {
                                    self->doRead();
                                } else {
                                    spdlog::error("[WebsocketSession] Handshake error: {}", ec.message());
                                    self->close();
                                }
                            });
                    } else {
                        // 返回 404 响应并关闭连接
                        http::response<http::string_body> res{http::status::not_found, self->req_.version()};
                        res.set(http::field::server, "HighPerformanceWebsocketServer");
                        res.body() = "Path not found";
                        res.prepare_payload();
                        http::async_write(self->ws_.next_layer(), res,
                            [self](beast::error_code, std::size_t) {
                                self->close();
                            });
                    }
                } else {
                    spdlog::error("[WebsocketSession] Handshake HTTP read error: {}", ec.message());
                    self->close();
                }
            });
    }

    // 发送消息（序列化后发送）
    void send(const Message& msg) {
        msgpack::sbuffer sbuf;
        msgpack::pack(sbuf, msg);
        auto data = std::make_shared<std::string>(sbuf.data(), sbuf.size());
        // 保证 send 操作在 strand 内执行，避免并发问题
        asio::post(strand_, [self = shared_from_this(), data]() {
            self->writeQueue_.push_back(data);
            if (!self->writing_) {
                self->doWrite();
            }
        });
    }

    // 关闭 WebSocket 连接，并触发 onClose 回调
    void close() {
        if (ws_.is_open()) {
            websocket::close_reason reason(websocket::close_code::normal);
            ws_.async_close(reason,
                [self = shared_from_this()](beast::error_code ec) {
                    if (ec) {
                        spdlog::error("[WebsocketSession] Close error: {}", ec.message());
                    } else {
                        spdlog::info("[WebsocketSession] Connection closed gracefully.");
                    }
                    // 通知上层管理：连接关闭事件
                    if (self->onClose) {
                        self->onClose(self);
                    }
                }
            );
        } else {
            spdlog::info("[WebsocketSession] WebSocket is already closed.");
            // 如果已经关闭，也调用回调通知上层
            if (onClose) {
                onClose(shared_from_this());
            }
        }
    }

    // 设置消息处理 Pipeline
    void setPipeline(std::shared_ptr<ChannelPipeline> pipeline) {
        pipeline_ = std::move(pipeline);
    }

    // 设置关闭回调，用于通知服务器移除该会话
    void setOnClose(std::function<void(std::shared_ptr<WebsocketSession>)> callback) {
        onClose = std::move(callback);
    }

private:
    // 循环读取客户端消息
    void doRead() {
        auto self = shared_from_this();
        ws_.async_read(buffer_,
            [this, self](beast::error_code ec, std::size_t bytes_transferred) {
                if (!ec) {
                    auto data = static_cast<const char*>(buffer_.cdata().data());
                    std::size_t size = buffer_.size();
                    try {
                        auto oh = msgpack::unpack(data, size);
                        Message msg;
                        oh.get().convert(msg);
                        if (pipeline_) {
                            pipeline_->process(self, msg);
                        }
                    } catch (const std::exception& e) {
                        spdlog::error("[WebsocketSession] Unpack error: {}", e.what());
                    }
                    buffer_.consume(bytes_transferred);
                    doRead();
                } else {
                    spdlog::error("[WebsocketSession] Read error: {}", ec.message());
                    close();
                }
            });
    }

    void doWrite() {
        if (writeQueue_.empty() || !ws_.is_open()) {
            writing_ = false;
            return;
        }

        writing_ = true;
        std::vector<asio::const_buffer> buffers;

        // 批量合并数据，减少 async_write 调用次数
        while (!writeQueue_.empty()) {
            auto data = writeQueue_.front();
            buffers.emplace_back(asio::buffer(*data));
            writeQueue_.pop_front();
            if (buffers.size() >= MAX_BATCH_SIZE) break; // 限制每次批量发送的数据量
        }

        ws_.async_write(buffers,
            asio::bind_executor(strand_,
                [self = shared_from_this()](beast::error_code ec, std::size_t /*bytes_transferred*/) {
                    if (!ec) {
                        self->doWrite();  // 继续发送队列中的下一个数据
                    } else {
                        spdlog::error("[WebsocketSession] Write error: {}", ec.message());
                        self->writing_ = false;
                    }
                }
            )
        );
    }

    bool handshake_verified_;
    http::request<http::string_body> req_;
    beast::flat_buffer buffer_;
    websocket::stream<beast::tcp_stream> ws_;
    std::shared_ptr<ChannelPipeline> pipeline_;
    // onClose 回调，当连接关闭时通知上层（例如 WebsocketServer 移除该会话）
    std::function<void(std::shared_ptr<WebsocketSession>)> onClose;
    websocket::stream<beast::basic_stream<tcp>>::executor_type strand_;
    std::deque<std::shared_ptr<std::string>> writeQueue_;
    bool writing_ = false;
    static constexpr std::size_t MAX_BATCH_SIZE = 5; // 每次最多批量发送 5 条消息
};
