/*
 * (c) Copyright, Real-Time Innovations, 2024.
 * All rights reserved.
 *
 * No duplications, whole or partial, manual or electronic, may be made
 * without express written permission.  Any such copies, or
 * revisions thereof, must display this notice unaltered.
 * This code contains trade secrets of Real-Time Innovations, Inc.
 */

#ifndef CIVETWEB_WEBSOCKET_HANDLER_H
#define CIVETWEB_WEBSOCKET_HANDLER_H

#include <CivetServer.h>
#include <civetweb.h>

#include <rti/apputils/http/server/CivetWebSocketConnection.hpp>
#include <rti/apputils/http/server/detail/WebSocketHandler.hpp>

namespace rti { namespace apputils { namespace websocket { namespace civetweb {

/**
 * This class implements the CivetWebSocketHandler interface, which is part
 * of the CivetWeb API. This class has a WebSocketHandler. It allows us
 * to isolate the WebSocket handling that CivetWeb does and the WebSocket
 * handling that our WebSocketAPI implements
 */
class CivetWebSocketHandlerImpl : public CivetWebSocketHandler {
public:
    CivetWebSocketHandlerImpl(WebSocketHandler& handler) : handler_(handler)
    { }

    /**
     * This method handles new client connections to the server. For such
     * purpose, it creates a new `WebSocketConnection` and attaches it to the
     * internal CivetWeb `mg_connection`. This mechanisms allows existing
     * connections to be reused in the future.
     * @param server CivetWeb server. Unused parameter.
     * @param conn CivetWeb connection.
     * @return Boolean value indicating the success or failure of the operation.
     */
    bool handleConnection(CivetServer *, const struct mg_connection *conn)
    {
        WebSocketConnection *client_conn =
                new CivetWebSocketConnection(const_cast<mg_connection *>(conn));

        handler_.handle_connection(*client_conn);
        mg_set_user_connection_data(const_cast<mg_connection *>(conn), client_conn);

        /*
         * client_conn is not being leaked here, it's stored indirectly
         * in the user connection data that can be get by calling
         * mg_get_user_connection_data as we can see on handleClose.
         */
        /* coverity[leaked_storage : FALSE] */
        return true;
    }
    /**
     * This method is called when a client connection has ended the handshake
     * protocol with success.
     * @param server CivetWeb server. Unused parameter.
     * @param conn CivetWeb connection
     */
    void handleReadyState(CivetServer *, struct mg_connection *conn)
    {
        WebSocketConnection *client = static_cast<WebSocketConnection *>(
                mg_get_user_connection_data(conn));
        handler_.handle_ready_state(*client);
    }

    /**
     * This method is called every time a client sends a message to the server
     * through the WebSocket connection.
     * @param server CivetWeb server. Unused parameter
     * @param conn CivetWeb connection
     * @param bits A byte with the kind of data (TEXT, PONG, ...)
     * @param data String containing the received message, it is not
     * null-terminated
     * @param data_size String size, the size of the string without the null
     * character.
     */
    bool handleData(
            CivetServer *,
            struct mg_connection *conn,
            int bits,
            char *data,
            size_t data_len)
    {
        StringBuffer msg;
        WebSocketConnection *client = static_cast<WebSocketConnection *>(
                mg_get_user_connection_data(conn));
        // The last byte of the bits say what kind of message it is
        switch (((unsigned char) bits) & 0x0F) {
        case MG_WEBSOCKET_OPCODE_TEXT:
            // This marks the end of the string
            *(data + data_len) = '\0';
            msg = StringBuffer(data);
            try {
                handler_.handle_data(*client, msg);
            } catch (const RequestException& ex) {
                DDSLog_exception(&RTI_LOG_ANY_FAILURE_s, ex.what());
                client->send_request_error(ex);
            } catch (const rti::apputils::http::HttpException& ex) {
                DDSLog_exception(&RTI_LOG_ANY_FAILURE_s, ex.what());
                client->send_error(ex);
            } catch (...) {
                std::string message = "Unexpected exception. The connection will be closed.";
                DDSLog_exception(&RTI_LOG_ANY_FAILURE_s, message.c_str());
                message.append(": ");
                message.append(RTI_FUNCTION_NAME);
                client->send(message);

                // Close the connection (Especially to avoid vulnerabilities)
                client->state(WebSocketStateKind::CLOSE);
                client->close();
            }
            break;
        case MG_WEBSOCKET_OPCODE_PING:
            // This marks the end of the string
            data[data_len] = '\0';
            msg = StringBuffer(data);
            // just send back the message data in a pong frame
            client->send_pong(msg);
            break;
        default:
            break;
        }

        return true;
    }

    /**
     * This method is called when a client closes the connection. The method
     * deletes the connection object.
     * @param server CivetWeb Server
     * @param conn CivetWeb connection
     */
    void handleClose(CivetServer *, const struct mg_connection *conn)
    {
        struct mg_context *ctx = mg_get_context(conn);
        WebSocketConnection *client =
                static_cast<WebSocketConnection *>(mg_get_user_connection_data(conn));

        if (client->state() != WebSocketStateKind::CLOSE){
            handler_.handle_close(*client);
        }

        mg_lock_context(ctx);
        delete client;
        mg_set_user_connection_data(const_cast<mg_connection *>(conn), NULL);
        mg_unlock_context(ctx);
    }

    ~CivetWebSocketHandlerImpl()
    {
    }

private:
    WebSocketHandler& handler_;
};

}}}}  // namespace rti::webdds::websocket::civetweb


#endif