#include <zephyr/kernel.h>
#include <zephyr/net/socket.h>

#include <stdio.h>
#include <string.h>

#include "tcp_server.h"

#define TCP_SERVER_PORT         5000

#define TCP_THREAD_STACK_SIZE   2048
#define TCP_THREAD_PRIORITY     6

#define TCP_RX_BUFFER_SIZE      256
#define TCP_PAYLOAD_MAX_SIZE    (50 * 1024)

K_THREAD_STACK_DEFINE(tcp_stack, TCP_THREAD_STACK_SIZE);
static struct k_thread tcp_thread_data;

static uint8_t payload_buffer[TCP_PAYLOAD_MAX_SIZE];


/* Function prototypes */
static void tcp_server_thread(void *p1, void *p2, void *p3);
static int recv_exact(int sock, uint8_t *buffer, size_t length);
static bool verify_payload(uint8_t *buffer, uint32_t length);
static int send_exact(int sock, const uint8_t *buffer, size_t length);


/* TCP server start */
void tcp_server_start(void)
{
    k_thread_create(&tcp_thread_data,
                    tcp_stack,
                    K_THREAD_STACK_SIZEOF(tcp_stack),
                    tcp_server_thread,
                    NULL, NULL, NULL,
                    TCP_THREAD_PRIORITY,
                    0,
                    K_NO_WAIT);
}


/* Receive exact number of bytes */
static int recv_exact(int sock, uint8_t *buffer, size_t length)
{
    size_t total_received = 0;
    int received;

    while (total_received < length) {

        received = zsock_recv(
            sock,
            &buffer[total_received],
            length - total_received,
            0);

        if (received > 0) {
            total_received += received;

        } else if (received == 0) {
            printk("TCP client closed connection\n");
            return -1;

        } else {
            printk("TCP receive failed: %d\n", errno);
            return -1;
        }
    }

    return 0;
}


/* Verify received payload pattern */
static bool verify_payload(uint8_t *buffer, uint32_t length)
{
    for (uint32_t i = 0; i < length; i++) {

        uint8_t expected = (uint8_t)(i % 256);

        if (buffer[i] != expected) {

            printk("Data mismatch at %u: RX=%02X, Expected=%02X\n",
                   i,
                   buffer[i],
                   expected);

            return false;
        }
    }

    return true;
}


/* TCP server thread */
static void tcp_server_thread(void *p1, void *p2, void *p3)
{
    int server_fd;
    int client_fd;
    int ret;

    uint64_t start_cycle;
    uint64_t end_cycle;
    uint64_t elapsed_us;

    uint8_t rx_buffer[TCP_RX_BUFFER_SIZE];

    uint32_t payload_length_net;
    uint32_t payload_length;

    uint32_t total_received;
    int received;

    struct sockaddr_in server_addr;
    struct sockaddr_in client_addr;

    char ack_buffer[64];
    int ack_length;

    socklen_t client_addr_len = sizeof(client_addr);

    ARG_UNUSED(p1);
    ARG_UNUSED(p2);
    ARG_UNUSED(p3);


    server_fd = zsock_socket(
        AF_INET,
        SOCK_STREAM,
        IPPROTO_TCP);

    if (server_fd < 0) {
        printk("TCP socket creation failed: %d\n", errno);
        return;
    }


    memset(&server_addr, 0, sizeof(server_addr));

    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(TCP_SERVER_PORT);
    server_addr.sin_addr.s_addr = htonl(INADDR_ANY);


    ret = zsock_bind(
        server_fd,
        (struct sockaddr *)&server_addr,
        sizeof(server_addr));

    if (ret < 0) {
        printk("TCP bind failed: %d\n", errno);
        zsock_close(server_fd);
        return;
    }


    ret = zsock_listen(server_fd, 1);

    if (ret < 0) {
        printk("TCP listen failed: %d\n", errno);
        zsock_close(server_fd);
        return;
    }


    printk("TCP server listening on port %d\n", TCP_SERVER_PORT);


    while (1) {

        client_addr_len = sizeof(client_addr);

        client_fd = zsock_accept(
            server_fd,
            (struct sockaddr *)&client_addr,
            &client_addr_len);

        if (client_fd < 0) {
            printk("TCP accept failed: %d\n", errno);
            continue;
        }


        printk("TCP client connected\n");


        /* Step 1: Receive 4-byte payload length */

        ret = recv_exact(
            client_fd,
            (uint8_t *)&payload_length_net,
            sizeof(payload_length_net));

        if (ret < 0) {
            zsock_close(client_fd);
            continue;
        }


        payload_length = ntohl(payload_length_net);

        printk("Payload length: %u bytes\n", payload_length);


        if (payload_length > TCP_PAYLOAD_MAX_SIZE) {

            printk("Payload too large\n");

            zsock_close(client_fd);
            continue;
        }


        /* Step 2: Receive payload */

        total_received = 0;

        start_cycle = k_cycle_get_64();

        while (total_received < payload_length) {

            uint32_t remain = payload_length - total_received;

            uint32_t recv_size =
                (remain > TCP_RX_BUFFER_SIZE) ?
                TCP_RX_BUFFER_SIZE :
                remain;


            received = zsock_recv(
                client_fd,
                rx_buffer,
                recv_size,
                0);


            if (received > 0) {

                memcpy(
                    &payload_buffer[total_received],
                    rx_buffer,
                    received);

                total_received += received;

            } else if (received == 0) {

                printk("TCP client closed connection\n");
                break;

            } else {

                printk("TCP receive failed: %d\n", errno);
                break;
            }
        }

        end_cycle = k_cycle_get_64();

        elapsed_us =
            k_cyc_to_us_floor64(end_cycle - start_cycle);


        /* Step 3: Verify payload */

        if (total_received == payload_length) {

            printk("Payload received successfully: %u bytes\n",
                   total_received);

            printk("Receive time: %llu us\n",
                    elapsed_us);       

            if (verify_payload(payload_buffer, payload_length)) {

                printk("Payload data verification OK\n");

                ack_length = snprintf(
                    ack_buffer,
                    sizeof(ack_buffer),
                    "ACK,%u,%llu\n",
                    payload_length,
                    elapsed_us);

                if (ack_length > 0) {

                    ret = send_exact(
                        client_fd,
                        (const uint8_t *)ack_buffer,
                        ack_length);

                    if (ret == 0) {

                        printk("ACK sent: %s", ack_buffer);

                    } else {

                        printk("ACK send failed\n");
                    }
                }                

            } else {

                printk("Payload data verification FAILED\n");
            }
        }


        zsock_close(client_fd);

        printk("TCP client disconnected\n");
    }
}

/* Send exact number of bytes */
static int send_exact(int sock, const uint8_t *buffer, size_t length)
{
    size_t total_sent = 0;
    int sent;

    while (total_sent < length) {

        sent = zsock_send(
            sock,
            &buffer[total_sent],
            length - total_sent,
            0);

        if (sent > 0) {

            total_sent += sent;

        } else {

            printk("TCP send failed: %d\n", errno);
            return -1;
        }
    }

    return 0;
}