#define AL_LOG_SECTION "multiplex"
//#define AL_LOG_ENABLE_TRACE
#include <al/log.h>
#include <al/atomic.h>

#include "multiplex.h"

bool nn_multiplex_socket_init(struct nn_multiplex_socket *multi, u8 type,
    bool (*connection_callback)(void *, u8, struct nn_packet_stream *), void *userdata)
{
    multi->sock.type = type;
    s32 flags = NNWT_SOCKET_NONBLOCKING;
    if (type == NNWT_SOCKET_TCP) flags |= NNWT_SOCKET_REUSE_ADDR;
    if (!nn_socket_init(&multi->sock, flags)) {
        return false;
    }
    multi->connection_callback = connection_callback;
    multi->userdata = userdata;
    return true;
}

static void socket_read_callback(struct ev_loop *loop, ev_io *w, s32 revents)
{
    struct nn_multiplex_connection *conn = (struct nn_multiplex_connection *)w->data;
    struct nn_multiplex_socket *multi = conn->multi;
    (void)revents;

    u8 id;
    ssize_t ret = nn_socket_read(&conn->sock, &id, 1);
    if (ret <= 0) {
        ev_io_stop(loop, &conn->event);
        nn_socket_close(&conn->sock);
        al_free(conn);
        if (ret == 0) log_warn("read() returned 0 on POLLIN.");
        return;
    }

    log_trace("ID byte: 0x%hhx.", id);
    switch (id) {
    case 0x48: // '(H)EAD'
        break;
    case 0x47: // '(G)ET'
        // Attempt WebSocket upgrade.
        break;
    case 0x50: // '(P)OST'
        break;
    }

    ev_io_stop(loop, &conn->event);

    struct nn_packet_stream *stream = al_alloc_object(struct nn_packet_stream);
    if (multi->connection_callback(multi->userdata, id, stream)) {
        nn_packet_stream_from_socket(stream, multi->loop, &conn->sock);
    } else {
        nn_socket_close(&conn->sock);
    }

    al_free(conn);
}

static void socket_connection_callback(struct ev_loop *loop, ev_io *w, s32 revents)
{
    struct nn_multiplex_socket *multi = (struct nn_multiplex_socket *)w->data;
    (void)revents;
    struct nn_multiplex_connection *conn = al_alloc_object(struct nn_multiplex_connection);
    if (!nn_socket_accept(&multi->sock, &conn->sock, NNWT_SOCKET_NONBLOCKING)) {
        al_free(conn);
        return;
    }
    conn->multi = multi;
    conn->event.data = conn;
    ev_io_init_n(&conn->event, socket_read_callback, nn_socket_get_fd(&conn->sock), EV_READ);
    ev_io_start(loop, &conn->event);
}

bool nn_multiplex_socket_listen(struct nn_multiplex_socket *multi, struct nn_event_loop *loop, str *addr, u16 port)
{
    if (!nn_socket_bind(&multi->sock, addr, port) || !nn_socket_listen(&multi->sock)) {
        return false;
    }
    multi->loop = loop;
    multi->event.data = multi;
    ev_io_init_n(&multi->event, socket_connection_callback, nn_socket_get_fd(&multi->sock), EV_READ);
    ev_io_start(multi->loop->ev, &multi->event);
    return true;
}

void nn_multiplex_socket_close(struct nn_multiplex_socket *multi)
{
    ev_io_stop(multi->loop->ev, &multi->event);
    nn_socket_shutdown(&multi->sock);
    nn_socket_close(&multi->sock);
    nn_socket_cleanup(&multi->sock);
}

static struct nn_multiplex_direct direct_global = { 0 };

void nn_multiplex_direct_init(bool (*connection_callback)(void *, u8, struct nn_packet_stream *), void *userdata)
{
    direct_global.connection_callback = connection_callback;
    direct_global.userdata = userdata;
}

static void queue_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct nn_multiplex_bridge *bridge = (struct nn_multiplex_bridge *)userdata;
    (void)stream;
    al_array_push(bridge->queue, packet);
}

static void direct_connect(struct nn_packet_stream *client, u8 id)
{
    struct nn_multiplex_bridge *bridge = al_alloc_object(struct nn_multiplex_bridge);
    bridge->broken = false;
    al_array_init(bridge->queue);
    bridge->bridge = al_alloc_object(struct nn_packet_stream);
    nn_packet_stream_init(bridge->bridge, NULL, NULL, NULL);
    // Server-side client, emulates the result of socket_accept().
    struct nn_packet_stream *cl = al_alloc_object(struct nn_packet_stream);
    nn_packet_stream_init(cl, NULL, NULL, NULL);
    cl->is_direct = true;
    client->is_direct = true;
    atomic_store(void)(&cl->direct, client, AL_ATOMIC_RELAXED);
    atomic_store(void)(&client->direct, cl, AL_ATOMIC_RELAXED);
    cl->bridge = bridge;
    cl->packet_callback = queue_packet_callback;
    cl->userdata = bridge;
    // The order of connect client -> global callback -> connect server-side client cannot change.
    nn_packet_stream_set_connected(client);
    direct_global.connection_callback(direct_global.userdata, id, cl);
    nn_packet_stream_set_connected(cl);
    // Resend any packets that might have been sent in client->connection_callback().
    struct nn_packet *packet;
    al_array_foreach(bridge->queue, i, packet) {
        al_assert(cl->packet_callback != queue_packet_callback);
        cl->packet_callback(cl->userdata, cl, packet);
        if (bridge->broken) break;
    }
    bridge->queue.count = 0;
    if (bridge->broken) {
        nn_multiplex_bridge_free(bridge);
    }
}

void nn_multiplex_direct_connect(struct nn_packet_stream *client, u8 id)
{
    client->id = id;
    direct_connect(client, id);
}

void nn_multiplex_direct_reconnect(struct nn_packet_stream *client)
{
    direct_connect(client, client->id);
}

void nn_multiplex_bridge_free(struct nn_multiplex_bridge *bridge)
{
    nn_packet_stream_free(bridge->bridge);
    al_free(bridge->bridge);
    al_array_free(bridge->queue);
    al_free(bridge);
}