/**
 *   This programis supplied to you by the Author in consideration of your agreement to the following terms, and your use or installation of
 *   The Software and the use of The Documentation constitutes acceptance of these terms.
 *   If you do not agree with these terms, please do not use or install The Software.
 *   The Author grants you a personal, non-exclusive license, under author's copyrights in this original software, to use The Software.
 *   Except as expressly stated in this notice, no other rights or licenses, express or implied, are granted by the Author, including but not limited to any
 *   patent rights that may be infringed by your derivative works or by other works in which The Software may be incorporated.
 *   The Software and the Documentation are provided by the Author on an "AS IS" basis.  THE AUTHOR MAKES NO WARRANTIES, EXPRESS OR IMPLIED, INCLUDING WITHOUT
 *   LIMITATION THE IMPLIED WARRANTIES OF NON-INFRINGEMENT, MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE, REGARDING THE SOFTWARE OR ITS USE AND OPERATION
 *   ALONE OR IN COMBINATION WITH YOUR PRODUCTS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, INDIRECT, INCIDENTAL OR CONSEQUENTIAL DAMAGES (INCLUDING,
 *   BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) ARISING IN ANY WAY OUT OF THE USE,
 *   REPRODUCTION AND MODIFICATION OF THE SOFTWARE AND OR OF THE DOCUMENTATION, HOWEVER CAUSED AND WHETHER UNDER THEORY OF CONTRACT, TORT (INCLUDING NEGLIGENCE),
 *   STRICT LIABILITY OR OTHERWISE, EVEN IF THE AUTHOR HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 **/

// For more information write to fabboco@gmail.com

#include "tcpip_client.h"
#include <string.h>
#include <stdlib.h>

#define DEBUG_printf
// #define DEBUG_printf printf

char _tcp_address[64];
u16_t _tcp_port;

/* Forward declarations */
err_t tcp_client_poll(void *arg, struct tcp_pcb *tpcb);
err_t tcp_client_connected(void *arg, struct tcp_pcb *tpcb, err_t err);
err_t tcp_client_sent(void *arg, struct tcp_pcb *tpcb, u16_t len);
err_t tcp_client_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err);
void tcp_client_err(void *arg, err_t err);

/* ---------------- INIT ---------------- */

TCP_CLIENT_T *tcp_client_init(const char *tcp_address, u16_t tcp_port)
{
    TCP_CLIENT_T *state = calloc(1, sizeof(TCP_CLIENT_T));
    if (!state)
    {
        DEBUG_printf("failed to allocate state\n");
        return NULL;
    }

    strcpy(&_tcp_address[0], tcp_address);
    _tcp_port = tcp_port;

    ip4addr_aton(tcp_address, &state->remote_addr);

    state->connected = false;
    state->complete = false;
    state->buffer_len = 0;
    state->sent_len = 0;
    state->run_count = 0;

    return state;
}

/* ---------------- OPEN CONNECTION ---------------- */

bool tcp_client_open(void *arg)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    DEBUG_printf("Connecting to %s port %u\n",
                 ip4addr_ntoa(&state->remote_addr),
                 _tcp_port);

    state->tcp_pcb = tcp_new_ip_type(IP_GET_TYPE(&state->remote_addr));

    if (!state->tcp_pcb)
    {
        DEBUG_printf("failed to create pcb\n");
        return false;
    }

    tcp_arg(state->tcp_pcb, state);
    tcp_poll(state->tcp_pcb, tcp_client_poll, 4);
    tcp_sent(state->tcp_pcb, tcp_client_sent);
    tcp_recv(state->tcp_pcb, tcp_client_recv);
    tcp_err(state->tcp_pcb, tcp_client_err);

    state->buffer_len = 0;

    cyw43_arch_lwip_begin();
    err_t err = tcp_connect(state->tcp_pcb,
                            &state->remote_addr,
                            _tcp_port,
                            tcp_client_connected);
    cyw43_arch_lwip_end();

    return err == ERR_OK;
}

/* ---------------- CLOSE ---------------- */

err_t tcp_client_close(void *arg)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    if (state->tcp_pcb)
    {
        tcp_arg(state->tcp_pcb, NULL);
        tcp_poll(state->tcp_pcb, NULL, 0);
        tcp_sent(state->tcp_pcb, NULL);
        tcp_recv(state->tcp_pcb, NULL);
        tcp_err(state->tcp_pcb, NULL);

        err_t err = tcp_close(state->tcp_pcb);
        if (err != ERR_OK)
        {
            DEBUG_printf("close failed %d, aborting\n", err);
            tcp_abort(state->tcp_pcb);
        }

        state->tcp_pcb = NULL;
    }

    return ERR_OK;
}

/* ---------------- RESULT WRAPPER ---------------- */

