/*
 * Copyright 2024 Rive
 */

#include "tcp_client.hpp"

#include "rive/math/math_types.hpp"
#include <stdio.h>
#include <string.h>

#ifndef EXTERNAL_TCP_CLIENT_DEFINITION

std::unique_ptr<TCPClient> TCPClient::Connect(
    const char* serverAddress /*server:port*/)
{
    bool success;
    auto tcpClient =
        std::unique_ptr<TCPClient>(new TCPClient(serverAddress, &success));
    if (success)
    {
        return tcpClient;
    }
    return nullptr;
}

static SOCKET invalid_socket()
{
#ifdef _WIN32
    return INVALID_SOCKET;
#else
    return -1;
#endif
}

static bool is_socket_valid(SOCKET sockfd)
{
#ifdef _WIN32
    return sockfd != INVALID_SOCKET;
#else
    return sockfd > 0;
#endif
}

static void close_socket(SOCKET sockfd)
{
#ifdef __EMSCRIPTEN__
    constexpr static int WEBSOCKET_NORMAL_CLOSURE = 1000;
    emscripten_websocket_close(sockfd, WEBSOCKET_NORMAL_CLOSURE, "finished");
#elif defined(_WIN32)
    closesocket(sockfd);
#else
    close(sockfd);
#endif
}

TCPClient::TCPClient(const char* serverAddress /*server:port*/, bool* success) :
    m_serverAddress(serverAddress), m_sockfd(invalid_socket())
{
    *success = false;

    char hostname[256];
    uint16_t port;
    if (sscanf(serverAddress, "%255[^:]:%hu", hostname, &port) != 2)
    {
        return;
    }

#ifdef __EMSCRIPTEN__
    // WASM has to use websockets instead of TCP.
    if (!emscripten_websocket_is_supported())
    {
        fprintf(stderr,
                "ERROR: WebSockets are not supported in this browser.\n");
        abort();
    }

    auto url = std::string("ws://") + serverAddress;
    EmscriptenWebSocketCreateAttributes attr = {
        .url = url.c_str(),
        .protocols = NULL,
        .createOnMainThread = EM_TRUE,
    };

    m_sockfd = emscripten_websocket_new(&attr);
#else
#ifdef _WIN32
    WSADATA wsaData;
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        fprintf(stderr, "WSAStartup() failed.\n");
        abort();
    }
#endif

    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(port);
    addr.sin_addr.s_addr = inet_addr(hostname);
    m_sockfd = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
#endif

    if (!is_socket_valid(m_sockfd))
    {
        fprintf(stderr,
                "Unable to create socket for %s:%u (%s)\n",
                hostname,
                port,
                strerror(GetLastError()));
        abort();
    }

#ifdef __EMSCRIPTEN__
    emscripten_websocket_set_onopen_callback(m_sockfd, this, OnWebSocketOpen);
    emscripten_websocket_set_onmessage_callback(m_sockfd,
                                                this,
                                                OnWebSocketMessage);
    emscripten_websocket_set_onerror_callback(m_sockfd, this, OnWebSocketError);

    while (m_webSocketStatus != WebSocketStatus::open)
    {
        if (m_webSocketStatus == WebSocketStatus::error)
        {
            fprintf(stderr,
                    "Unable to connect to WebSocket at %s\n",
                    url.c_str());
            abort();
        }
        // Busy wait until our connection is established.
        //
        // NOTE: emscripten_sleep() is an async operation that yields control
        // back to the browser to process.
        emscripten_sleep(10);
    }
#else
    if (connect(m_sockfd, (struct sockaddr*)&addr, sizeof(addr)) < 0)
    {
        fprintf(stderr,
                "Unable to connect to TCP server at %s:%u (%s)\n",
                hostname,
                port,
                strerror(GetLastError()));
        abort();
    }
#endif

    *success = true;
}

TCPClient::~TCPClient()
{
    if (is_socket_valid(m_sockfd))
    {
        close_socket(m_sockfd);
    }
#ifdef _WIN32
    WSACleanup();
#endif
}

std::unique_ptr<TCPClient> TCPClient::clone() const
{
    auto clone = TCPClient::Connect(serverAddress());
    if (clone == nullptr)
    {
        fprintf(stderr, "Failed to clone connection to %s\n", serverAddress());
        abort();
    }
    return clone;
}

#ifdef __EMSCRIPTEN__
// Derived from:
// https://stackoverflow.com/questions/180947/base64-decode-snippet-in-c
static std::string base64_encode(const char* buf, unsigned int bufLen)
{
    static const std::string base64_chars =
        "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789+/";

    std::string ret;
    ret.reserve(((bufLen + 2) / 3) * 4); // 8/6 of the byte length,
                                         // rounded to a multiple of 4.
    int i = 0;
    int j = 0;
    char char_array_3[3];
    char char_array_4[4];

    while (bufLen--)
    {
        char_array_3[i++] = *(buf++);
        if (i == 3)
        {
            char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
            char_array_4[1] = ((char_array_3[0] & 0x03) << 4) +
                              ((char_array_3[1] & 0xf0) >> 4);
            char_array_4[2] = ((char_array_3[1] & 0x0f) << 2) +
                              ((char_array_3[2] & 0xc0) >> 6);
            char_array_4[3] = char_array_3[2] & 0x3f;

            for (i = 0; (i < 4); i++)
                ret += base64_chars[char_array_4[i]];
            i = 0;
        }
    }

    if (i)
    {
        for (j = i; j < 3; j++)
            char_array_3[j] = '\0';

        char_array_4[0] = (char_array_3[0] & 0xfc) >> 2;
        char_array_4[1] =
            ((char_array_3[0] & 0x03) << 4) + ((char_array_3[1] & 0xf0) >> 4);
        char_array_4[2] =
            ((char_array_3[1] & 0x0f) << 2) + ((char_array_3[2] & 0xc0) >> 6);
        char_array_4[3] = char_array_3[2] & 0x3f;

        for (j = 0; (j < i + 1); j++)
            ret += base64_chars[char_array_4[j]];

        while ((i++ < 3))
            ret += '=';
    }

    return ret;
}

