/**
 * @file        tcp_echo_server.c
 *
 * @brief       TCP echo server example using raw API
 *
 * @version     V1.0.0
 *
 * @date        2025-05-08
 *
 * @attention
 *
 *  Copyright (C) 2025 Geehy Semiconductor
 *
 *  You may not use this file except in compliance with the
 *  GEEHY COPYRIGHT NOTICE (GEEHY SOFTWARE PACKAGE LICENSE).
 *
 *  The program is only for reference, which is distributed in the hope
 *  that it will be useful and instructional for customers to develop
 *  their software. Unless required by applicable law or agreed to in
 *  writing, the program is distributed on an "AS IS" BASIS, WITHOUT
 *  ANY WARRANTY OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the GEEHY SOFTWARE PACKAGE LICENSE for the governing permissions
 *  and limitations under the License.
 */

/* This file was modified by Geehy Semiconductor */

/* Includes ***************************************************************/
#include "lwip/opt.h"
#include "lwip/debug.h"
#include "lwip/stats.h"
#include "lwip/tcp.h"
#include "tcp_echo_server.h"
#include "log.h"
#include <stdlib.h>
#include <string.h>

#if LWIP_TCP

/* Private macro **********************************************************/

/* Private typedef ********************************************************/
enum tcps_echo_states
{
    ES_NONE = 0,
    ES_ACCEPTED,
    ES_RECEIVED,
    ES_CLOSING
};

struct tcps_echo_state
{
    u8_t state;
    u8_t retries;
    struct tcp_pcb *pcb;
    /* pbuf (chain) to recycle */
    struct pbuf *p;
};

/* Private variables ******************************************************/
static const char* tag = "tcp_echo_server";
struct tcp_pcb *tcps_echo_pcb;
struct tcps_echo_state * tcps_echo_es;

/* Private function prototypes ********************************************/
err_t tcps_echo_accept(void *arg, struct tcp_pcb *newpcb, err_t err);
err_t tcps_echo_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err);
void tcps_echo_error(void *arg, err_t err);
err_t tcps_echo_poll(void *arg, struct tcp_pcb *tpcb);
err_t tcps_echo_sent(void *arg, struct tcp_pcb *tpcb, u16_t len);
void tcps_echo_send(struct tcp_pcb *tpcb, struct tcps_echo_state *es);
void tcps_echo_close(struct tcp_pcb *tpcb, struct tcps_echo_state *es);
void tcp_printf_rec_data(struct pbuf *p);

/* External variables *****************************************************/

/* External functions *****************************************************/

/**
 * @brief   tcp server init
 *
 * @param   tcps_port: server port
 *
 * @retval  None
 */
void tcps_echo_init(uint16_t tcps_port)
{
    tcps_echo_pcb = tcp_new();
    if (tcps_echo_pcb != NULL)
    {
        err_t err;

        err = tcp_bind(tcps_echo_pcb, IP_ADDR_ANY, tcps_port);
        if (err == ERR_OK)
        {
            tcps_echo_pcb = tcp_listen(tcps_echo_pcb);
            tcp_accept(tcps_echo_pcb, tcps_echo_accept);
        }
        else
        {
            /* abort? output diagnostic? */
        }
    }
    else
    {
        /* abort? output diagnostic? */
    }
}

/**
 * @brief   Printf TCP Receive Data
 *
 * @param   p: pointer on the tcp pcb received pbuf
 *
 * @retval  None
 */
void tcp_printf_rec_data(struct pbuf *p)
{
    char recbuf[256] = {0};

    memcpy(recbuf,p->payload,p->len);
    LOGI(tag, "Received message:%s\r\n",recbuf);
}

/**
 * @brief   LwIP tcp_accept callback function
 *
 * @param   newpcb: point to the newly created tcp pcb
 *
 * @retval  None
 */
