/**
 * @file tcp_server.c
 * @brief Indexed lwIP RAW TCP server manager.
 */

#include "tcp_server.h"

#include "../Drivers/LwIP/src/include/lwip/pbuf.h"
#include "../Drivers/LwIP/src/include/lwip/tcp.h"

#include "SEGGER_RTT.h"

#include <string.h>

typedef struct {
    struct tcp_pcb *listen_pcb;
    struct tcp_pcb *client_pcb;
    uint8_t rx_ring[TCP_SERVER_RX_BUFFER_SIZE];
    uint16_t rx_head;
    uint16_t rx_tail;
    struct pbuf *hold_pbuf;
    uint16_t hold_offset;
    uint8_t index;
    tcp_server_instance_config_t config;
    tcp_server_status_t status;
} tcp_server_ctx_t;

static tcp_server_ctx_t g_servers[TCP_SERVER_INSTANCE_COUNT];

static uint16_t ring_free(uint16_t head, uint16_t tail, uint16_t size)
{
    return (head >= tail) ? (uint16_t)(size - head + tail - 1u) : (uint16_t)(tail - head - 1u);
}

static uint16_t ring_used(uint16_t head, uint16_t tail, uint16_t size)
{
    return (head >= tail) ? (uint16_t)(head - tail) : (uint16_t)(size - tail + head);
}

static void tcp_server_reset_rx_state(tcp_server_ctx_t *ctx)
{
    if (ctx == NULL) {
        return;
    }
    if (ctx->hold_pbuf != NULL) {
        pbuf_free(ctx->hold_pbuf);
        ctx->hold_pbuf = NULL;
    }
    ctx->hold_offset = 0u;
    ctx->rx_head = 0u;
    ctx->rx_tail = 0u;
}

static void tcp_server_fill_ring_from_pbuf(tcp_server_ctx_t *ctx)
{
    struct pbuf *q;
    uint16_t offset;

    if (ctx == NULL || ctx->hold_pbuf == NULL) {
        return;
    }

    q = ctx->hold_pbuf;
    offset = ctx->hold_offset;
    while (q != NULL && offset >= q->len) {
        offset = (uint16_t)(offset - q->len);
        q = q->next;
    }

    while (q != NULL) {
        const uint8_t *src = (const uint8_t *)q->payload;
        for (uint16_t i = offset; i < q->len; ++i) {
            if (ring_free(ctx->rx_head, ctx->rx_tail, TCP_SERVER_RX_BUFFER_SIZE) == 0u) {
                ctx->hold_offset = (uint16_t)(ctx->hold_offset + i - offset);
                return;
            }
            ctx->rx_ring[ctx->rx_head] = src[i];
            ctx->rx_head = (uint16_t)((ctx->rx_head + 1u) % TCP_SERVER_RX_BUFFER_SIZE);
            ctx->status.rx_bytes++;
        }
        ctx->hold_offset = (uint16_t)(ctx->hold_offset + q->len - offset);
        offset = 0u;
        q = q->next;
    }

    pbuf_free(ctx->hold_pbuf);
    ctx->hold_pbuf = NULL;
    ctx->hold_offset = 0u;
}

static err_t tcp_server_on_recv(void *arg, struct tcp_pcb *pcb, struct pbuf *p, err_t err)
{
    tcp_server_ctx_t *ctx = (tcp_server_ctx_t *)arg;

    if (ctx == NULL) {
        if (p != NULL) {
            pbuf_free(p);
        }
        return ERR_ARG;
    }
    if (err != ERR_OK) {
        if (p != NULL) {
            pbuf_free(p);
        }
        return err;
    }
    if (p == NULL) {
        tcp_arg(pcb, NULL);
        tcp_recv(pcb, NULL);
        tcp_sent(pcb, NULL);
        tcp_err(pcb, NULL);
        tcp_abort(pcb);
        ctx->client_pcb = NULL;
        ctx->status.state = ctx->config.enabled ? TCP_SERVER_STATE_LISTENING : TCP_SERVER_STATE_IDLE;
        return ERR_ABRT;
    }

    if (ctx->hold_pbuf != NULL) {
        ctx->status.errors++;
        return ERR_MEM;
    }

    pbuf_ref(p);
    ctx->hold_pbuf = p;
    ctx->hold_offset = 0u;
    pbuf_free(p);
    tcp_server_fill_ring_from_pbuf(ctx);
    return ERR_OK;
}

