/**
 * @file WSServer.cpp
 * @brief Multi-client WebSocket server implementation.
 */

#include "WSServer.h"
#include "WSFrame.h"
#include "SHA1.h"
#include "Base64.h"
#include "AdvancedErrorManagement.h"
#include "Select.h"
#include "Sleep.h"
#include "Threads.h"
#include "TimeoutType.h"
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

namespace StreamHub {

using MARTe::TimeoutType;
using MARTe::Sleep;

/*---------------------------------------------------------------------------*/
/* Thread trampolines                                                         */
/*---------------------------------------------------------------------------*/

struct AcceptThreadArg { WSServer *srv; };
struct ClientThreadArg { WSServer *srv; uint32 slot; };

static void AcceptThreadEntry(const void *arg) {
    WSServer *srv = reinterpret_cast<AcceptThreadArg *>(
                        const_cast<void *>(arg))->srv;
    delete reinterpret_cast<const AcceptThreadArg *>(arg);
    srv->AcceptLoop();
}

static void ClientReadEntry(const void *arg) {
    const ClientThreadArg *a = reinterpret_cast<const ClientThreadArg *>(arg);
    WSServer *srv  = a->srv;
    uint32    slot = a->slot;
    delete a;
    srv->ClientReadLoop(slot);
}

/*---------------------------------------------------------------------------*/
/* Tiny JSON helpers (no library, only what we need)                         */
/*---------------------------------------------------------------------------*/

/** Locate the first occurrence of pattern in s, return pointer or NULL. */
static const char *FindSubstr(const char *s, const char *pattern) {
    return strstr(s, pattern);
}

/*---------------------------------------------------------------------------*/
/* WSServer                                                                   */
/*---------------------------------------------------------------------------*/

WSServer::WSServer()
    : numClients(0u),
      liveReadThreads(0u),
      callback(static_cast<WSCommandCallback *>(0)),
      running(false),
      numAllowedOrigins(0u),
      acceptTid(MARTe::InvalidThreadIdentifier) {

    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        clients[i].sock   = static_cast<BasicTCPSocket *>(0);
        clients[i].active = false;
        clients[i].readTid = MARTe::InvalidThreadIdentifier;
    }
    for (uint32 i = 0u; i < WS_MAX_ORIGINS; i++) {
        allowedOrigins[i][0] = '\0';
    }
}

bool WSServer::AddAllowedOrigin(const char *origin) {
    if ((origin == static_cast<const char *>(0)) || (origin[0] == '\0')) {
        return false;
    }
    if (numAllowedOrigins >= WS_MAX_ORIGINS) { return false; }
    if (strlen(origin) >= WS_MAX_ORIGIN_LEN) { return false; }
    strcpy(allowedOrigins[numAllowedOrigins], origin);
    numAllowedOrigins++;
    return true;
}

WSServer::~WSServer() {
    (void) Stop();
}

bool WSServer::Start(uint16 port, WSCommandCallback *cb) {
    if (running) { return true; }
    callback = cb;
    running  = true;

    if (!tcpListener.Open()) {
        REPORT_ERROR_STATIC(MARTe::ErrorManagement::FatalError,
                            "WSServer: Failed to open TCP listener socket.");
        running = false;
        return false;
    }
    if (!tcpListener.Listen(port, static_cast<MARTe::int32>(WS_MAX_CLIENTS))) {
        REPORT_ERROR_STATIC(MARTe::ErrorManagement::FatalError,
                            "WSServer: Failed to listen on port %u.", static_cast<uint32>(port));
        running = false;
        return false;
    }

    AcceptThreadArg *arg = new AcceptThreadArg();
    arg->srv = this;
    acceptTid = MARTe::Threads::BeginThread(
        reinterpret_cast<MARTe::ThreadFunctionType>(AcceptThreadEntry),
        arg, MARTe::THREADS_DEFAULT_STACKSIZE, "WSAccept");

    REPORT_ERROR_STATIC(MARTe::ErrorManagement::Information,
                        "WSServer: Listening on port %u.", static_cast<uint32>(port));
    return true;
}

