#include "server.h"
#include "handlers.h"
#include "list.h"

static inline u32 get_incremental_id(struct lia_server *server)
{
    server->increment = al_u32_inc_wrap(server->increment);
    return server->increment;
}

bool lia_server_init(struct lia_server *server, struct nn_event_loop *loop)
{
    server->loop = loop;
    server->increment = 0;
    al_array_init(server->nodes);
    al_array_init(server->dormant_connections);
    return true;
}

static void remove_dormant_connection(struct lia_server *server, struct nn_packet_stream *stream)
{
    al_array_remove(server->dormant_connections, stream);
}

static void data_packet_sent_callback(void *userdata, struct nn_packet *packet)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    nn_packet_pool_lock(&conn->pool);
    nn_packet_pool_return(&conn->pool, packet);
    nn_packet_pool_unlock(&conn->pool);
}

static void data_packets_sent_callback(void *userdata, struct nn_packet **packets, u32 count)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    nn_packet_pool_lock(&conn->pool);
    for (u32 i = 0; i < count; i++) {
        if (packets[i]) {
            nn_packet_pool_return(&conn->pool, packets[i]);
        }
    }
    nn_packet_pool_unlock(&conn->pool);
}

static void packet_pool_callback(void *userdata, struct nn_packet *packet)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    if (!nn_packet_stream_send_packet(conn->stream, packet)) {
        nn_packet_pool_lock(&conn->pool);
        nn_packet_pool_return(&conn->pool, packet);
        nn_packet_pool_unlock(&conn->pool);
    }
}

static nn_thread_result NNWT_THREADCALL handler_thread(void *userdata)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    nn_thread_set_name("liana_handler");

    if (conn->seek_pos != LIANA_TIMESTAMP_INVALID) {
        conn->handler->seek(conn->handler, conn->seek_pos);
        conn->seek_pos = LIANA_TIMESTAMP_INVALID;
    }

    for (;;) {
        struct nn_packet *packet = nn_packet_pool_get(&conn->pool);
        if (!packet) {
            return 0;
        }

        conn->handler->step(conn->handler);
        conn->handler->write_packet(conn->handler, packet);

#ifdef LIANA_SERVER_LOOP
        if (conn->node->duration > 0 && conn->handler->status == CAMU_ERR_EOF) {
            nn_packet_pool_lock(&conn->pool);
            nn_packet_pool_return(&conn->pool, packet);
            nn_packet_pool_unlock(&conn->pool);
            continue;
        }
#endif

        nn_packet_pool_submit(&conn->pool, packet);

        // Check status after submitting so the EOF packet gets sent.
        if (conn->handler->status != CAMU_OK) break;
    }

    nn_packet_pool_flush(&conn->pool);

    return 0;
}

static void discard_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    (void)userdata;
    // @TODO: This should invalidate the connection instead of asserting.
    nn_packet_stream_return_packet(stream, packet);
    al_assert(false);
}

static bool should_free_node(struct lia_node *node)
{
    return node->closed && !node->handler && !node->requests.count && !node->connections.count;
}

static void free_node(struct lia_node *node)
{
    struct lia_server *server = node->server;
    cch_entry_free(&node->entry);
    al_assert(!node->requests.count);
    al_array_free(node->requests);
    al_assert(!node->connections.count);
    al_array_free(node->connections);
    al_array_remove(server->nodes, node);
    al_free(node);
}

static void free_connection(struct lia_node_connection *conn)
{
    struct lia_node *node = conn->node;
    al_assert(conn->handler);
    conn->handler->free(&conn->handler);
    cch_entry_return_handle(conn->node->entry, &conn->handle);
    nn_packet_pool_free(&conn->pool);
    bool removed = al_array_remove(node->connections, conn);
    al_assert(removed);
    al_free(conn);
    if (should_free_node(node)) {
        free_node(node);
    }
}

static void free_connection_stream(struct lia_node_connection *conn)
{
    nn_packet_stream_free(conn->stream);
    al_free(conn->stream);
    conn->stream = NULL;
    conn->ref = false;
}