static err_t tcp_server_on_sent(void *arg, struct tcp_pcb *pcb, u16_t len)
{
    tcp_server_ctx_t *ctx = (tcp_server_ctx_t *)arg;
    (void)pcb;
    if (ctx != NULL) {
        ctx->status.tx_bytes += len;
    }
    return ERR_OK;
}

static void tcp_server_on_err(void *arg, err_t err)
{
    tcp_server_ctx_t *ctx = (tcp_server_ctx_t *)arg;
    if (ctx == NULL) {
        return;
    }
    tcp_server_reset_rx_state(ctx);
    ctx->client_pcb = NULL;
    ctx->status.state = ctx->config.enabled ? TCP_SERVER_STATE_LISTENING : TCP_SERVER_STATE_IDLE;
    ctx->status.errors++;
    SEGGER_RTT_printf(0, "TCP server[%u] connection error=%d\r\n", ctx->index, (int)err);
}

static err_t tcp_server_on_accept(void *arg, struct tcp_pcb *newpcb, err_t err)
{
    tcp_server_ctx_t *ctx = (tcp_server_ctx_t *)arg;

    if (ctx == NULL || err != ERR_OK) {
        return (ctx == NULL) ? ERR_ARG : err;
    }
    if (ctx->client_pcb != NULL) {
        tcp_abort(newpcb);
        return ERR_ABRT;
    }

    ctx->client_pcb = newpcb;
    ctx->status.state = TCP_SERVER_STATE_CONNECTED;
    ctx->status.connections++;
    tcp_nagle_disable(newpcb);
    tcp_arg(newpcb, ctx);
    tcp_recv(newpcb, tcp_server_on_recv);
    tcp_sent(newpcb, tcp_server_on_sent);
    tcp_err(newpcb, tcp_server_on_err);
    return ERR_OK;
}

int tcp_server_init_all(void)
{
    memset(g_servers, 0, sizeof(g_servers));
    for (uint8_t i = 0; i < TCP_SERVER_INSTANCE_COUNT; ++i) {
        g_servers[i].index = i;
        g_servers[i].status.state = TCP_SERVER_STATE_IDLE;
    }
    return 0;
}

int tcp_server_config(uint8_t instance, const tcp_server_instance_config_t *config)
{
    if (instance >= TCP_SERVER_INSTANCE_COUNT || config == NULL) {
        return -1;
    }
    g_servers[instance].config = *config;
    return 0;
}

int tcp_server_start(uint8_t instance)
{
    struct tcp_pcb *pcb;
    err_t err;
    tcp_server_ctx_t *ctx;

    if (instance >= TCP_SERVER_INSTANCE_COUNT) {
        return -1;
    }
    ctx = &g_servers[instance];
    if (!ctx->config.enabled) {
        ctx->status.state = TCP_SERVER_STATE_IDLE;
        return 0;
    }
    if (ctx->listen_pcb != NULL) {
        return 0;
    }

    pcb = tcp_new_ip_type(IPADDR_TYPE_V4);
    if (pcb == NULL) {
        ctx->status.errors++;
        return -1;
    }

    err = tcp_bind(pcb, IP_ANY_TYPE, ctx->config.port);
    if (err != ERR_OK) {
        tcp_abort(pcb);
        ctx->status.errors++;
        return -1;
    }

    ctx->listen_pcb = tcp_listen_with_backlog(pcb, 1);
    if (ctx->listen_pcb == NULL) {
        ctx->status.errors++;
        return -1;
    }

    tcp_arg(ctx->listen_pcb, ctx);
    tcp_accept(ctx->listen_pcb, tcp_server_on_accept);
    ctx->status.state = TCP_SERVER_STATE_LISTENING;
    return 0;
}

int tcp_server_stop(uint8_t instance)
{
    tcp_server_ctx_t *ctx;

    if (instance >= TCP_SERVER_INSTANCE_COUNT) {
        return -1;
    }
    ctx = &g_servers[instance];

    if (ctx->client_pcb != NULL) {
        tcp_server_reset_rx_state(ctx);
        tcp_arg(ctx->client_pcb, NULL);
        tcp_recv(ctx->client_pcb, NULL);
        tcp_sent(ctx->client_pcb, NULL);
        tcp_err(ctx->client_pcb, NULL);
        tcp_abort(ctx->client_pcb);
        ctx->client_pcb = NULL;
    }
    if (ctx->listen_pcb != NULL) {
        tcp_arg(ctx->listen_pcb, NULL);
        tcp_accept(ctx->listen_pcb, NULL);
        if (tcp_close(ctx->listen_pcb) != ERR_OK) {
            tcp_abort(ctx->listen_pcb);
        }
        ctx->listen_pcb = NULL;
    }

    ctx->status.state = TCP_SERVER_STATE_IDLE;
    tcp_server_reset_rx_state(ctx);
    return 0;
}

