// Copyright (c) 2025 Huawei Technologies Co., Ltd.
// openUBMC is licensed under Mulan PSL v2.
// You can use this software according to the terms and conditions of the Mulan PSL v2.
// You may obtain a copy of Mulan PSL v2 at:
//         http://license.coscl.org.cn/MulanPSL2
// THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
// EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
// MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
// See the Mulan PSL v2 for more details.
//
// Description: 实现KVM客户端websocket连接

#include "WebSocketClient.hpp"

WebSocketClient::WebSocketClient()
    : io_context_(), ssl_context_(ssl::context::tlsv13_client), resolver_(net::make_strand(io_context_)), buffer_(),
      heartbeat_manager_(io_context_, std::chrono::seconds(HEARTBEAT_INTERVAL),
                         [this](const std::vector<uint8_t>& message) {
                             SendHeartbeat(message);
                         })
{}

WebSocketClient::~WebSocketClient() = default;

void WebSocketClient::Connect(std::string host, std::string port, std::string uri)
{
    // 如果 ws_ 已经存在，先销毁
    ws_.reset();

    // 创建新的 ws_ 对象
    if (io_context_.stopped()) {
        io_context_.restart();
    }
    ws_.emplace(net::make_strand(io_context_), ssl_context_);
    if (!SSL_set_tlsext_host_name(ws_->next_layer().native_handle(), host.c_str())) {
        beast::error_code ec{static_cast<int>(::ERR_get_error()), net::error::get_ssl_category()};
        on_error(std::runtime_error("SSL_set_tlsext_host_name failed: " + ec.message()));
        return;
    }

    host_ = host;
    port_ = port;
    uri_  = uri;

    resolver_.async_resolve(host, port, beast::bind_front_handler(&WebSocketClient::on_resolve, this));
    ws_->binary(true);
}

void WebSocketClient::Disconnect()
{
    try {
        beast::error_code ec;
        ws_->close(websocket::close_code::normal, ec);
        if (ec) {
            if (ec == net::error::operation_aborted || ec == net::error::not_connected || ec == net::error::eof) {
                // 断开连接及停止I/O服务时会出现Operation Canceled, end of failed, not
                // connected的错误，属于正常行为，因此忽略该报错
                return;
            }
            on_error(std::runtime_error("Error during disconnect: " + ec.message()));
        }
        ws_.reset();
        io_context_.stop();
        buffer_.consume(buffer_.size());
    } catch (const std::exception& e) {
        on_error(e);
    }
}

void WebSocketClient::Send(const std::vector<uint8_t>& binary_data)
{
    {
        std::lock_guard<std::mutex> lock(sendMutex_);
        sendQueue_.push(binary_data);
    }

    // 如果当前没有正在进行的写操作，则启动队列处理
    {
        std::lock_guard<std::mutex> lock(sendMutex_);
        if (!isWriting_) {
            isWriting_ = true;

            // 延迟调用 ProcessSendQueue，避免递归调用
            net::post(io_context_, [this]() {
                ProcessSendQueue();
            });
        }
    }
}

void WebSocketClient::ProcessSendQueue()
{
    std::shared_ptr<std::vector<uint8_t>> data;

    {
        std::unique_lock<std::mutex> lock(sendMutex_);
        if (sendQueue_.empty()) {
            isWriting_ = false; // 如果队列为空，重置写操作标志
            return;
        }

        data = std::make_shared<std::vector<uint8_t>>(std::move(sendQueue_.front()));
        sendQueue_.pop();
    }

    // 启动第一个异步写操作
    ws_->async_write(net::buffer(*data),
                     [this, data](const boost::system::error_code& ec, std::size_t bytes_transferred) {
        on_write(ec, bytes_transferred);
    });
}

void WebSocketClient::Run()
{
    try {
        if (io_context_.stopped()) {
            io_context_.restart();
        }
        io_context_.run();
    } catch (const std::exception& e) {
        std::cerr << "Error in io_context: " << e.what() << std::endl;
        on_error(e);
    }
}