err_t tcp_result(void *arg, int status)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    if (status == 0)
        DEBUG_printf("success\n");
    else
        DEBUG_printf("call failed %d\n", status);

    state->complete = true;

    return tcp_client_close(arg);
}

/* ---------------- POLL ---------------- */

err_t tcp_client_poll(void *arg, struct tcp_pcb *tpcb)
{
    (void)arg;
    (void)tpcb;

    return ERR_OK;
}

/* ---------------- ERROR ---------------- */

void tcp_client_err(void *arg, err_t err)
{
    if (err != ERR_ABRT)
    {
        printf("tcp_client_err %d\n", err);
        tcp_result(arg, err);
    }
}

/* ---------------- CONNECTED ---------------- */

err_t tcp_client_connected(void *arg, struct tcp_pcb *tpcb, err_t err)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    if (err != ERR_OK)
    {
        DEBUG_printf("connect failed %d\n", err);
        return tcp_result(arg, err);
    }

    state->connected = true;
    DEBUG_printf("Connected!\n");

    return ERR_OK;
}

/* ---------------- SENT CALLBACK ---------------- */

err_t tcp_client_sent(void *arg, struct tcp_pcb *tpcb, u16_t len)
{
    (void)arg;
    (void)tpcb;

    return ERR_OK;
}

/* ---------------- RECEIVE ---------------- */

err_t tcp_client_recv(void *arg, struct tcp_pcb *tpcb,
                      struct pbuf *p, err_t err)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    if (!p)
    {
        state->connected = false;
        return ERR_OK;
    }

    uint16_t copy_len = p->len;

    if (copy_len > BUF_SIZE)
        copy_len = BUF_SIZE;

    memcpy(state->buffer, p->payload, copy_len);

    state->buffer_len = copy_len;
    state->data_ready = true;

    tcp_recved(tpcb, p->tot_len);
    pbuf_free(p);

    return ERR_OK;
}

err_t tcp_client_send(void *arg)
{
    TCP_CLIENT_T *state = (TCP_CLIENT_T *)arg;

    if (!state || !state->tcp_pcb)
        return ERR_CONN;

    cyw43_arch_lwip_begin();

    // wait until buffer is available
    if (tcp_sndbuf(state->tcp_pcb) < state->buffer_len)
    {
        DEBUG_printf("TX not ready (sndbuf too small)\n");
        return ERR_MEM;
    }

    err_t err = tcp_write(state->tcp_pcb, state->buffer, state->buffer_len, TCP_WRITE_FLAG_COPY);

    if (err != ERR_OK)
    {
        DEBUG_printf("tcp_write failed REAL = %d\n", err);
        cyw43_arch_lwip_end();
        return err;
    }

    err = tcp_output(state->tcp_pcb);

    if (err != ERR_OK)
    {
        DEBUG_printf("tcp_output failed %d\n", err);
        cyw43_arch_lwip_end();
        return err;
    }

    cyw43_arch_lwip_end();

    return ERR_OK;
}

err_t tcp_client_send_and_wait(
    TCP_CLIENT_T *state,
    const uint8_t *tx,
    uint16_t tx_len,
    char *rx_out,
    uint16_t rx_max,
    uint32_t timeout_ms)
{
    if (!state || !state->tcp_pcb || !state->connected)
        return ERR_CONN;

    DEBUG_printf("TX %u bytes\n", tx_len);

    cyw43_arch_lwip_begin();

    err_t err = tcp_write(state->tcp_pcb,
                          tx,
                          tx_len,
                          TCP_WRITE_FLAG_COPY);

    if (err != ERR_OK)
    {
        DEBUG_printf("tcp_write failed %d\n", err);
        cyw43_arch_lwip_end();
        return err;
    }

    err = tcp_output(state->tcp_pcb);

    cyw43_arch_lwip_end();

    if (err != ERR_OK)
    {
        DEBUG_printf("tcp_output failed %d\n", err);
        return err;
    }

    /* ---------------- WAIT RESPONSE ---------------- */

    uint32_t start = to_ms_since_boot(get_absolute_time());

    state->buffer_len = 0;

    while (1)
    {
        // IMPORTANT: no cyw43_arch_poll() in thread-safe background mode

        if (state->buffer_len > 0)
        {
            uint16_t n = state->buffer_len;

            if (n >= rx_max)
                n = rx_max - 1;

            memcpy(rx_out, state->buffer, n);
            rx_out[n] = '\0';

            return ERR_OK;
        }

        if ((to_ms_since_boot(get_absolute_time()) - start) > timeout_ms)
        {
            DEBUG_printf("timeout waiting response\n");
            return ERR_TIMEOUT;
        }

        sleep_ms(5);
    }
}