int tcp_server_send(uint8_t instance, const uint8_t *data, uint16_t len)
{
    err_t err;
    tcp_server_ctx_t *ctx;

    if (instance >= TCP_SERVER_INSTANCE_COUNT || data == NULL || len == 0u) {
        return -1;
    }
    ctx = &g_servers[instance];
    if (ctx->client_pcb == NULL) {
        return -1;
    }
    if (tcp_sndbuf(ctx->client_pcb) < len) {
        ctx->status.errors++;
        return 0;
    }

    err = tcp_write(ctx->client_pcb, data, len, TCP_WRITE_FLAG_COPY);
    if (err == ERR_MEM) {
        ctx->status.errors++;
        return 0;
    }
    if (err != ERR_OK) {
        ctx->status.errors++;
        return -1;
    }
    err = tcp_output(ctx->client_pcb);
    if (err == ERR_MEM) {
        ctx->status.errors++;
        return 0;
    }
    if (err != ERR_OK) {
        ctx->status.errors++;
        return -1;
    }
    return (int)len;
}

int tcp_server_recv(uint8_t instance, uint8_t *data, uint16_t max_len)
{
    uint16_t copied = 0u;
    tcp_server_ctx_t *ctx;

    if (instance >= TCP_SERVER_INSTANCE_COUNT || data == NULL || max_len == 0u) {
        return -1;
    }
    ctx = &g_servers[instance];
    tcp_server_fill_ring_from_pbuf(ctx);
    while (copied < max_len && ctx->rx_tail != ctx->rx_head) {
        data[copied++] = ctx->rx_ring[ctx->rx_tail];
        ctx->rx_tail = (uint16_t)((ctx->rx_tail + 1u) % TCP_SERVER_RX_BUFFER_SIZE);
    }
    if (copied > 0u && ctx->client_pcb != NULL) {
        tcp_recved(ctx->client_pcb, copied);
    }
    return (int)copied;
}

uint16_t tcp_server_rx_available(uint8_t instance)
{
    if (instance >= TCP_SERVER_INSTANCE_COUNT) {
        return 0u;
    }
    tcp_server_fill_ring_from_pbuf(&g_servers[instance]);
    return ring_used(g_servers[instance].rx_head, g_servers[instance].rx_tail, TCP_SERVER_RX_BUFFER_SIZE);
}

uint16_t tcp_server_peek(uint8_t instance, uint8_t *data, uint16_t max_len)
{
    uint16_t copied = 0u;
    uint16_t tail;
    tcp_server_ctx_t *ctx;

    if (instance >= TCP_SERVER_INSTANCE_COUNT || data == NULL || max_len == 0u) {
        return 0u;
    }

    ctx = &g_servers[instance];
    tcp_server_fill_ring_from_pbuf(ctx);
    tail = ctx->rx_tail;
    while (copied < max_len && tail != ctx->rx_head) {
        data[copied++] = ctx->rx_ring[tail];
        tail = (uint16_t)((tail + 1u) % TCP_SERVER_RX_BUFFER_SIZE);
    }
    return copied;
}

void tcp_server_drop(uint8_t instance, uint16_t len)
{
    tcp_server_ctx_t *ctx;
    uint16_t dropped = 0u;

    if (instance >= TCP_SERVER_INSTANCE_COUNT || len == 0u) {
        return;
    }

    ctx = &g_servers[instance];
    while (dropped < len && ctx->rx_tail != ctx->rx_head) {
        ctx->rx_tail = (uint16_t)((ctx->rx_tail + 1u) % TCP_SERVER_RX_BUFFER_SIZE);
        dropped++;
    }
    if (dropped > 0u && ctx->client_pcb != NULL) {
        tcp_recved(ctx->client_pcb, dropped);
    }
    tcp_server_fill_ring_from_pbuf(ctx);
}

bool tcp_server_is_connected(uint8_t instance)
{
    return (instance < TCP_SERVER_INSTANCE_COUNT) && (g_servers[instance].client_pcb != NULL);
}

void tcp_server_get_status(uint8_t instance, tcp_server_status_t *status)
{
    if (instance < TCP_SERVER_INSTANCE_COUNT && status != NULL) {
        *status = g_servers[instance].status;
    }
}

void tcp_server_poll(void)
{
    for (uint8_t i = 0; i < TCP_SERVER_INSTANCE_COUNT; ++i) {
        tcp_server_fill_ring_from_pbuf(&g_servers[i]);
    }
}