bool WSServer::Stop() {
    if (!running) { return true; }
    running = false;

    /* Close all client connections — their read threads wake out of select()
     * and unwind through FreeSlot. */
    (void) clientsMutex.FastLock();
    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        if (clients[i].active && (clients[i].sock != static_cast<BasicTCPSocket *>(0))) {
            clients[i].sock->Close();
        }
    }
    clientsMutex.FastUnLock();

    /* The accept loop polls WaitConnection with a 500 ms timeout, so it is out
     * of the listener by now. */
    Sleep::MSec(600u);
    tcpListener.Close();

    /* Wait for the read threads: they hold pointers to the sockets freed
     * below.  Bounded — leaking a socket at exit beats deleting one that a
     * wedged thread is still reading from. */
    static const uint32 kReadJoinMs = 3000u;
    for (uint32 waited = 0u; waited < kReadJoinMs; waited += 20u) {
        (void) clientsMutex.FastLock();
        const uint32 live = liveReadThreads;
        clientsMutex.FastUnLock();
        if (live == 0u) { break; }
        Sleep::MSec(20u);
    }

    /* Free any remaining slots */
    (void) clientsMutex.FastLock();
    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        if (clients[i].sock != static_cast<BasicTCPSocket *>(0)) {
            delete clients[i].sock;
            clients[i].sock   = static_cast<BasicTCPSocket *>(0);
            clients[i].active = false;
        }
    }
    numClients = 0u;
    clientsMutex.FastUnLock();
    return true;
}

/*---------------------------------------------------------------------------*/
/* Accept loop                                                                */
/*---------------------------------------------------------------------------*/

void WSServer::AcceptLoop() {
    while (running) {
        BasicTCPSocket *client = tcpListener.WaitConnection(
            TimeoutType(500u), static_cast<BasicTCPSocket *>(0));
        if (client == static_cast<BasicTCPSocket *>(0)) {
            /* Timeout — loop again to check running flag */
            continue;
        }

        /* HTTP upgrade */
        if (!UpgradeHTTP(client)) {
            client->Close();
            delete client;
            continue;
        }

        uint32 slot = AllocSlot(client);
        if (slot == WS_MAX_CLIENTS) {
            /* No room */
            static const char *full = "HTTP/1.1 503 Service Unavailable\r\n\r\n";
            uint32 l = static_cast<uint32>(strlen(full));
            (void) client->Write(full, l);
            client->Close();
            delete client;
            continue;
        }

        /* Notify StreamHub */
        if (callback != static_cast<WSCommandCallback *>(0)) {
            callback->OnWSClientConnected();
        }

        /* Start per-client read thread */
        (void) clientsMutex.FastLock();
        liveReadThreads++;
        clientsMutex.FastUnLock();

        ClientThreadArg *arg = new ClientThreadArg();
        arg->srv  = this;
        arg->slot = slot;
        clients[slot].readTid = MARTe::Threads::BeginThread(
            reinterpret_cast<MARTe::ThreadFunctionType>(ClientReadEntry),
            arg, MARTe::THREADS_DEFAULT_STACKSIZE, "WSClient");
    }
}

/*---------------------------------------------------------------------------*/
/* HTTP upgrade                                                               */
/*---------------------------------------------------------------------------*/

