/**
 * @file tcp_client.c
 * @brief lwIP RAW TCP client for the UART3 bridge.
 */

#include "tcp_client.h"

#include "main.h"

#include "SEGGER_RTT.h"

#include "lwip/ip_addr.h"
#include "lwip/pbuf.h"
#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;
    uint32_t next_retry_ms;
    tcp_client_config_t config;
    tcp_client_status_t status;
} tcp_client_ctx_t;

static tcp_client_ctx_t g_client;

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 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;
    struct pbuf *q;

    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);
        if (tcp_close(pcb) != ERR_OK) {
            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_OK;
    }

    for (q = p; q != NULL; q = q->next) {
        const uint8_t *src = (const uint8_t *)q->payload;
        for (uint16_t i = 0; i < q->len; ++i) {
            if (ring_free(ctx->rx_head, ctx->rx_tail, TCP_CLIENT_RX_BUFFER_SIZE) == 0u) {
                ctx->status.errors++;
                break;
            }
            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++;
        }
    }

    tcp_recved(pcb, p->tot_len);
    pbuf_free(p);
    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;
    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;
    }

    ctx->pcb = NULL;
    ctx->status.state = TCP_CLIENT_STATE_DISCONNECTED;
    ctx->status.errors++;
    ctx->next_retry_ms = HAL_GetTick() + ctx->config.reconnect_interval_ms;
    SEGGER_RTT_printf(0, "TCP client error=%d, reconnect scheduled\r\n", (int)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 (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;
        SEGGER_RTT_printf(0, "TCP client connect callback failed err=%d\r\n", (int)err);
        return err;
    }

    ctx->pcb = pcb;
    ctx->status.state = TCP_CLIENT_STATE_CONNECTED;
    SEGGER_RTT_printf(0,
                      "TCP client connected to %u.%u.%u.%u:%u\r\n",
                      ctx->config.server_ip[0], ctx->config.server_ip[1],
                      ctx->config.server_ip[2], ctx->config.server_ip[3],
                      ctx->config.server_port);
    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(const tcp_client_config_t *config)
{
    memset(&g_client, 0, sizeof(g_client));
    g_client.config.server_ip[0] = 192u;
    g_client.config.server_ip[1] = 168u;
    g_client.config.server_ip[2] = 31u;
    g_client.config.server_ip[3] = 1u;
    g_client.config.local_port = TCP_CLIENT_DEFAULT_PORT;
    g_client.config.server_port = TCP_CLIENT_DEFAULT_PORT;
    g_client.config.auto_reconnect = true;
    g_client.config.reconnect_interval_ms = TCP_CLIENT_RECONNECT_DELAY_MS;
    g_client.status.state = TCP_CLIENT_STATE_IDLE;

    if (config != NULL) {
        g_client.config = *config;
    }

    return 0;
}

int tcp_client_connect(void)
{
    struct tcp_pcb *pcb;
    ip_addr_t remote_addr;
    err_t err;

    if (g_client.pcb != NULL) {
        return 0;
    }

    pcb = tcp_new_ip_type(IPADDR_TYPE_V4);
    if (pcb == NULL) {
        g_client.status.errors++;
        g_client.status.state = TCP_CLIENT_STATE_ERROR;
        SEGGER_RTT_WriteString(0, "TCP client connect failed: no PCB\r\n");
        return -1;
    }

    if (g_client.config.local_port != 0u) {
        err = tcp_bind(pcb, IP_ANY_TYPE, g_client.config.local_port);
        if (err != ERR_OK) {
            tcp_abort(pcb);
            g_client.status.state = TCP_CLIENT_STATE_DISCONNECTED;
            g_client.status.errors++;
            g_client.next_retry_ms = HAL_GetTick() + g_client.config.reconnect_interval_ms;
            SEGGER_RTT_printf(0, "TCP client bind failed err=%d\r\n", (int)err);
            return -1;
        }
    }

    IP_ADDR4(&remote_addr,
             g_client.config.server_ip[0],
             g_client.config.server_ip[1],
             g_client.config.server_ip[2],
             g_client.config.server_ip[3]);

    g_client.status.state = TCP_CLIENT_STATE_CONNECTING;
    tcp_arg(pcb, &g_client);
    tcp_err(pcb, tcp_client_on_err);
    err = tcp_connect(pcb, &remote_addr, g_client.config.server_port, tcp_client_on_connected);
    if (err != ERR_OK) {
        tcp_err(pcb, NULL);
        tcp_abort(pcb);
        g_client.status.state = TCP_CLIENT_STATE_DISCONNECTED;
        g_client.status.errors++;
        g_client.next_retry_ms = HAL_GetTick() + g_client.config.reconnect_interval_ms;
        SEGGER_RTT_printf(0, "TCP client connect start failed err=%d\r\n", (int)err);
        return -1;
    }

    g_client.pcb = pcb;
    SEGGER_RTT_printf(0,
                      "TCP client connecting to %u.%u.%u.%u:%u\r\n",
                      g_client.config.server_ip[0], g_client.config.server_ip[1],
                      g_client.config.server_ip[2], g_client.config.server_ip[3],
                      g_client.config.server_port);
    return 0;
}

int tcp_client_disconnect(void)
{
    if (g_client.pcb != NULL) {
        tcp_arg(g_client.pcb, NULL);
        tcp_recv(g_client.pcb, NULL);
        tcp_sent(g_client.pcb, NULL);
        tcp_err(g_client.pcb, NULL);
        tcp_abort(g_client.pcb);
        g_client.pcb = NULL;
    }

    g_client.status.state = TCP_CLIENT_STATE_DISCONNECTED;
    SEGGER_RTT_WriteString(0, "TCP client disconnected\r\n");
    return 0;
}

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

    if (g_client.pcb == NULL || data == NULL || len == 0u) {
        return -1;
    }

    if ((g_client.pcb->flags & TF_RXCLOSED) != 0u) {
        g_client.status.errors++;
        return -1;
    }

    if (tcp_sndbuf(g_client.pcb) < len) {
        return 0;
    }

    err = tcp_write(g_client.pcb, data, len, TCP_WRITE_FLAG_COPY);
    if (err != ERR_OK) {
        g_client.status.errors++;
        return -1;
    }

    err = tcp_output(g_client.pcb);
    if (err != ERR_OK) {
        g_client.status.errors++;
        return -1;
    }

    return (int)len;
}