err_t tcps_echo_accept(void *arg, struct tcp_pcb *newpcb, err_t err)
{
    err_t ret_err;
    struct tcps_echo_state *es;

    LWIP_UNUSED_ARG(arg);
    LWIP_UNUSED_ARG(err);

    /* Unless this pcb should have NORMAL priority, set its priority now.
        When running out of pcbs, low priority pcbs can be aborted to create
        new pcbs of higher priority. */
    tcp_setprio(newpcb, TCP_PRIO_MIN);

    es = (struct tcps_echo_state *)mem_malloc(sizeof(struct tcps_echo_state));
    tcps_echo_es = es;

    if (es != NULL)
    {
        es->state = ES_ACCEPTED;
        es->pcb = newpcb;
        es->retries = 0;
        es->p = NULL;

        /* pass newly allocated es to our callbacks */
        tcp_arg(newpcb, es);
        tcp_recv(newpcb, tcps_echo_recv);
        tcp_err(newpcb, tcps_echo_error);
        tcp_poll(newpcb, tcps_echo_poll, 0);

        ret_err = ERR_OK;
    }
    else
    {
        /* close tcp connection */
        tcps_echo_close(newpcb, es);
        ret_err = ERR_MEM;
    }

    return ret_err;
}

/**
 * @brief   LwIP tcp_recv callback function
 *
 * @param   None
 *
 * @retval  None
 */
err_t tcps_echo_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err)
{
    struct tcps_echo_state *es;
    err_t ret_err;

    LWIP_ASSERT("arg != NULL",arg != NULL);
    es = (struct tcps_echo_state *)arg;
    if (p == NULL)
    {
        /* remote host closed connection */
        es->state = ES_CLOSING;
        if(es->p == NULL)
        {
            /* we're done sending, close it */
            tcps_echo_close(tpcb, es);
        }
        else
        {
            /* we're not done yet */
            tcp_sent(tpcb, tcps_echo_sent);
            tcps_echo_send(tpcb, es);
        }
        ret_err = ERR_OK;
    }
    else if(err != ERR_OK)
    {
        /* cleanup, for unkown reason */
        if (p != NULL)
        {
            es->p = NULL;
            pbuf_free(p);
        }
        ret_err = err;
    }
    else if(es->state == ES_ACCEPTED)
    {
        /* first data chunk in p->payload */
        es->state = ES_RECEIVED;
        /* store reference to incoming pbuf (chain) */
        es->p = p;
        /* install send completion notifier */
        tcp_sent(tpcb, tcps_echo_sent);

        /** Printf TCP Receive Data */
        tcp_printf_rec_data(p);

        tcps_echo_send(tpcb, es);
        ret_err = ERR_OK;
    }
    else if (es->state == ES_RECEIVED)
    {
        /* read some more data */
        if(es->p == NULL)
        {
            es->p = p;

            /** Printf TCP Receive Data */
            tcp_printf_rec_data(p);

            tcp_sent(tpcb, tcps_echo_sent);
            tcps_echo_send(tpcb, es);
        }
        else
        {
            struct pbuf *ptr;

            /* chain pbufs to the end of what we recv'ed previously  */
            ptr = es->p;
            pbuf_chain(ptr,p);
        }
        ret_err = ERR_OK;
    }
    else if(es->state == ES_CLOSING)
    {
        /* odd case, remote side closing twice, trash data */
        tcp_recved(tpcb, p->tot_len);
        es->p = NULL;
        pbuf_free(p);
        ret_err = ERR_OK;
    }
    else
    {
        /* unkown es->state, trash data  */
        tcp_recved(tpcb, p->tot_len);
        es->p = NULL;
        pbuf_free(p);
        ret_err = ERR_OK;
    }

    return ret_err;
}

/**
 * @brief   LwIP tcp_err callback function
 *
 * @param   None
 *
 * @retval  None
 */
void tcps_echo_error(void *arg, err_t err)
{
    struct tcps_echo_state *es;

    LWIP_UNUSED_ARG(err);

    es = (struct tcps_echo_state *)arg;
    if (es != NULL)
    {
        mem_free(es);
    }
}