bool WSServer::UpgradeHTTP(BasicTCPSocket *sock) {
    /* Read raw HTTP request until \r\n\r\n */
    static const uint32 kMaxHdr = 4096u;
    char hdrBuf[4096];
    uint32 totalRead = 0u;

    while (totalRead < kMaxHdr - 1u) {
        uint32 rem = kMaxHdr - 1u - totalRead;
        bool ok = sock->Read(hdrBuf + totalRead, rem,
                             TimeoutType(3000u));
        if (!ok || rem == 0u) { return false; }
        totalRead += rem;
        hdrBuf[totalRead] = '\0';
        if (strstr(hdrBuf, "\r\n\r\n") != static_cast<char *>(0)) { break; }
    }

    /* Origin validation (CSWSH / CSRF defence, RFC 6455 §10.2).
     * If an Origin header is present it must either be on the configured
     * allowlist or its host must match the Host header (same-origin).
     * Non-browser clients (no Origin) are allowed. */
    const char *originHdr = FindSubstr(hdrBuf, "Origin:");
    if (originHdr != static_cast<const char *>(0)) {
        originHdr += 7; /* skip "Origin:" */
        while (*originHdr == ' ') { originHdr++; }

        /* Full origin value "scheme://host[:port]", for the allowlist. */
        char originFull[WS_MAX_ORIGIN_LEN];
        uint32 ofLen = 0u;
        while ((originHdr[ofLen] != '\r') && (originHdr[ofLen] != '\n') &&
               (originHdr[ofLen] != '\0') && (ofLen < (WS_MAX_ORIGIN_LEN - 1u))) {
            originFull[ofLen] = originHdr[ofLen];
            ofLen++;
        }
        originFull[ofLen] = '\0';
        bool allowed = false;
        for (uint32 i = 0u; (i < numAllowedOrigins) && !allowed; i++) {
            if (strcmp(originFull, allowedOrigins[i]) == 0) { allowed = true; }
        }

        /* Extract the host part of Origin: "scheme://host[:port]" */
        char originHost[256];
        uint32 ohLen = 0u;
        const char *op = originHdr;
        /* Skip scheme:// */
        const char *schemeEnd = strstr(op, "://");
        if (schemeEnd != static_cast<const char *>(0)) { op = schemeEnd + 3; }
        while (*op != '\r' && *op != '\n' && *op != '\0' &&
               *op != '/' && ohLen < 255u) {
            originHost[ohLen++] = *op++;
        }
        originHost[ohLen] = '\0';

        /* Extract Host header value */
        const char *hostHdr = FindSubstr(hdrBuf, "Host:");
        if (!allowed && (hostHdr != static_cast<const char *>(0))) {
            hostHdr += 5; /* skip "Host:" */
            while (*hostHdr == ' ') { hostHdr++; }
            char hostVal[256];
            uint32 hvLen = 0u;
            while (*hostHdr != '\r' && *hostHdr != '\n' &&
                   *hostHdr != '\0' && hvLen < 255u) {
                hostVal[hvLen++] = *hostHdr++;
            }
            hostVal[hvLen] = '\0';
            if (strcmp(originHost, hostVal) != 0) {
                /* Cross-origin — reject the upgrade */
                const char *forbidden =
                    "HTTP/1.1 403 Forbidden\r\n"
                    "Content-Type: text/plain\r\n"
                    "Connection: close\r\n"
                    "\r\nOrigin not allowed\r\n";
                uint32 forbLen = static_cast<uint32>(strlen(forbidden));
                (void) sock->Write(forbidden, forbLen);
                return false;
            }
        }
    }

    /* Find Sec-WebSocket-Key */
    const char *keyHdr = FindSubstr(hdrBuf, "Sec-WebSocket-Key:");
    if (keyHdr == static_cast<const char *>(0)) { return false; }
    keyHdr += 18; /* skip "Sec-WebSocket-Key:" */
    while (*keyHdr == ' ') { keyHdr++; }

    char key[64];
    uint32 keyLen = 0u;
    while (keyHdr[keyLen] != '\r' && keyHdr[keyLen] != '\n' &&
           keyHdr[keyLen] != '\0' && keyLen < 63u) {
        key[keyLen] = keyHdr[keyLen];
        keyLen++;
    }
    key[keyLen] = '\0';

    /* Compute Accept = Base64(SHA1(key + GUID)) */
    static const char *kGUID = "258EAFA5-E914-47DA-95CA-C5AB0DC85B11";
    char concat[128];
    snprintf(concat, sizeof(concat), "%s%s", key, kGUID);

    uint8 digest[20];
    SHA1(reinterpret_cast<const uint8 *>(concat),
         static_cast<uint32>(strlen(concat)), digest);

    char accept[64];
    Base64Encode(digest, 20u, accept);

    /* Send 101 response */
    char resp[512];
    snprintf(resp, sizeof(resp),
        "HTTP/1.1 101 Switching Protocols\r\n"
        "Upgrade: websocket\r\n"
        "Connection: Upgrade\r\n"
        "Sec-WebSocket-Accept: %s\r\n"
        "\r\n", accept);
    uint32 respLen = static_cast<uint32>(strlen(resp));
    return sock->Write(resp, respLen);
}

/*---------------------------------------------------------------------------*/
/* Client read loop                                                           */
/*---------------------------------------------------------------------------*/