void WebSocketClient::StartHeartbeat(std::chrono::seconds interval)
{
    heartbeat_manager_.SetInterval(interval);
    heartbeat_manager_.Start();
}

void WebSocketClient::StopHeartbeat()
{
    heartbeat_manager_.Stop();
}

void WebSocketClient::SendHeartbeat(const std::vector<uint8_t>& message)
{
    ws_->async_write(net::buffer(message), [this](const boost::system::error_code& ec, std::size_t) {
        if (ec) {
            on_error(std::runtime_error("Failed to send heartbeat: " + ec.message()));
            return;
        }
    });
}

void WebSocketClient::on_resolve(const boost::system::error_code& ec, tcp::resolver::results_type results)
{
    if (ec) {
        on_error(std::runtime_error("Resolve failed: " + ec.message()));
        return;
    }

    beast::get_lowest_layer(*ws_).async_connect(results, beast::bind_front_handler(&WebSocketClient::on_connect, this));
}

void WebSocketClient::on_connect(const boost::system::error_code& ec, const tcp::endpoint& ep)
{
    if (ec) {
        on_error(std::runtime_error("Connect failed: " + ec.message()));
        return;
    }

    host_ += ':' + std::to_string(ep.port());
    ws_->next_layer().async_handshake(ssl::stream_base::client,
                                      beast::bind_front_handler(&WebSocketClient::on_ssl_handshake, this));
}

void WebSocketClient::on_ssl_handshake(const boost::system::error_code& ec)
{
    if (ec) {
        on_error(std::runtime_error("SSL handshake failed: " + ec.message()));
        return;
    }

    beast::get_lowest_layer(*ws_).expires_never();

    ws_->set_option(websocket::stream_base::decorator([](websocket::request_type& req) {
        req.set(http::field::user_agent, std::string(BOOST_BEAST_VERSION_STRING) + " websocket-client-async-ssl");
    }));
    ws_->read_message_max(MAX_READ_BUFFER);
    ws_->write_buffer_bytes(MAX_WRITE_BUFFER);

    ws_->async_handshake(host_, uri_, beast::bind_front_handler(&WebSocketClient::on_handshake, this));
}

void WebSocketClient::on_handshake(const boost::system::error_code& ec)
{
    if (ec) {
        on_error(std::runtime_error("Handshake failed: " + ec.message()));
        return;
    }

    do_read();
}

void WebSocketClient::do_read()
{
    ws_->async_read(buffer_, beast::bind_front_handler(&WebSocketClient::on_read, this));
}

void WebSocketClient::on_read(const boost::system::error_code& ec, std::size_t)
{
    if (ec) {
        if (ec == net::error::operation_aborted || ec == net::error::eof) {
            // 断开连接及停止I/O服务时会出现Operation Canceled或End Of File错误，属于正常行为，因此忽略该报错
            return;
        }
        if (ec != websocket::error::closed) {
            on_error(std::runtime_error("Read failed: " + ec.message()));
        } else {
            on_close(ws_->reason());
        }
        return;
    }

    do_read();
    OnMessage(beast::buffers_to_string(buffer_.data()));
    buffer_.consume(buffer_.size());
}

void WebSocketClient::on_write(beast::error_code ec, std::size_t bytes_transferred)
{
    if (ec) {
        on_error(std::runtime_error("Write failed: " + ec.message()));
        {
            std::lock_guard<std::mutex> lock(sendMutex_);
            isWriting_ = false; // 写操作失败时重置标志
        }
        return;
    }

    // 延迟调用 ProcessSendQueue，避免递归调用
    net::post(io_context_, [this]() {
        ProcessSendQueue();
    });
}

void WebSocketClient::on_close(const websocket::close_reason& cr)
{
    OnClose(cr);
    StopHeartbeat();
    io_context_.stop();
}

void WebSocketClient::on_error(const std::exception& e)
{
    OnError(e);
    StopHeartbeat();
    io_context_.stop();
}