#include "rpc2.h"

bool nn_rpc_init(struct nn_rpc *rpc, struct nn_event_loop *loop,
    void (*connection_callback)(void *, struct nn_rpc_connection *),
    void (*connection_closed_callback)(void *, struct nn_rpc_connection *), void *userdata)
{
    rpc->loop = loop;
    rpc->increment = 0;
    al_array_init(rpc->commands);
    rpc->conn = NULL;
    al_array_init(rpc->connections);
    rpc->connection_callback = connection_callback;
    rpc->connection_closed_callback = connection_closed_callback;
    rpc->userdata = userdata;
    return true;
}

void nn_rpc_add_command(struct nn_rpc *rpc, struct nn_rpc_command *command)
{
    al_array_push(rpc->commands, *command);
}

static void packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct nn_rpc_connection *conn = (struct nn_rpc_connection *)userdata;
    al_assert(conn->stream == stream);
    u32 id = nn_packet_read_u32(packet);
    s8 op = nn_packet_read_s8(packet);
    if (op == -1) { // Response
        struct nn_rpc_callback *callback;
        al_array_foreach_ptr(conn->callbacks, i, callback) {
            if (callback->id == id) {
                callback->callback(callback->userdata, conn, packet);
                al_array_remove_at(conn->callbacks, i);
                return;
            }
        }
    } else { // Command
        struct nn_rpc_command *command;
        al_array_foreach_ptr(conn->rpc->commands, i, command) {
            if (command->op == op) {
                struct nn_packet *rpacket = nn_packet_create();
                nn_packet_write_u32(rpacket, id);
                nn_packet_write_s8(rpacket, -1);
                if (command->callback(command->userdata, conn, packet, rpacket)) {
                    conn->outgoing++;
                    nn_packet_stream_send_packet(stream, rpacket);
                } else {
                    nn_packet_free(rpacket);
                }
                return;
            }
        }
    }
    nn_packet_stream_return_packet(stream, packet);
}

static void packet_sent_callback(void *userdata, struct nn_packet *packet)
{
    struct nn_rpc_connection *conn = (struct nn_rpc_connection *)userdata;
    nn_packet_free(packet);
    al_assert(conn->outgoing > 0);
    conn->outgoing--;
    if (conn->flushing && conn->outgoing == 0) {
        nn_packet_stream_disconnect(conn->stream);
    }
}

static bool stream_connection_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct nn_rpc_connection *conn = (struct nn_rpc_connection *)userdata;
    struct nn_rpc *rpc = conn->rpc;
    if (stream->sock.type == NNWT_SOCKET_TCP) {
        nn_packet_stream_set_nodelay(stream, 1);
    }
    stream->packet_callback = packet_callback;
    stream->packet_sent_callback = packet_sent_callback;
    stream->userdata = conn;
    rpc->connection_callback(rpc->userdata, conn);
    // Only add connections when acting as a server.
    if (!rpc->conn) al_array_push(rpc->connections, conn);
    return true;
}

static void stream_connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct nn_rpc_connection *conn = (struct nn_rpc_connection *)userdata;
    al_assert(stream == conn->stream);
    conn->rpc->connection_closed_callback(conn->rpc->userdata, conn);
}

static inline void init_rpc_connection(struct nn_rpc *rpc, struct nn_rpc_connection *conn)
{
    conn->rpc = rpc;
    al_array_init(conn->callbacks);
    conn->outgoing = 0;
    conn->flushing = false;
}

void nn_rpc_add_stream(struct nn_rpc *rpc, struct nn_packet_stream *stream)
{
    struct nn_rpc_connection *conn = al_alloc_object(struct nn_rpc_connection);
    init_rpc_connection(rpc, conn);
    conn->stream = stream;
    conn->stream->connection_callback = stream_connection_callback;
    conn->stream->connection_closed_callback = stream_connection_closed_callback;
    conn->stream->userdata = conn;
}

struct nn_rpc_connection *nn_rpc_prepare_client(struct nn_rpc *rpc)
{
    al_assert(!rpc->conn);
    struct nn_rpc_connection *conn = al_alloc_object(struct nn_rpc_connection);
    init_rpc_connection(rpc, conn);
    conn->stream = al_alloc_object(struct nn_packet_stream);
    nn_packet_stream_init(conn->stream, stream_connection_callback, stream_connection_closed_callback, conn);
    rpc->conn = conn;
    al_array_push(rpc->connections, conn);
    return rpc->conn;
}

void nn_rpc_connect(struct nn_rpc *rpc, u8 id, u8 type, str *addr, u16 port)
{
    al_assert(rpc->conn);
    nn_packet_stream_connect(rpc->conn->stream, rpc->loop, id, type, addr, port);
}

void nn_rpc_reconnect(struct nn_rpc *rpc, str *addr, u16 port)
{
    al_assert(rpc->conn);
    nn_packet_stream_reconnect(rpc->conn->stream, addr, port);
}

// @TODO: Pack opcode into a u32, reduce ID by 8 bits (make struct with bitmask)
struct nn_packet *nn_rpc_get_packet(struct nn_rpc *rpc, s8 op)
{
    struct nn_packet *packet = nn_packet_create();
    nn_packet_write_u32(packet, (rpc->increment = al_u32_inc_wrap(rpc->increment)));
    nn_packet_write_s8(packet, op);
    return packet;
}

void nn_rpc_free(struct nn_rpc *rpc)
{
    struct nn_rpc_connection *conn;
    al_array_foreach(rpc->connections, i, conn) {
        nn_packet_stream_free(conn->stream);
        al_array_free(conn->callbacks);
        al_free(conn->stream);
        al_free(conn);
    }
    al_array_free(rpc->connections);
    al_array_free(rpc->commands);
}

void nn_rpc_connection_command(struct nn_rpc_connection *conn, struct nn_packet *packet,
    void (*callback)(void *, struct nn_rpc_connection *conn, struct nn_packet *), void *userdata)
{
    if (callback) {
        al_array_push(conn->callbacks, ((struct nn_rpc_callback){
            .id = nn_packet_get_u32(packet, NNWT_PACKET_HEADER_LENGTH),
            .callback = callback,
            .userdata = userdata
        }));
    }
    conn->outgoing++;
    nn_packet_stream_send_packet(conn->stream, packet);
}

void nn_rpc_conn_flush(struct nn_rpc_connection *conn)
{
    if (conn->outgoing != 0) {
        conn->flushing = true;
    } else {
        nn_packet_stream_disconnect(conn->stream);
    }
}

void nn_rpc_conn_disconnect(struct nn_rpc_connection *conn)
{
    nn_packet_stream_disconnect(conn->stream);
}