/**
 * @brief   LwIP tcp_poll callback function
 *
 * @param   None
 *
 * @retval  None
 */
err_t tcps_echo_poll(void *arg, struct tcp_pcb *tpcb)
{
    err_t ret_err;
    struct tcps_echo_state *es;

    es = (struct tcps_echo_state *)arg;
    if (es != NULL)
    {
        if (es->p != NULL)
        {
            /* there is a remaining pbuf (chain)  */
            tcp_sent(tpcb, tcps_echo_sent);
            tcps_echo_send(tpcb, es);
        }
        else
        {
            /* no remaining pbuf (chain)  */
            if(es->state == ES_CLOSING)
            {
                tcps_echo_close(tpcb, es);
            }
        }
        ret_err = ERR_OK;
    }
    else
    {
        /* nothing to be done */
        tcp_abort(tpcb);
        ret_err = ERR_ABRT;
    }
    return ret_err;
}

/**
 * @brief   LwIP tcp_sent callback function
 *
 * @param   None
 *
 * @retval  None
 */
err_t tcps_echo_sent(void *arg, struct tcp_pcb *tpcb, u16_t len)
{
    struct tcps_echo_state *es;

    LWIP_UNUSED_ARG(len);

    es = (struct tcps_echo_state *)arg;
    es->retries = 0;

    if(es->p != NULL)
    {
        /* still got pbufs to send */
        tcp_sent(tpcb, tcps_echo_sent);
        tcps_echo_send(tpcb, es);
    }
    else
    {
        /* no more pbufs to send */
        if(es->state == ES_CLOSING)
        {
            tcps_echo_close(tpcb, es);
        }
    }

    return ERR_OK;
}

/**
 * @brief   tcp send data function
 *
 * @param   tpcb: point to the tcp pcb
 *
 * @param   es: point to tcps_echo_state structure
 *
 * @retval  None
 */
void tcps_echo_send(struct tcp_pcb *tpcb, struct tcps_echo_state *es)
{
    struct pbuf *ptr;
    err_t wr_err = ERR_OK;

    while ((wr_err == ERR_OK) &&
            (es->p != NULL) &&
            (es->p->len <= tcp_sndbuf(tpcb)))
    {
        ptr = es->p;

        /* enqueue data for transmission */
        wr_err = tcp_write(tpcb, ptr->payload, ptr->len, 1);
        if (wr_err == ERR_OK)
        {
            u16_t plen;
            u8_t freed;

            plen = ptr->len;
            /* continue with next pbuf in chain (if any) */
            es->p = ptr->next;
            if(es->p != NULL)
            {
                /* new reference! */
                pbuf_ref(es->p);
            }
            /* chop first pbuf from chain */
            do
            {
                /* try hard to free pbuf */
                freed = pbuf_free(ptr);
            }
            while(freed == 0);
            /* we can read more data now */
            tcp_recved(tpcb, plen);
        }
        else if(wr_err == ERR_MEM)
        {
            /* we are low on memory, try later / harder, defer to poll */
            es->p = ptr;
        }
        else
        {
            /* other problem ?? */
        }
    }
}

/**
 * @brief   closes tcp pcb function
 *
 * @param   tpcb: point to the tcp pcb
 *
 * @param   es: point to tcps_echo_state structure
 *
 * @retval  None
 */
void tcps_echo_close(struct tcp_pcb *tpcb, struct tcps_echo_state *es)
{
    tcp_arg(tpcb, NULL);
    tcp_sent(tpcb, NULL);
    tcp_recv(tpcb, NULL);
    tcp_err(tpcb, NULL);
    tcp_poll(tpcb, NULL, 0);

    if (es != NULL)
    {
        mem_free(es);
    }

    tcp_close(tpcb);
}

#endif /* LWIP_TCP */
