#include "tcp_debug.h"
#include "pico/cyw43_arch.h"
#include "lwip/tcp.h"
#include "pico/stdio/driver.h"
#include "pico/sync.h"
#include <string.h>
#include <stdio.h>

#define TCP_PORT 4242
#define DEBUG_BUF_SIZE 2048
#define TCP_IN_BUF_SIZE 20000 // Increased to hold full image + commands

// Circular buffer for debug output
static char debug_buf[DEBUG_BUF_SIZE];
static volatile int write_idx = 0;
static volatile int read_idx = 0;

// Circular buffer for TCP input
static char tcp_in_buf[TCP_IN_BUF_SIZE];
static volatile int in_write_idx = 0;
static volatile int in_read_idx = 0;
static critical_section_t tcp_in_crit;

static struct tcp_pcb *server_pcb = NULL;
static struct tcp_pcb *client_pcb = NULL;

// Helper to write to buffer
static void buffer_write(const char *data, int len) {
    for (int i = 0; i < len; i++) {
        int next_write = (write_idx + 1) % DEBUG_BUF_SIZE;
        if (next_write != read_idx) { // Only write if not full
            debug_buf[write_idx] = data[i];
            write_idx = next_write;
        }
    }
}

// stdio driver callback
static void tcp_out_chars(const char *buf, int len) {
    // Handle CRLF conversion manually if needed, or just write raw
    // The issue is likely that \n is sent without \r, causing "staircase" effect in raw terminals
    for (int i = 0; i < len; i++) {
        if (buf[i] == '\n') {
            char cr = '\r';
            buffer_write(&cr, 1);
        }
        buffer_write(&buf[i], 1);
    }
}

static int tcp_in_chars(char *buf, int len) {
    int count = 0;
    critical_section_enter_blocking(&tcp_in_crit);
    while (count < len && in_read_idx != in_write_idx) {
        buf[count++] = tcp_in_buf[in_read_idx];
        in_read_idx = (in_read_idx + 1) % TCP_IN_BUF_SIZE;
    }
    critical_section_exit(&tcp_in_crit);
    return count > 0 ? count : PICO_ERROR_NO_DATA;
}

static stdio_driver_t tcp_driver = {
    .out_chars = tcp_out_chars,
    .out_flush = NULL,
    .in_chars = tcp_in_chars,
#if PICO_STDIO_ENABLE_CRLF_SUPPORT
    .crlf_enabled = false
#endif
};

// TCP callbacks
static err_t tcp_server_poll(void *arg, struct tcp_pcb *tpcb) {
    if (tpcb != client_pcb) return ERR_OK;

    // Check if there is data in the ring buffer to send
    if (write_idx != read_idx) {
        int len_to_send = 0;
        char send_buf[256]; // Temporary buffer for one TCP write chunk
        
        // Copy from ring buffer to linear buffer
        while (write_idx != read_idx && len_to_send < sizeof(send_buf)) {
            send_buf[len_to_send++] = debug_buf[read_idx];
            read_idx = (read_idx + 1) % DEBUG_BUF_SIZE;
        }

        if (len_to_send > 0) {
            err_t err = tcp_write(tpcb, send_buf, len_to_send, TCP_WRITE_FLAG_COPY);
            if (err == ERR_OK) {
                tcp_output(tpcb);
            } else {
                // If write failed (e.g. buffer full), rewind read_idx (simple retry logic)
                // Note: This is a simplification. In a real robust system we might drop data or handle partial writes.
                // For now, we just lose the data if TCP is full to avoid blocking the system.
            }
        }
    }
    return ERR_OK;
}

static tcp_connection_callback_t g_conn_callback = NULL;

void tcp_set_connection_callback(tcp_connection_callback_t callback) {
    g_conn_callback = callback;
}

static err_t tcp_server_recv(void *arg, struct tcp_pcb *tpcb, struct pbuf *p, err_t err) {
    if (!p) {
        tcp_close(tpcb);
        client_pcb = NULL;
        printf("TCP debug client disconnected\n");
        if (g_conn_callback) g_conn_callback(false);
        return ERR_OK;
    }
    
    // Process received data into input buffer
    if (p->tot_len > 0) {
        critical_section_enter_blocking(&tcp_in_crit);
        struct pbuf *q = p;
        while (q != NULL) {
            char *data = (char *)q->payload;
            for (int i = 0; i < q->len; i++) {
                int next_write = (in_write_idx + 1) % TCP_IN_BUF_SIZE;
                if (next_write != in_read_idx) {
                    tcp_in_buf[in_write_idx] = data[i];
                    in_write_idx = next_write;
                }
            }
            q = q->next;
        }
        critical_section_exit(&tcp_in_crit);
        tcp_recved(tpcb, p->tot_len);
    }
    
    pbuf_free(p);
    return ERR_OK;
}

static err_t tcp_server_accept(void *arg, struct tcp_pcb *newpcb, err_t err) {
    if (client_pcb) {
        // Already have a client, reject or close old?
        // Let's close the old one and accept the new one (last one wins)
        tcp_abort(client_pcb);
        if (g_conn_callback) g_conn_callback(false);
    }
    
    client_pcb = newpcb;
    tcp_recv(newpcb, tcp_server_recv);
    tcp_poll(newpcb, tcp_server_poll, 1); // Poll every 500ms (coarse grain) -> actually arg is interval in 500ms steps? 
    // Wait, poll interval is passed to tcp_poll. 1 means every 500ms.
    // We might want faster polling for debug output.
    // But we can't poll faster than the lwIP timer.
    // However, we can also trigger writes from the main loop if we wanted, but polling is safer for threading.
    
    printf("TCP debug client connected\n");
    if (g_conn_callback) g_conn_callback(true);
    return ERR_OK;
}

bool tcp_debug_init(void) {
    if (server_pcb) return true;

    critical_section_init(&tcp_in_crit);

    // Register stdio driver
    stdio_set_driver_enabled(&tcp_driver, true);

    cyw43_arch_lwip_begin();
    server_pcb = tcp_new_ip_type(IPADDR_TYPE_ANY);
    if (!server_pcb) {
        cyw43_arch_lwip_end();
        return false;
    }

    err_t err = tcp_bind(server_pcb, IP_ANY_TYPE, TCP_PORT);
    if (err != ERR_OK) {
        cyw43_arch_lwip_end();
        return false;
    }

    server_pcb = tcp_listen(server_pcb);
    tcp_accept(server_pcb, tcp_server_accept);
    cyw43_arch_lwip_end();

    printf("TCP debug server listening on port %d\n", TCP_PORT);
    return true;
}