/**
 * @file tcp_client.c
 * @brief Indexed lwIP RAW TCP client manager.
 */

#include "tcp_client.h"

#include "../Core/Inc/main.h"
#include "../Drivers/LwIP/src/include/lwip/ip_addr.h"
#include "../Drivers/LwIP/src/include/lwip/pbuf.h"
#include "../Drivers/LwIP/src/include/lwip/tcp.h"

#include <string.h>

typedef struct {
    struct tcp_pcb *pcb;
    uint8_t rx_ring[TCP_CLIENT_RX_BUFFER_SIZE];
    uint16_t rx_head;
    uint16_t rx_tail;
    struct pbuf *hold_pbuf;
    uint16_t hold_offset;
    uint32_t next_retry_ms;
    uint8_t index;
    tcp_client_instance_config_t config;
    tcp_client_status_t status;
} tcp_client_ctx_t;

static tcp_client_ctx_t g_clients[TCP_CLIENT_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_client_reset_rx_state(tcp_client_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_client_fill_ring_from_pbuf(tcp_client_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_CLIENT_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_CLIENT_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_client_on_recv(void *arg, struct tcp_pcb *pcb, struct pbuf *p, err_t err)
{
    tcp_client_ctx_t *ctx = (tcp_client_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->pcb = NULL;
        ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
        ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
        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_client_fill_ring_from_pbuf(ctx);
    return ERR_OK;
}

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

static void tcp_client_on_err(void *arg, err_t err)
{
    tcp_client_ctx_t *ctx = (tcp_client_ctx_t *)arg;
    if (ctx == NULL) {
        return;
    }
    tcp_client_reset_rx_state(ctx);
    ctx->pcb = NULL;
    ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
    ctx->status.errors++;
    ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
    (void)err;
}

static err_t tcp_client_on_connected(void *arg, struct tcp_pcb *pcb, err_t err)
{
    tcp_client_ctx_t *ctx = (tcp_client_ctx_t *)arg;

    if (ctx == NULL) {
        return ERR_ARG;
    }
    if (err != ERR_OK) {
        ctx->pcb = NULL;
        ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
        ctx->status.errors++;
        ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
        return err;
    }

    ctx->pcb = pcb;
    ctx->status.state = TCP_CLIENT_STATE_CONNECTED;
    tcp_nagle_disable(pcb);
    tcp_arg(pcb, ctx);
    tcp_recv(pcb, tcp_client_on_recv);
    tcp_sent(pcb, tcp_client_on_sent);
    tcp_err(pcb, tcp_client_on_err);
    return ERR_OK;
}

int tcp_client_init_all(void)
{
    memset(g_clients, 0, sizeof(g_clients));
    for (uint8_t i = 0; i < TCP_CLIENT_INSTANCE_COUNT; ++i) {
        g_clients[i].index = i;
        g_clients[i].status.state = TCP_CLIENT_STATE_IDLE;
        g_clients[i].config.reconnect_interval_ms = TCP_CLIENT_RECONNECT_DELAY_MS;
        g_clients[i].config.auto_reconnect = true;
    }
    return 0;
}

int tcp_client_config(uint8_t instance, const tcp_client_instance_config_t *config)
{
    if (instance >= TCP_CLIENT_INSTANCE_COUNT || config == NULL) {
        return -1;
    }
    g_clients[instance].config = *config;
    return 0;
}

int tcp_client_connect(uint8_t instance)
{
    struct tcp_pcb *pcb;
    ip_addr_t remote_addr;
    err_t err;
    tcp_client_ctx_t *ctx;

    if (instance >= TCP_CLIENT_INSTANCE_COUNT) {
        return -1;
    }
    ctx = &g_clients[instance];
    if (!ctx->config.enabled) {
        return 0;
    }
    if (ctx->pcb != NULL) {
        return 0;
    }

    pcb = tcp_new_ip_type(IPADDR_TYPE_V4);
    if (pcb == NULL) {
        ctx->status.errors++;
        ctx->status.state = TCP_CLIENT_STATE_ERROR;
        return -1;
    }
    if (ctx->config.local_port != 0u) {
        err = tcp_bind(pcb, IP_ANY_TYPE, ctx->config.local_port);
        if (err != ERR_OK) {
            tcp_abort(pcb);
            ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
            ctx->status.errors++;
            ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
            return -1;
        }
    }

    IP_ADDR4(&remote_addr,
             ctx->config.remote_ip[0],
             ctx->config.remote_ip[1],
             ctx->config.remote_ip[2],
             ctx->config.remote_ip[3]);

    ctx->status.state = TCP_CLIENT_STATE_CONNECTING;
    tcp_arg(pcb, ctx);
    tcp_err(pcb, tcp_client_on_err);
    err = tcp_connect(pcb, &remote_addr, ctx->config.remote_port, tcp_client_on_connected);
    if (err != ERR_OK) {
        tcp_err(pcb, NULL);
        tcp_abort(pcb);
        ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
        ctx->status.errors++;
        ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
        return -1;
    }

    ctx->pcb = pcb;
    return 0;
}

int tcp_client_disconnect(uint8_t instance)
{
    tcp_client_ctx_t *ctx;

    if (instance >= TCP_CLIENT_INSTANCE_COUNT) {
        return -1;
    }
    ctx = &g_clients[instance];
    if (ctx->pcb != NULL) {
        tcp_client_reset_rx_state(ctx);
        tcp_arg(ctx->pcb, NULL);
        tcp_recv(ctx->pcb, NULL);
        tcp_sent(ctx->pcb, NULL);
        tcp_err(ctx->pcb, NULL);
        tcp_abort(ctx->pcb);
        ctx->pcb = NULL;
    }
    ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
    tcp_client_reset_rx_state(ctx);
    return 0;
}

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

    if (instance >= TCP_CLIENT_INSTANCE_COUNT || data == NULL || len == 0u) {
        return -1;
    }
    ctx = &g_clients[instance];
    if (ctx->pcb == NULL) {
        return -1;
    }
    if (tcp_sndbuf(ctx->pcb) < len) {
        ctx->status.errors++;
        return 0;
    }
    err = tcp_write(ctx->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->pcb);
    if (err == ERR_MEM) {
        ctx->status.errors++;
        return 0;
    }
    if (err != ERR_OK) {
        ctx->status.errors++;
        return -1;
    }
    return (int)len;
}

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

    if (instance >= TCP_CLIENT_INSTANCE_COUNT || data == NULL || max_len == 0u) {
        return -1;
    }
    ctx = &g_clients[instance];
    tcp_client_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_CLIENT_RX_BUFFER_SIZE);
    }
    if (copied > 0u && ctx->pcb != NULL) {
        tcp_recved(ctx->pcb, copied);
    }
    return (int)copied;
}

