#include <stdio.h>
#include <stdint.h>

#include "pico/cyw43_arch.h"
#include "hardware/watchdog.h"

#include "lwip/def.h"
#include "lwip/tcp.h"

#include "modbus.h"

#include "logging.h"

#define POLL_TIME_S 5

struct internal_state {
    struct __attribute__((__packed__)) {
        struct moddbus_tcp_packet header;
        struct modbus_request body;
    } request;
    struct modbus_response* response;

    bool done;

    struct tcp_pcb* conn;
};
static struct internal_state state = {0};

static uint16_t get_identifier(void)
{
    static uint16_t global_identifier = 0;

    return global_identifier++;
}

static err_t modbus_client_close(void* arg, err_t err)
{
    struct internal_state* state = arg;

    if (state->conn != NULL) {
        LogDebug("cleaning up connection");

        tcp_arg(state->conn, NULL);
        tcp_poll(state->conn, NULL, 0);
        tcp_sent(state->conn, NULL);
        tcp_recv(state->conn, NULL);
        tcp_err(state->conn, NULL);
        cyw43_arch_lwip_begin();
        err = tcp_close(state->conn);
        cyw43_arch_lwip_end();
        if (err != ERR_OK) {
            LogError("close failed %d, calling abort\n", err);
            tcp_abort(state->conn);
            err = ERR_ABRT;
        }
        state->conn = NULL;
    }
    state->done = true;
    return err;
}

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

    LogDebug("modbus_client_poll");
    return modbus_client_close(arg, ERR_OK);
}

static err_t modbus_client_recv(void *arg, struct tcp_pcb* conn, struct pbuf *p, err_t err)
{
    struct internal_state* state = arg;

    if (!p)
    {
        LogDebug("connected closed");
        return modbus_client_close(arg, ERR_OK);
    }
    else if (err != ERR_OK)
    {
        LogInfo("something went wrong (%i), cleaning up early", err);
        if (p != NULL)
        {
            pbuf_free(p);
        }

        return modbus_client_close(arg, err);
    }

    struct __attribute__((__packed__)) response {
        struct moddbus_tcp_packet header;
        struct modbus_response body;
    };

    struct response* response = (struct response*)p->payload;

    if (response->header.identifier != state->request.header.identifier)
    {
        LogError("invalid identifier received (expected %i, found %i", ntohs(state->request.header.identifier), ntohs(response->header.identifier));
        goto cleanup;
    }

    if (response->header.protocol != state->request.header.protocol)
    {
        LogError("invalid protocol received (expected %i, found %i", ntohs(state->request.header.protocol), ntohs(response->header.protocol));
        goto cleanup;
    }

    if (response->header.unit != state->request.header.unit)
    {
        LogError("invalid unit received (expected %i, found %i", state->request.header.unit, response->header.unit);
        goto cleanup;
    }

    if (response->body.code != state->request.body.code)
    {
        LogError("invalid response code received (expected %i, found %i", state->request.body.code, response->body.code);
        goto cleanup;
    }

    size_t body_size = ntohs(response->header.length) - 1;

    state->response = malloc(body_size);
    if (!state->response)
    {
        LogError("failed to allocate response buffer");
        goto cleanup;
    }

    LogDebug("copying into response");
    memcpy(state->response, &response->body, body_size);

cleanup:
    tcp_recved(conn, p->tot_len);
    LogDebug("packet marked as received");

    LogDebug("freed package");
    pbuf_free(p);

    state->done = true;

    return ERR_OK;
}

static void modbus_client_err(void *arg, err_t err) {
    if (err != ERR_ABRT) {
        LogError("connection errored");
        LogDebug("tcp_client_err %d", err);
        modbus_client_close(arg, ERR_OK);
    }
}

err_t modbus_open_connection(void)
{
    err_t err = ERR_OK;

    if (!state.conn)
    {
        LogInfo("creating new connection");
        state.conn = tcp_new_ip_type(IPADDR_TYPE_V4);

        tcp_arg(state.conn, &state);
        tcp_poll(state.conn, modbus_client_poll, POLL_TIME_S * 2);
        tcp_recv(state.conn, modbus_client_recv);
        tcp_err(state.conn, modbus_client_err);

        ip_addr_t dest;
        IP4_ADDR(&dest, 192, 168, 178, 65);

        LogDebug("starting connection");
        cyw43_arch_lwip_begin();
        err = tcp_connect(state.conn, &dest, MODBUS_PORT, NULL);
        cyw43_arch_lwip_end();
    }

    return err;
}

err_t modbus_close_connection(void)
{
    err_t err = modbus_client_close(&state, ERR_OK);
    if (err == ERR_OK)
    {
        while (state.conn != NULL)
        {
            sleep_ms(1000);
        }
    }

    return err;
}

err_t modbus_read_registers(uint16_t address, uint16_t count, struct modbus_response** response_out)
{

    err_t err;
    state.response = NULL;
    state.done = false;

    if (!state.conn)
    {
        LogInfo("creating new connection");
        err = modbus_open_connection();
        if (err != ERR_OK)
        {
            LogInfo("connection failed (%i)", err);
            return modbus_client_close(&state, ERR_OK);
        }
    }
    else
    {
        LogInfo("reusing connection");
    }


    LogDebug("setting up connection parameters");
    state.request.header.identifier = htons(get_identifier());
    state.request.header.protocol = htons(MODBUS_PROTOCOL);
    state.request.header.length = htons(sizeof(state.request.body) + 1);
    state.request.header.unit = 0x01;

    state.request.body.code = 0x04;
    state.request.body.address = htons(address-1);
    state.request.body.count = htons(count);

    LogDebug("writing request");
    cyw43_arch_lwip_begin();
    err = tcp_write(state.conn, &state.request, sizeof(state.request), 0);
    cyw43_arch_lwip_end();
    if (err != ERR_OK)
    {
        LogError("failed to write request (%i)", err);
        return modbus_client_close(&state, ERR_OK);
    }

    LogDebug("sending request");
    cyw43_arch_lwip_begin();
    err = tcp_output(state.conn);
    cyw43_arch_lwip_end();
    if (err != ERR_OK)
    {
        LogError("failed to send request (%i)", err);
        return modbus_client_close(&state, ERR_OK);
    }

    LogDebug("waiting until response is given");
    while (state.done != true)
    {
        watchdog_update();
        sleep_ms(1000);
    }

    LogDebug("request done");

    *response_out = state.response;
    state.response = NULL;

    return ERR_OK;
}