void WSServer::ClientReadLoop(uint32 slotIdx) {
    WSClientSlot &slot = clients[slotIdx];
    BasicTCPSocket *sock = slot.sock;

    /* Receive buffer: WS_MAX_RECV_PAYLOAD + max header (14: 2 + 8 ext-length +
     * 4 mask) + 1 spare byte for in-place NUL-termination of the payload. */
    static const uint32 kRecvBuf = WS_MAX_RECV_PAYLOAD + 14u + 1u;
    uint8 *buf = new uint8[kRecvBuf];
    uint32 filled = 0u;

    while (running && slot.active) {
        uint32 want = kRecvBuf - filled;
        if (want == 0u) {
            /* Buffer full — discard old frame (shouldn't happen with reasonable clients) */
            filled = 0u;
            continue;
        }

        /* Wait for readability before reading.  BasicTCPSocket::Read reports a
         * timeout and a closed peer identically (false, zero bytes), so polling
         * it on its own cannot end the loop: once the client goes away recv
         * returns immediately and forever, and the thread spins at 100% CPU
         * until it starves the rest of the hub.  select() tells the two apart —
         * readable followed by no data is end of stream.  A wait consumes the
         * handle set, hence a fresh Select each pass. */
        MARTe::Select sel;
        if (!sel.AddReadHandle(*sock)) { break; }
        const MARTe::int32 ready = sel.WaitUntil(TimeoutType(500u));
        if (ready == 0) { continue; }   /* idle client — recheck running */
        if (ready < 0) { break; }       /* socket closed or errored */

        /* Readable: this returns at once, and only fails at end of stream. */
        if (!sock->Read(reinterpret_cast<char *>(buf + filled), want,
                        TimeoutType(500u))) {
            break;
        }
        filled += want;

        /* Parse as many complete frames as possible */
        uint32 consumed = 0u;
        while (consumed < filled) {
            const uint8 *frameStart = buf + consumed;
            uint32 available = filled - consumed;

            WSFrameHeader hdr;
            if (!WSParseHeader(frameStart, available, hdr)) {
                break; /* incomplete header — need more data */
            }
            if (hdr.payloadLen > static_cast<uint64>(WS_MAX_RECV_PAYLOAD)) {
                /* Payload too large — close connection */
                goto client_done;
            }
            uint32 plen = static_cast<uint32>(hdr.payloadLen);
            uint32 totalFrameBytes = hdr.headerSize + plen;
            if (available < totalFrameBytes) {
                break; /* incomplete payload — need more data */
            }

            /* Unmask payload */
            uint8 *payload = const_cast<uint8 *>(frameStart) + hdr.headerSize;
            if (hdr.masked) {
                WSUnmask(payload, plen, hdr.maskKey);
            }

            /* Handle opcode */
            if (hdr.opcode == WS_OPCODE_CLOSE) {
                goto client_done;
            } else if (hdr.opcode == WS_OPCODE_PING) {
                /* Send pong */
                (void) slot.writeMutex.FastLock();
                SendFrame(slot, WS_OPCODE_PONG, payload, plen);
                slot.writeMutex.FastUnLock();
            } else if (hdr.opcode == WS_OPCODE_TEXT) {
                if (callback != static_cast<WSCommandCallback *>(0)) {
                    /* NUL-terminate in place for the callback, but restore the
                     * byte afterwards: when several client frames coalesce in
                     * one TCP read, payload[plen] is the first header byte of
                     * the NEXT frame. */
                    uint8 savedByte = payload[plen];
                    payload[plen] = '\0'; /* safe: buf has extra byte */
                    callback->OnWSCommand(reinterpret_cast<const char *>(payload), plen, slotIdx);
                    payload[plen] = savedByte;
                }
            }
            /* Binary frames from client are ignored (server only sends binary) */

            consumed += totalFrameBytes;
        }

        /* Compact buffer: move unconsumed bytes to front */
        if (consumed > 0u && consumed < filled) {
            memmove(buf, buf + consumed, filled - consumed);
        }
        filled = (consumed <= filled) ? (filled - consumed) : 0u;
    }

client_done:
    delete[] buf;
    if (callback != static_cast<WSCommandCallback *>(0)) {
        callback->OnWSClientDisconnected();
    }
    FreeSlot(slotIdx);

    (void) clientsMutex.FastLock();
    if (liveReadThreads > 0u) { liveReadThreads--; }
    clientsMutex.FastUnLock();
}

