#include <stdio.h>
#include <stdlib.h>
#include <aos/aos.h>
#include <aos/simpleslab.h>
#include <aos/deferred.h>
#include <aos/aos_rpc.h>
#include <aos/ump_binding.h>
#include <aos/ump_chan.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

#define ECHO_TEST_MESSAGE "Hello Client! Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum. Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt mollit anim id est laborum."

__attribute__((__used__))
static void handle_payload_recv(void *arg, size_t payload_size, void *payload) {
    printf("Echo server: Received payload(%lu): %.*s\n", payload_size, payload_size, payload);
}

__attribute__((__used__))
static void handle_header_recv(void *arg, size_t header_size, void *header, size_t payload_size) {
    printf("Echo server: Received header: %u\n", *(uint8_t *)header);
    printf("Echo server: Payload is %lu bytes\n", payload_size);
    void *payload = malloc(payload_size);
    assert(payload != NULL);
    ump_recv_payload(arg, payload, handle_payload_recv, NULL);
}

__attribute__((__used__))
static void handle_send(void *arg, struct ump_send_queue_entry *entry) {
    printf("Echo server: Message sent!\n");
}

static errval_t connect_server (void *arg, struct capref cap) {
    errval_t err;

    printf("Echo server: incoming connection\n");

    struct ump_send_chan *send_chan;
    struct ump_recv_chan *recv_chan;
    err = ump_chan_init(UMP_ROLE_SERVER, &send_chan, &recv_chan, 1, cap, get_default_waitset());
    if (err_is_fail(err)) return err;

    uint8_t *value_to_send = malloc(sizeof(uint8_t));
    *value_to_send = 123;
    struct ump_send_queue_entry *entry = malloc(sizeof(struct ump_send_queue_entry));
    printf("Echo server: Sending Message\n");
    ump_send(send_chan, entry, 1, value_to_send, sizeof(ECHO_TEST_MESSAGE), ECHO_TEST_MESSAGE, handle_send, NULL);

    printf("Echo server: Receiving Message\n");
    ump_recv_header(recv_chan, handle_header_recv, recv_chan);

    return SYS_ERR_OK;
}


// net 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 net_listen_callback (void *arg, errval_t err, uint16_t src_port) {
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to listen on UDP");
        return;
    }
    printf("Echo server: Listening on UDP port %d.\n", src_port);
}

static void net_send_callback (void *arg, errval_t err) {
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to send packet");
    }
}

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,
        PORT_ECHO, 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);
}

static void send_again (void *arg) {
    static struct ump_net_call call;
    ump_net_udp_send(&call, 0xc0a80204, PORT_ECHO, 7000, 12, "hello again\n", net_send_callback, NULL);
}

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

    printf("Echo server: registering\n");
    struct ump_binding_server server;
    err = ump_binding_register(&server, UMP_SERVER_ECHO, connect_server, NULL);
    if (err_is_fail(err)) USER_PANIC_ERR(err, "Failed to register UMP server");
    printf("Echo server: listening\n");


    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;
    ump_net_udp_listen(&call, PORT_ECHO, handle_udp_recv, NULL, net_listen_callback, NULL);

    struct ump_net_call call2;
    ump_net_udp_send(&call2, 0xc0a80204, PORT_ECHO, 7000, 6, "hello\n", net_send_callback, NULL);

    struct deferred_event dev;
    deferred_event_init(&dev);
    err = deferred_event_register(&dev, default_ws, 1000000,
                                  MKCLOSURE(send_again, NULL));
    if (err_is_fail(err)) return err;


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