static void disable_connection_and_wait(struct lia_node_connection *conn)
{
    nn_packet_pool_disable(&conn->pool);
    cch_handle_disable(&conn->handle);
    nn_thread_join(&conn->thread);
}

static void enable_connection(struct lia_node_connection *conn)
{
    cch_handle_enable(&conn->handle);
    nn_packet_pool_enable(&conn->pool);
}

static void data_connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    (void)stream;
    // We will never be here if init_thread() blocks or fails.
    disable_connection_and_wait(conn);
    free_connection_stream(conn);
    if (conn->disconnected) {
        free_connection(conn);
    } else {
        enable_connection(conn);
    }
}

static void start_connection_handler(struct lia_node_connection *conn, u32 mask)
{
    struct nn_packet_stream *stream = conn->stream;
    conn->handler->subscribe(conn->handler, mask);
    stream->packet_callback = discard_packet_callback;
    stream->packet_sent_callback = data_packet_sent_callback;
    stream->packets_sent_callback = data_packets_sent_callback;
    stream->connection_closed_callback = data_connection_closed_callback;
    nn_thread_create(&conn->thread, handler_thread, conn);
}

static void subscribe_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    u32 mask = nn_packet_read_u32(packet);
    nn_packet_stream_return_packet(stream, packet);
    start_connection_handler(conn, mask);
}

static void subscribe_packet_sent_callback(void *userdata, struct nn_packet *packet)
{
    (void)userdata;
    nn_packet_free(packet);
}

static void subscribe_connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    (void)stream;
    free_connection_stream(conn);
    free_connection(conn);
}

static void handle_connection(struct lia_node_connection *conn, struct nn_packet *packet)
{
    struct nn_packet_stream *stream = conn->stream;

    u32 mask = nn_packet_read_u32(packet);
    u64 seek_pos = nn_packet_read_u64(packet);

    al_assert(!conn->ref);
    conn->ref = true;

    // If mask is already set, this is a reconnect. So, always seek to flush
    // the handler. Otherwise and if seek_pos is 0, don't seek as seeking to 0
    // on a new stream can skip data.
    if (mask != 0 || seek_pos > 0) {
        conn->seek_pos = seek_pos;
    }

    if (mask == 0) {
        struct nn_packet *rpacket = nn_packet_create();
        nn_packet_write_u32(rpacket, conn->id);
        nn_packet_write_str(rpacket, cch_entry_get_handler(conn->node->entry));
        conn->handler->write_info(conn->handler, rpacket);
        stream->packet_callback = subscribe_packet_callback;
        stream->packet_sent_callback = subscribe_packet_sent_callback;
        stream->connection_closed_callback = subscribe_connection_closed_callback;
        nn_packet_stream_send_packet(stream, rpacket);
    } else {
        start_connection_handler(conn, mask);
        conn->handler->subscribe(conn->handler, mask);
    }

    nn_packet_stream_return_packet(stream, packet);
}

static void connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_server *server = (struct lia_server *)userdata;
    remove_dormant_connection(server, stream);
    nn_packet_stream_free(stream);
    al_free(stream);
}

static void packet_sent_callback(void *userdata, struct nn_packet *packet)
{
    (void)userdata;
    nn_packet_free(packet);
}

static void demote_and_disconnect_stream(struct lia_server *server, struct nn_packet_stream *stream)
{
    // Discard queue based on the currently set packet_sent_callback.
    // This should always be the expected behavior but here it's important to
    // not lose packets that belong to the packet pool.
    nn_packet_stream_discard_queue(stream);
    stream->userdata = server;
    stream->connection_closed_callback = connection_closed_callback;
    stream->packet_callback = discard_packet_callback;
    stream->packet_sent_callback = packet_sent_callback;
    stream->packets_sent_callback = NULL;
    nn_packet_stream_disconnect(stream);
}