uint16_t tcp_client_rx_available(uint8_t instance)
{
    if (instance >= TCP_CLIENT_INSTANCE_COUNT) {
        return 0u;
    }
    tcp_client_fill_ring_from_pbuf(&g_clients[instance]);
    return ring_used(g_clients[instance].rx_head, g_clients[instance].rx_tail, TCP_CLIENT_RX_BUFFER_SIZE);
}

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

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

    ctx = &g_clients[instance];
    tcp_client_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_CLIENT_RX_BUFFER_SIZE);
    }
    return copied;
}

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

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

    ctx = &g_clients[instance];
    while (dropped < len && ctx->rx_tail != ctx->rx_head) {
        ctx->rx_tail = (uint16_t)((ctx->rx_tail + 1u) % TCP_CLIENT_RX_BUFFER_SIZE);
        dropped++;
    }
    if (dropped > 0u && ctx->pcb != NULL) {
        tcp_recved(ctx->pcb, dropped);
    }
    tcp_client_fill_ring_from_pbuf(ctx);
}

bool tcp_client_is_connected(uint8_t instance)
{
    return (instance < TCP_CLIENT_INSTANCE_COUNT) &&
           (g_clients[instance].pcb != NULL) &&
           (g_clients[instance].status.state == TCP_CLIENT_STATE_CONNECTED);
}

void tcp_client_get_status(uint8_t instance, tcp_client_status_t *status)
{
    if (instance < TCP_CLIENT_INSTANCE_COUNT && status != NULL) {
        *status = g_clients[instance].status;
    }
}

void tcp_client_poll(void)
{
    uint32_t now = HAL_GetTick();

    for (uint8_t i = 0; i < TCP_CLIENT_INSTANCE_COUNT; ++i) {
        tcp_client_ctx_t *ctx = &g_clients[i];
        tcp_client_fill_ring_from_pbuf(ctx);
        if (!ctx->config.enabled || !ctx->config.auto_reconnect || tcp_client_is_connected(i)) {
            continue;
        }
        if ((ctx->pcb != NULL) && (ctx->status.state == TCP_CLIENT_STATE_CONNECTING)) {
            continue;
        }
        if (now >= ctx->next_retry_ms) {
            ctx->status.reconnect_count++;
            ctx->next_retry_ms = now + ctx->config.reconnect_interval_ms;
            (void)tcp_client_connect(i);
        }
    }
}