#include <stdio.h>
#include <stdlib.h>
#include <aos/aos.h>
#include <aos/waitset_chan.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>
#include <netutil/ip.h>

#define BUF_SIZE (1500 - 20 - 8)
struct send_buf {
    struct ump_net_call call;
    uint8_t payload[BUF_SIZE];
    size_t send_len;
    struct waitset_chanstate chan;
};

static uint16_t local_port;
static uint16_t remote_port;
static uint32_t remote_ip;

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

static void net_recv_packet (void *arg, size_t payload_size, void *payload) {
    printf("%.*s", payload_size, arg);
    free(arg);
}

static void handle_udp_recv (void *arg, struct ump_net_ev_udp_recv *udp_recv, size_t payload_size) {
    void *payload = malloc(payload_size);
    if (payload == NULL) {
        printf("ERROR: malloc failed\n");
        ump_net_recv_payload(NULL, NULL, NULL);
    } else {
        if (remote_ip == 0) remote_ip = udp_recv->src_ip;
        if (remote_port == 0) remote_port = udp_recv->src_port;
        ump_net_recv_payload(payload, net_recv_packet, payload);
    }
}

static void send_callback (void *arg, errval_t err) {
    struct send_buf *send_buf = arg;
    free(send_buf);
    if (err_is_fail(err)) {
        DEBUG_ERR(err, "failed to send packet");
    }
}

static void udp_send(void *arg) {
    struct send_buf *send_buf = arg;
    ump_net_udp_send(
        &send_buf->call,
        remote_ip,
        local_port, remote_port,
        send_buf->send_len, &send_buf->payload,
        send_callback, send_buf
    );
}

static struct send_buf *create_send_buffer (void) {
    struct send_buf *send_buf = malloc(sizeof(struct send_buf));
    if (send_buf == NULL) {
        printf("ERROR: malloc failed\n");
        return NULL;
    }
    waitset_chanstate_init(&send_buf->chan, CHANTYPE_OTHER);
    send_buf->send_len = 0;
    return send_buf;
}

static struct waitset_chanstate exitchan;
static struct ump_net_call stopcall;

static void udp_stop_callback (void *arg, errval_t err) {
    exit(0);
}

static void exit_process(void *arg) {
    ump_net_udp_listen_stop(&stopcall, local_port, udp_stop_callback, NULL);
}

static int input_reader (void *arg) {
    errval_t err;
    struct aos_rpc *rpc = aos_rpc_get_serial_channel();
    struct waitset *default_ws = get_default_waitset();
    struct send_buf *send_buf = create_send_buffer();
    if (send_buf == NULL) return 1;
    while (true) {
        char in;
        err = aos_rpc_serial_getchar(rpc, &in);
        if (err_is_fail(err)) {
            DEBUG_ERR(err, "failed to read");
            return 1;
        }
        if (in == '\x03') { // Ctrl+C
            waitset_chan_trigger_closure(default_ws, &exitchan, MKCLOSURE(exit_process, NULL));
            return 0;
        }
        if (in == '\0') in = '\n';
        send_buf->payload[send_buf->send_len] = in;
        send_buf->send_len++;
        if (send_buf->send_len == sizeof(send_buf->payload) || in == '\n') {
            if (remote_port == 0 || remote_ip == 0) {
                printf("Error: Can't send, remote ip/port not provided and no packet received yet.\n");
                send_buf->send_len = 0;
            } else {
                waitset_chan_trigger_closure(default_ws, &send_buf->chan, MKCLOSURE(udp_send, send_buf));
                send_buf = create_send_buffer();
                if (send_buf == NULL) return 1;
            }
        }
    }
}

int main (int argc, char *argv[]) {
    errval_t err;
    struct waitset *default_ws = get_default_waitset();
    if (argc < 2) {
        printf("Usage: netcat local_port [remote_ip] [remote_port]\n");
        return 1;
    }
    local_port = strtol(argv[1], NULL, 10);

    if (argc > 2) {
        int	a1, a2, a3, a4;
        int parsed = sscanf(argv[2], "%d.%d.%d.%d.", &a1, &a2, &a3, &a4);
        if (parsed != 4) {
            printf("Error: failed to parse IP address\n");
            return 1;
        }
        remote_ip = MK_IP(a1, a2, a3, a4);
        remote_port = local_port;
    }
    if (argc > 3) {
        remote_port = strtol(argv[3], NULL, 10);
    }

    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, local_port, handle_udp_recv, NULL, udp_listen_callback, NULL);

    thread_create(input_reader, NULL);

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