#include <stdio.h>
#include <stdlib.h>
#include <aos/aos.h>
#include <aos/simpleslab.h>
#include <aos/deferred.h>
#include <aos/ump_net.h>
#include <aos/ump_net_client.h>

// Echo Protocol: https://datatracker.ietf.org/doc/html/rfc862
#define PORT_ECHO 7

static uint16_t listen_port = PORT_ECHO;

// UDP echo server

#define BUF_SIZE 1500
struct echo_buf {
    uint32_t src_ip;
    uint16_t src_port;
    struct ump_net_call call;
    uint8_t payload[BUF_SIZE];
};
static struct simpleslab_allocator slabs;
static uint8_t slab_buf[sizeof(struct echo_buf) * 128];


static void udp_listen_callback (void *arg, errval_t err, uint16_t src_port) {
    if (err_is_fail(err)) {
        printf("Failed to listen on UDP port %d:\n", listen_port);
        err_print_calltrace(err);
        return;
    }
    DEBUG_PRINTF("Listening on UDP port %d.\n", src_port);
}

static void net_reply_callback (void *arg, errval_t err) {
    struct echo_buf *echo_buf = arg;
    simpleslab_free(&slabs, echo_buf);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to send packet");
    }
}

static struct ump_net_ev_udp_recv net_udp_recv;

static void net_recv_packet (void *arg, size_t payload_size, void *payload) {
    struct echo_buf *echo_buf = arg;
    //printf("Echo server: Received payload: '%.*s'\n", payload_size, &echo_buf->payload);

    // Send reply
    ump_net_udp_send(
        &echo_buf->call,
        net_udp_recv.src_ip,
        listen_port, net_udp_recv.src_port,
        payload_size, &echo_buf->payload,
        net_reply_callback, echo_buf
    );
}

static void handle_udp_recv (void *arg, struct ump_net_ev_udp_recv *udp_recv, size_t payload_size) {
    if (payload_size > BUF_SIZE) {
        printf("Echo server: ERROR: Packet too big\n");
    } else {
        struct echo_buf *echo_buf = simpleslab_alloc(&slabs);
        if (echo_buf == NULL) {
            printf("Echo server: WARN: All buffers full, dropping packet\n");
        } else {
            net_udp_recv = *udp_recv;
            ump_net_recv_payload(&echo_buf->payload, net_recv_packet, echo_buf);
            return;
        }
    }
    ump_net_recv_payload(NULL, NULL, NULL);
}

// TCP echo server

#define TCP_MAX_CONNECTIONS 100
struct tcp_conn {
    struct tcp_tcb tcb;
    uint8_t rcv_buf[4*4096];
    uint8_t snd_buf[4*4096];
};

static int tcp_connection_count = 0;

static void tcp_listen_callback (void *arg, errval_t err, uint16_t src_port) {
    if (err_is_fail(err)) {
        printf("Failed to listen on TCP port %d:\n", listen_port);
        err_print_calltrace(err);
        return;
    }
    DEBUG_PRINTF("Listening on TCP port %d.\n", src_port);
}

static void tcp_receive_available_handler (void *arg) {
    struct tcp_conn *conn = arg;
    uint8_t buf[2048];
    while (true) {
        uint32_t read_size = ump_net_tcp_receive(&conn->tcb, buf, sizeof(buf));
        if (read_size == 0) break;
        debug_printf("TCP received: '%.*s'\n", read_size, buf);
    }
}

static void tcp_send_available_handler (void *arg) {
    // struct tcp_conn *conn = arg;
    debug_printf("tcp_send_available_handler\n");
}

static void tcp_destroy_handler (void *arg) {
    struct tcp_conn *conn = arg;
    debug_printf("TCP connection destroyed for %x, port %d\n", conn->tcb.remote_ip, conn->tcb.remote_port);
    free(conn);
    tcp_connection_count--;
}

static void handle_tcp_connect (
    void *arg, struct ump_net_ev_tcp_recv *tcp_recv,
    uint32_t remote_ip, uint16_t remote_port, uint16_t local_port
) {
    debug_printf("Received connection attempt from %x, port %d\n", remote_ip, remote_port);
    if (tcp_connection_count > TCP_MAX_CONNECTIONS) {
        ump_net_tcp_refuse(tcp_recv);
    } else {
        struct tcp_conn *conn = malloc(sizeof(struct tcp_conn));
        if (conn == NULL) {
            ump_net_tcp_refuse(tcp_recv);
        } else {
            tcp_connection_count++;
            ump_net_tcp_accept(
                tcp_recv, &conn->tcb,
                tcp_receive_available_handler,
                tcp_send_available_handler,
                tcp_destroy_handler,
                conn,
                conn->rcv_buf, sizeof(conn->rcv_buf),
                conn->snd_buf, sizeof(conn->snd_buf)
            );
        }
    }
}

int main (int argc, char *argv[]) {
    errval_t err;
    struct waitset *default_ws = get_default_waitset();

    if (argc >= 2) {
        listen_port = strtol(argv[1], NULL, 10);
    }

    simpleslab_init(&slabs, sizeof(struct echo_buf), slab_buf, sizeof(slab_buf));

    err = ump_net_init(default_ws);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to connect to net server");

    struct ump_net_call call_udp_listen;
    ump_net_udp_listen(&call_udp_listen, listen_port, handle_udp_recv, NULL, udp_listen_callback, NULL);

    struct ump_net_call call_tcp_listen;
    ump_net_tcp_listen(&call_tcp_listen, listen_port, handle_tcp_connect, NULL, tcp_listen_callback, NULL);

    while (true) {
        err = event_dispatch(default_ws);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "in event_dispatch");
            abort();
        }
    }
}