static void signal_callback(void *userdata)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    struct lia_node *node = conn->node;
    struct lia_server *server = node->server;

    nn_thread_join(&conn->thread);
    nn_signal_stop(&conn->signal);
    al_array_remove(node->requests, conn);

    struct nn_packet_stream *stream = conn->stream;
    struct nn_packet *packet = conn->packet;
    conn->packet = NULL;

    if (!packet || conn->errored) {
        conn->handler->free(&conn->handler);
        cch_entry_return_handle(node->entry, &conn->handle);
    }

    if (!packet) {
        // Connection was closed before init was done.
        if (should_free_node(node)) {
            free_node(node);
        }
        nn_packet_stream_free(stream);
        al_free(stream);
        al_free(conn);
    } else if (conn->errored) {
        // We must return the packet before disconnecting.
        nn_packet_stream_return_packet(stream, packet);
        al_free(conn);
        conn = NULL;
        demote_and_disconnect_stream(server, stream);
    } else {
        conn->id = get_incremental_id(server);
        al_array_push(node->connections, conn);
        nn_packet_pool_init(&conn->pool, 1024, server->loop, packet_pool_callback, conn);
        handle_connection(conn, packet);
    }
}

static nn_thread_result NNWT_THREADCALL init_thread(void *userdata)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    nn_thread_set_name("liana_node_init");
    if (!conn->handler->init(conn->handler, &conn->handle)) {
        conn->errored = true;
    }
    nn_signal_send(&conn->signal);
    return 0;
}

static struct lia_node *get_node_from_id(struct lia_server *server, u32 id)
{
    struct lia_node *node;
    al_array_foreach(server->nodes, i, node) {
        if (node->id == id) return node;
    }
    return NULL;
}

static struct lia_node_connection *get_connection_from_id(struct lia_node *node, u32 id)
{
    struct lia_node_connection *conn;
    al_array_foreach(node->connections, i, conn) {
        if (conn->id == id) return conn;
    }
    return NULL;
}

static void pre_init_connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_node_connection *conn = (struct lia_node_connection *)userdata;
    struct nn_packet *packet = conn->packet;
    // Checked in signal_callback and will signal to cleanup the connection.
    conn->packet = NULL;
    nn_packet_stream_return_packet(stream, packet);
}

static void packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct lia_server *server = (struct lia_server *)userdata;

    // We got a packet, this connection is no longer dormant.
    remove_dormant_connection(server, stream);

    u32 node_id = nn_packet_read_u32(packet);
    u32 connection_id = nn_packet_read_u32(packet);

    struct lia_node *node = get_node_from_id(server, node_id);
    if (!node) {
        goto err;
    }

    struct lia_node_connection *conn = NULL;
    if (connection_id == 0) {
        conn = al_alloc_object(struct lia_node_connection);
        stream->userdata = conn;
        conn->node = node;
        conn->stream = stream;
        conn->packet = packet;
        nn_signal_init(&conn->signal, server->loop, signal_callback, conn);
        nn_signal_start(&conn->signal);
        cch_entry_get_handle(node->entry, &conn->handle);
        conn->handler = lia_handler_by_name(cch_entry_get_handler(node->entry))->create_server_handler();
        conn->errored = false;
        conn->disconnected = false;
        conn->ref = false;
        conn->seek_pos = LIANA_TIMESTAMP_INVALID;
        stream->packet_callback = discard_packet_callback;
        stream->connection_closed_callback = pre_init_connection_closed_callback;
        al_array_push(node->requests, conn);
        nn_thread_create(&conn->thread, init_thread, conn);
    } else {
        if ((conn = get_connection_from_id(node, connection_id))) {
            if (conn->ref) {
                // Cleanup the existing connection's handler and demote it's stream.
                // The stream was likely already disconnected client-side but it's still safe
                // to disconnect it here to be sure.
                // An alternative to this could be to create a new connection if conn->ref.
                disable_connection_and_wait(conn);
                demote_and_disconnect_stream(server, conn->stream);
                conn->ref = false;
                enable_connection(conn);
            }
            stream->userdata = conn;
            al_assert(conn->node == node);
            conn->stream = stream;
            handle_connection(conn, packet);
        } else {
            goto err;
        }
    }

    return;