/*---------------------------------------------------------------------------*/
/* Broadcast helpers                                                          */
/*---------------------------------------------------------------------------*/

void WSServer::BroadcastText(const char *json, uint32 len) {
    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        if (!clients[i].active) { continue; }
        (void) clients[i].writeMutex.FastLock();
        if (clients[i].active) {
            (void) SendFrame(clients[i], WS_OPCODE_TEXT,
                             reinterpret_cast<const uint8 *>(json), len);
        }
        clients[i].writeMutex.FastUnLock();
    }
}

void WSServer::BroadcastBinary(const uint8 *buf, uint32 len) {
    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        if (!clients[i].active) { continue; }
        (void) clients[i].writeMutex.FastLock();
        if (clients[i].active) {
            (void) SendFrame(clients[i], WS_OPCODE_BINARY, buf, len);
        }
        clients[i].writeMutex.FastUnLock();
    }
}

void WSServer::SendText(uint32 slotIdx, const char *json, uint32 len) {
    if (slotIdx >= WS_MAX_CLIENTS) { return; }
    WSClientSlot &slot = clients[slotIdx];
    if (!slot.active) { return; }
    (void) slot.writeMutex.FastLock();
    if (slot.active) {
        (void) SendFrame(slot, WS_OPCODE_TEXT,
                         reinterpret_cast<const uint8 *>(json), len);
    }
    slot.writeMutex.FastUnLock();
}

uint32 WSServer::ClientCount() const {
    (void) clientsMutex.FastLock();
    uint32 c = numClients;
    clientsMutex.FastUnLock();
    return c;
}

/*---------------------------------------------------------------------------*/
/* Internal helpers                                                           */
/*---------------------------------------------------------------------------*/

bool WSServer::SendFrame(WSClientSlot &slot, uint8 opcode,
                         const uint8 *payload, uint32 payloadLen) {
    if (!slot.active || (slot.sock == static_cast<BasicTCPSocket *>(0))) {
        return false;
    }
    uint8 hdrBuf[10];
    uint32 hdrLen = WSEncodeHeader(hdrBuf, opcode,
                                   static_cast<uint64>(payloadLen));

    /* Write header */
    uint32 sz = hdrLen;
    if (!slot.sock->Write(reinterpret_cast<const char *>(hdrBuf), sz)) {
        return false;
    }
    /* Write payload */
    sz = payloadLen;
    if (payloadLen > 0u) {
        if (!slot.sock->Write(reinterpret_cast<const char *>(payload), sz)) {
            return false;
        }
    }
    return true;
}

uint32 WSServer::AllocSlot(BasicTCPSocket *sock) {
    (void) clientsMutex.FastLock();
    uint32 idx = WS_MAX_CLIENTS;
    for (uint32 i = 0u; i < WS_MAX_CLIENTS; i++) {
        if (!clients[i].active) {
            clients[i].sock   = sock;
            clients[i].active = true;
            clients[i].readTid = MARTe::InvalidThreadIdentifier;
            idx = i;
            numClients++;
            break;
        }
    }
    clientsMutex.FastUnLock();
    return idx;
}

void WSServer::FreeSlot(uint32 idx) {
    if (idx >= WS_MAX_CLIENTS) { return; }
    /* HI-5: acquire writeMutex before modifying active/sock to prevent
     * use-after-free when BroadcastText/BroadcastBinary are iterating. */
    (void) clients[idx].writeMutex.FastLock();
    (void) clientsMutex.FastLock();
    if (clients[idx].active) {
        clients[idx].active = false;
        if (clients[idx].sock != static_cast<BasicTCPSocket *>(0)) {
            clients[idx].sock->Close();
            delete clients[idx].sock;
            clients[idx].sock = static_cast<BasicTCPSocket *>(0);
        }
        if (numClients > 0u) { numClients--; }
    }
    clientsMutex.FastUnLock();
    clients[idx].writeMutex.FastUnLock();
}

} /* namespace StreamHub */