int tcp_client_recv(uint8_t *data, uint16_t max_len, uint32_t timeout_ms)
{
    uint16_t copied = 0u;
    (void)timeout_ms;

    if (data == NULL || max_len == 0u) {
        return -1;
    }

    while (copied < max_len && g_client.rx_tail != g_client.rx_head) {
        data[copied++] = g_client.rx_ring[g_client.rx_tail];
        g_client.rx_tail = (uint16_t)((g_client.rx_tail + 1u) % TCP_CLIENT_RX_BUFFER_SIZE);
    }

    return (int)copied;
}

bool tcp_client_is_connected(void)
{
    return (g_client.pcb != NULL) && (g_client.status.state == TCP_CLIENT_STATE_CONNECTED);
}

int tcp_client_set_server(const uint8_t *ip, uint16_t port)
{
    if (ip == NULL || port == 0u) {
        return -1;
    }

    memcpy(g_client.config.server_ip, ip, 4u);
    g_client.config.server_port = port;
    return 0;
}

void tcp_client_get_status(tcp_client_status_t *status)
{
    if (status != NULL) {
        *status = g_client.status;
    }
}

void *tcp_client_get_rx_stream(void)
{
    return NULL;
}

void *tcp_client_get_tx_stream(void)
{
    return NULL;
}

void tcp_client_task(void *argument)
{
    (void)argument;
}

void tcp_client_poll(void)
{
    uint32_t now;

    if (!g_client.config.auto_reconnect || tcp_client_is_connected()) {
        return;
    }

    if ((g_client.pcb != NULL) && (g_client.status.state == TCP_CLIENT_STATE_CONNECTING)) {
        return;
    }

    now = HAL_GetTick();
    if (now >= g_client.next_retry_ms) {
        g_client.status.reconnect_count++;
        g_client.next_retry_ms = now + g_client.config.reconnect_interval_ms;
        SEGGER_RTT_printf(0,
                          "TCP client reconnect attempt %lu\r\n",
                          g_client.status.reconnect_count);
        (void)tcp_client_connect();
    }
}