err:
    // Return packet before disconnecting.
    nn_packet_stream_return_packet(stream, packet);
    nn_packet_stream_disconnect(stream);
}

static bool connection_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_server *server = (struct lia_server *)userdata;
    stream->packet_callback = packet_callback;
    stream->packet_sent_callback = packet_sent_callback;
    al_array_push(server->dormant_connections, stream);
    return true;
}

void lia_server_add_stream(struct lia_server *server, struct nn_packet_stream *stream)
{
    stream->connection_callback = connection_callback;
    stream->connection_closed_callback = connection_closed_callback;
    stream->userdata = server;
}

struct lia_node *lia_server_create_node(struct lia_server *server, struct cch_entry *entry)
{
    struct lia_node *node = al_alloc_object(struct lia_node);
    node->id = get_incremental_id(server);
    node->entry = entry;
    al_array_init(node->requests);
    al_array_init(node->connections);
    node->closed = false;
    node->duration = LIANA_TIMESTAMP_INVALID;
    node->server = server;
    al_array_push(server->nodes, node);
    return node;
}

static nn_thread_result NNWT_THREADCALL init_duration_thread(void *userdata)
{
    struct lia_node *node = (struct lia_node *)userdata;
    nn_thread_set_name("liana_init_dur");
    if (!node->handler->init(node->handler, &node->handle)) {
        node->errored = true;
    } else {
        node->duration = node->handler->get_duration(node->handler);
    }
    nn_signal_send(&node->signal);
    return 0;
}

static void duration_signal_callback(void *userdata)
{
    struct lia_node *node = (struct lia_node *)userdata;
    nn_signal_stop(&node->signal);
    nn_thread_join(&node->thread);
    node->handler->free(&node->handler);
    cch_entry_return_handle(node->entry, &node->handle);
    if (should_free_node(node)) {
        free_node(node);
    } else {
        if (node->errored) {
            al_assert(node->duration == LIANA_TIMESTAMP_INVALID);
            node->callback(node->userdata, LIANA_NODE_ERRORED, &node->duration);
        } else {
            node->callback(node->userdata, LIANA_NODE_DURATION, &node->duration);
        }
    }
}

void lia_node_get_duration(struct lia_node *node)
{
    // It is vital we don't block the loop during handler->init().
    nn_signal_init(&node->signal, node->server->loop, duration_signal_callback, node);
    nn_signal_start(&node->signal);
    cch_entry_get_handle(node->entry, &node->handle);
    node->handler = lia_handler_by_name(cch_entry_get_handler(node->entry))->create_server_handler();
    nn_thread_create(&node->thread, init_duration_thread, node);
}

void lia_node_close(struct lia_node *node)
{
    node->closed = true;
    if (should_free_node(node)) {
        free_node(node);
    } else {
        struct lia_node_connection *conn;
        al_array_foreach_rev(node->requests, i, conn) {
            nn_packet_stream_disconnect(conn->stream);
        }
        al_array_foreach_rev(node->connections, i, conn) {
            if (conn->stream) {
                conn->disconnected = true;
                nn_packet_stream_disconnect(conn->stream);
            } else {
                free_connection(conn);
            }
        }
    }
}

void lia_server_close(struct lia_server *server)
{
    struct nn_packet_stream *conn;
    al_array_foreach_rev(server->dormant_connections, i, conn) {
        nn_packet_stream_disconnect(conn);
    }
}

void lia_server_free(struct lia_server *server)
{
    // Assuming we joined on the event loop, server->nodes should be empty.
    al_assert(!server->nodes.count);
    al_array_free(server->nodes);
    al_assert(!server->dormant_connections.count);
    al_array_free(server->dormant_connections);
}