uint32_t TCPClient::send(const char* data, uint32_t size)
{
    // Base64-encode our data and send it as text, in order to avoid landmines
    // in binary data transmission.
    std::string base64 = base64_encode(data, size);
    if (emscripten_websocket_send_utf8_text(m_sockfd, base64.c_str()) < 0)
    {
        fprintf(stderr,
                "Failed to send %zu base64 chars to websocket.\n",
                base64.size());
        abort();
    }
    return size;
}
#else
uint32_t TCPClient::send(const char* data, uint32_t size)
{
    size_t sent = ::send(m_sockfd, data, size, 0);
    if (sent == -1)
    {
        fprintf(stderr, "Failed to send %u bytes to server.\n", size);
        abort();
    }
    return rive::math::lossless_numeric_cast<uint32_t>(sent);
}
#endif

uint32_t TCPClient::recv(char* buff, uint32_t size)
{
#ifdef __EMSCRIPTEN__
    while (m_serverMessages.size() < size)
    {
        // Busy wait until the server sends us our data. This isn't ideal, but
        // we always expect immediate responses with our testing protocol.
        //
        // NOTE: emscripten_sleep() is an async operation that yields control
        // back to the browser to process.
        emscripten_sleep(10);
    }
    std::copy(m_serverMessages.begin(), m_serverMessages.begin() + size, buff);
    m_serverMessages.erase(m_serverMessages.begin(),
                           m_serverMessages.begin() + size);
    return size;
#else
    return rive::math::lossless_numeric_cast<uint32_t>(
        ::recv(m_sockfd, buff, size, 0));
#endif
}

void TCPClient::sendall(const void* data, size_t size)
{
    const char* cdata = reinterpret_cast<const char*>(data);
    while (size != 0)
    {
        uint32_t sent = send(cdata,
                             rive::math::lossless_numeric_cast<uint32_t>(
                                 std::min<size_t>(size, 4096)));
        size -= sent;
        cdata += sent;
    }
}

void TCPClient::recvall(void* buff, size_t size)
{
    char* cbuff = reinterpret_cast<char*>(buff);
    while (size != 0)
    {
        uint32_t read = recv(cbuff,
                             rive::math::lossless_numeric_cast<uint32_t>(
                                 std::min<size_t>(size, 4096)));
        size -= read;
        cbuff += read;
    }
}

void TCPClient::send4(uint32_t value)
{
    uint32_t netValue = htonl(value);
    sendall(&netValue, 4);
}

uint32_t TCPClient::recv4()
{
    uint32_t netValue;
    recvall(&netValue, 4);
    return ntohl(netValue);
}

void TCPClient::sendHandshake() { send4(HANDSHAKE_TOKEN); }

void TCPClient::recvHandshake()
{
    uint32_t token = recv4();
    if (token != HANDSHAKE_TOKEN)
    {
        fprintf(stderr, "Bad handshake\n");
        abort();
    }
}

void TCPClient::sendString(const std::string& str)
{
    uint32_t length = rive::math::lossless_numeric_cast<uint32_t>(str.length());
    send4(length);
    send(str.c_str(), length);
}

std::string TCPClient::recvString()
{
    std::string str;
    uint32_t length = recv4();
    str.resize(length);
    recvall(str.data(), length);
    return str;
}

#ifdef __EMSCRIPTEN__
EM_BOOL TCPClient::OnWebSocketOpen(int eventType,
                                   const EmscriptenWebSocketOpenEvent* e,
                                   void* userData)
{
    auto this_ = reinterpret_cast<TCPClient*>(userData);
    assert(e->socket == this_->m_sockfd);
    this_->m_webSocketStatus = WebSocketStatus::open;
    return EM_TRUE;
}

EM_BOOL TCPClient::OnWebSocketMessage(int eventType,
                                      const EmscriptenWebSocketMessageEvent* e,
                                      void* userData)
{
    auto this_ = reinterpret_cast<TCPClient*>(userData);
    assert(e->socket == this_->m_sockfd);
    this_->m_serverMessages.insert(this_->m_serverMessages.end(),
                                   e->data,
                                   e->data + e->numBytes);
    return EM_TRUE;
}

EM_BOOL TCPClient::OnWebSocketError(int eventType,
                                    const EmscriptenWebSocketErrorEvent* e,
                                    void* userData)
{
    auto this_ = reinterpret_cast<TCPClient*>(userData);
    assert(e->socket == this_->m_sockfd);
    this_->m_webSocketStatus = WebSocketStatus::error;
    return EM_TRUE;
}
#endif

#endif