#include <al/atomic.h>

#include "multiplex.h"
#include "packet_stream.h"

enum {
    PACKET_STREAM_CONNECTING = 0,
    PACKET_STREAM_CONNECTED,
    PACKET_STREAM_DISCONNECTING,
    PACKET_STREAM_DISCONNECTED,
    PACKET_STREAM_DEINIT
};

static inline void init_io_state(struct nn_packet_stream *stream)
{
    stream->corked = true;
    stream->in.packet = nn_packet_create();
    stream->in.have_header = false;
    stream->in.index = 0;
    al_array_init(stream->out.queue);
    stream->out.packet = NULL;
    stream->out.running = false;
}

static void start_write_internal(struct nn_packet_stream *stream)
{
    stream->out.running = true;
    ev_io_start(stream->loop->ev, &stream->wevent);
}

static void stop_write_internal(struct nn_packet_stream *stream)
{
    stream->out.running = false;
    ev_io_stop(stream->loop->ev, &stream->wevent);
}

static void discard_out_queue_internal(struct nn_packet_stream *stream)
{
    // stream->out.running can be true here. The stream will stop the
    // write event by itself if the out queue is empty.
    if (stream->out.packet) {
        stream->packet_sent_callback(stream->userdata, stream->out.packet);
        stream->out.packet = NULL;
    }
    if (stream->packets_sent_callback) {
        stream->packets_sent_callback(stream->userdata,
            stream->out.queue.data, stream->out.queue.count);
    } else {
        struct nn_packet *packet;
        al_array_foreach(stream->out.queue, i, packet) {
            stream->packet_sent_callback(stream->userdata, packet);
        }
    }
    stream->out.queue.count = 0;
}

static void stop_internal(struct nn_packet_stream *stream)
{
    stream->connect = PACKET_STREAM_DISCONNECTED;
    nn_socket_close(&stream->sock);
    // Mark the stream corked so reconnect() can be consistent
    // with connect() and from_socket().
    stream->corked = true;
    nn_packet_reset(stream->in.packet);
    stream->in.have_header = false;
    stream->in.index = 0;
    discard_out_queue_internal(stream);
    stream->connection_closed_callback(stream->userdata, stream);
}

// The simplest control flow requires the event loop to not be making decisions based
// on (inaccurate) estimates of the amount of data available on the socket.
// Therefore, we expect it to only call stream_read/write_callback() once per iteration.
// For write, that means write as much data as possible reguardless of how many packets that would be.
// For read, we have to consider that reading as much data as possible will indefinitely block the loop.
// That could easily cause an unresponsive application. Multiple stacked packet_callback()'s may also
// complicate control flow. For sending large amounts of data, a potential optimization is to batch more
// data into a single "packet" (refering to an nn_packet from this API, not a primitive network packet).
#define EV_REPEAT 0x04

static void stream_read_callback(struct ev_loop *loop, ev_io *w, s32 revents)
{
    (void)loop;
    struct nn_packet_stream *stream = (struct nn_packet_stream *)w->data;
    // EV_READ can mean any of POLLIN, POLLERR, or POLLHUP.
    al_assert(revents & EV_READ); // Assert on EV_ERROR.

read_more: {
    struct nn_buffer *buffer = &stream->in.packet->buffer;
    u8 *ptr = nn_buffer_get_ptr(buffer, stream->in.index);
    u32 size = stream->in.have_header ? nn_packet_get_size(stream->in.packet) : NNWT_PACKET_HEADER_LENGTH;
    ssize_t ret = nn_socket_read(&stream->sock, ptr, size - stream->in.index);
    // Stop on EOF, error or connect = DISCONNECTING.
    if (ret == 0 || (ret < 0 && !nn_socket_eagain(ret)) || stream->connect == PACKET_STREAM_DISCONNECTING) {
        ev_io_stop(stream->loop->ev, &stream->revent);
        if (stream->out.running) {
            stop_write_internal(stream);
        }
        stop_internal(stream);
        return;
    }
    if (ret < 0) {
        // If this callback came from the loop, read/recv() shouldn't return EAGAIN.
        al_assert(revents & EV_REPEAT);
        return;
    }

    stream->in.index += ret;

    al_assert(stream->in.have_header || stream->in.index <= NNWT_PACKET_HEADER_LENGTH);
    if (!stream->in.have_header && stream->in.index == NNWT_PACKET_HEADER_LENGTH) {
        stream->in.have_header = true;
        size = nn_packet_get_size(stream->in.packet);
        buffer->size = size;
        nn_buffer_ensure_space(buffer, size);
        revents |= EV_REPEAT;
        goto read_more;
    }

    if (stream->in.have_header && stream->in.index >= size) {
        struct nn_packet *packet = stream->in.packet;
        stream->in.packet = nn_packet_create();
        stream->in.index = 0;
        stream->in.have_header = false;
        stream->packet_callback(stream->userdata, stream, packet);
        // Stream could be invalid at this point.
    }
}
}

static bool set_stream_connected(struct nn_packet_stream *stream)
{
    stream->connect = PACKET_STREAM_CONNECTED;
    al_assert(stream->corked);
    if (stream->connection_callback(stream->userdata, stream)) {
        stream->corked = false;
    }
    al_assert(stream->packet_callback && stream->packet_sent_callback);
    return !stream->corked;
}

static void stream_write_callback(struct ev_loop *loop, ev_io *w, s32 revents)
{
    (void)loop;
    struct nn_packet_stream *stream = (struct nn_packet_stream *)w->data;
    // EV_WRITE can mean any of POLLOUT, POLLERR, or POLLHUP.
    al_assert(revents & EV_WRITE); // Assert on EV_ERROR.

    if (UNLIKELY(stream->connect == PACKET_STREAM_CONNECTING)) {
        ssize_t ret = nn_socket_write(&stream->sock, &stream->id, 1);
        if (ret < 0) {
            // EV_WRITE may be signaled even if any subsequent write()/send() will error.
            nn_socket_check_error(ret);
            stop_write_internal(stream);
            stop_internal(stream);
            return;
        } else if (ret == 0) {
            // Try again on next POLLOUT, not sure if this can actually happen.
            return;
        }
        if (set_stream_connected(stream)) {
            ev_io_start(stream->loop->ev, &stream->revent);
        }
    }

write_more: {
    if (!stream->out.packet) {
        if (stream->out.queue.count > 0) {
            al_array_pop_at(stream->out.queue, 0, stream->out.packet);
            stream->packet_dequeued_callback(stream->userdata, stream->out.packet);
            stream->out.index = 0;
        } else {
            stop_write_internal(stream);
            return;
        }
    }

    u8 *ptr = nn_buffer_get_ptr(&stream->out.packet->buffer, stream->out.index);
    u32 size = nn_packet_get_size(stream->out.packet);
    ssize_t ret = nn_socket_write(&stream->sock, ptr, size - stream->out.index);
    al_assert(ret != 0); // If this came from the loop, write/send() shouldn't return 0.
    if (ret < 0) { // Don't stop_write_internal() here.
        // On error, if not EAGAIN, expect and wait for stream_read_callback() to stop the stream.
        if (nn_socket_eagain(ret)) al_assert(revents & EV_REPEAT);
        return;
    }

    stream->out.index += ret;

    if (stream->out.index >= size) {
        stream->packet_sent_callback(stream->userdata, stream->out.packet);
        stream->out.packet = NULL;
        revents |= EV_REPEAT;
        goto write_more;
    }
}
}

void nn_packet_stream_init(struct nn_packet_stream *stream,
    bool (*connection_callback)(void *, struct nn_packet_stream *),
    void (*connection_closed_callback)(void *, struct nn_packet_stream *), void *userdata)
{
    init_io_state(stream);
    stream->sock.type = NNWT_SOCKET_INVALID;
    stream->revent.data = stream;
    stream->wevent.data = stream;
    ev_init_n(&stream->revent, stream_read_callback);
    ev_init_n(&stream->wevent, stream_write_callback);
    stream->is_direct = false;
    stream->bridge = NULL;
    stream->connection_callback = connection_callback;
    stream->connection_closed_callback = connection_closed_callback;
    stream->packet_callback = NULL;
    stream->packet_dequeued_callback = NULL;
    stream->packet_sent_callback = NULL;
    stream->packets_sent_callback = NULL;
    stream->userdata = userdata;
    stream->connect = PACKET_STREAM_DISCONNECTED;
}

void nn_packet_stream_from_socket(struct nn_packet_stream *stream, struct nn_event_loop *loop, struct nn_socket *sock)
{
    al_assert(stream->connection_callback && stream->connection_closed_callback);
    stream->loop = loop;
    stream->sock = *sock;
    nn_socket_set_blocking(&stream->sock, false);
    init_io_state(stream);
    stream->revent.data = stream;
    stream->wevent.data = stream;
    s32 fd = nn_socket_get_fd(&stream->sock);
    ev_io_init_n(&stream->revent, stream_read_callback, fd, EV_READ);
    ev_io_init_n(&stream->wevent, stream_write_callback, fd, EV_WRITE);
    if (set_stream_connected(stream)) {
        ev_io_start(stream->loop->ev, &stream->revent);
    }
}

//#define SPORADIC_CONNECTION_FAILURE
#ifdef SPORADIC_CONNECTION_FAILURE
#include <al/random.h>
#endif

static bool do_connect_internal(struct nn_packet_stream *stream, str *addr, u16 port)
{
    bool init = nn_socket_init(&stream->sock, NNWT_SOCKET_NONBLOCKING);
#ifdef SPORADIC_CONNECTION_FAILURE
    bool error = al_random_int(0, 30) == 17;
    if (error || !init || !nn_socket_connect(&stream->sock, addr, port)) {
#else
    if (!init || !nn_socket_connect(&stream->sock, addr, port)) {
#endif
        nn_socket_close(&stream->sock);
        stream->connection_closed_callback(stream->userdata, stream);
        return false;
    }
    s32 fd = nn_socket_get_fd(&stream->sock);
    ev_io_set(&stream->revent, fd, EV_READ);
    ev_io_set(&stream->wevent, fd, EV_WRITE);
    stream->connect = PACKET_STREAM_CONNECTING;
    start_write_internal(stream);
    return true;
}

bool nn_packet_stream_connect(struct nn_packet_stream *stream, struct nn_event_loop *loop,
    u8 id, u8 type, str *addr, u16 port)
{
    stream->id = id;
    stream->loop = loop;
    stream->sock.type = type;
    return do_connect_internal(stream, addr, port);
}

// Keep in mind this stream will have had it's state reset in stop_internal().
bool nn_packet_stream_reconnect(struct nn_packet_stream *stream, str *addr, u16 port)
{
    if (stream->connect == PACKET_STREAM_CONNECTING) {
        stop_write_internal(stream);
        stop_internal(stream);
    }
    al_assert(stream->connect == PACKET_STREAM_DISCONNECTED);
    return do_connect_internal(stream, addr, port);
}

bool nn_packet_stream_set_connected(struct nn_packet_stream *stream)
{
    return set_stream_connected(stream);
}

void nn_packet_stream_cork(struct nn_packet_stream *stream, bool cork)
{
    if (stream->is_direct) return;
    if (stream->connect == PACKET_STREAM_CONNECTED) {
        if (cork && !stream->corked) {
            ev_io_stop(stream->loop->ev, &stream->revent);
        } else if (!cork && stream->corked) {
            ev_io_start(stream->loop->ev, &stream->revent);
        }
        stream->corked = cork;
    }
}

bool nn_packet_stream_send_packet(struct nn_packet_stream *stream, struct nn_packet *packet)
{
    if (stream->is_direct) {
        stream->packet_dequeued_callback(stream->userdata, packet);
        struct nn_packet_stream *direct = atomic_load(void)(&stream->direct, AL_ATOMIC_RELAXED);
        direct->packet_callback(direct->userdata, direct, packet);
        return true;
    }
    if (stream->connect == PACKET_STREAM_DISCONNECTING) {
        return false;
    }
    al_assert(stream->connect == PACKET_STREAM_CONNECTED);
    al_array_push(stream->out.queue, packet);
    if (!stream->out.running) {
        start_write_internal(stream);
    }
    return true;
}

void nn_packet_stream_discard_queue(struct nn_packet_stream *stream)
{
    discard_out_queue_internal(stream);
}

void nn_packet_stream_return_packet(struct nn_packet_stream *stream, struct nn_packet *packet)
{
    if (stream->is_direct) {
        struct nn_packet_stream *direct = atomic_load(void)(&stream->direct, AL_ATOMIC_RELAXED);
        direct->packet_sent_callback(direct->userdata, packet);
    } else {
        nn_packet_free(packet);
    }
}

void nn_packet_stream_return_packets(struct nn_packet_stream *stream, struct nn_packet **packets, u32 count)
{
    if (stream->is_direct) {
        struct nn_packet_stream *direct = atomic_load(void)(&stream->direct, AL_ATOMIC_RELAXED);
        if (direct->packets_sent_callback) {
            direct->packets_sent_callback(direct->userdata, packets, count);
        } else {
            for (u32 i = 0; i < count; i++) {
                direct->packet_sent_callback(direct->userdata, packets[i]);
            }
        }
    } else {
        for (u32 i = 0; i < count; i++) {
            nn_packet_free(packets[i]);
        }
    }
}

static void bridge_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    (void)userdata;
    al_assert(stream->is_direct);
    struct nn_packet_stream *direct = stream->direct;
    direct->packet_sent_callback(direct->userdata, packet);
}

static inline void install_bridge(struct nn_packet_stream *stream, struct nn_packet_stream *direct, struct nn_packet_stream *bridge)
{
    bridge->is_direct = true;
    bridge->direct = stream;
    // Update all callbacks that could happen post-connection because
    // direct->connection_closed_callback() could run the event loop after
    // we update stream->direct to bridge.
    bridge->packet_callback = bridge_packet_callback;
    bridge->packet_dequeued_callback = direct->packet_dequeued_callback;
    bridge->packet_sent_callback = direct->packet_sent_callback;
    bridge->packets_sent_callback = direct->packets_sent_callback;
    bridge->userdata = direct->userdata;
    // nn_packet_stream_return_packet() needs to be thread-safe. Non-atomically updating
    // the value of stream->direct here would break that.
    atomic_store(void)(&stream->direct, bridge, AL_ATOMIC_RELEASE);
}

void nn_packet_stream_disconnect(struct nn_packet_stream *stream)
{
    al_assert(stream->connect != PACKET_STREAM_DISCONNECTED);
    al_assert(stream->connect != PACKET_STREAM_DISCONNECTING);
    if (stream->is_direct) {
        struct nn_packet_stream *direct = atomic_load(void)(&stream->direct, AL_ATOMIC_ACQUIRE);
        // After direct->connection_closed_callback(), direct could be freed.
        // So use the bridge to finish the disconnect() on this side of the stream.
        struct nn_multiplex_bridge *bridge = NULL;
        if (stream->bridge) {
            bridge = stream->bridge;
            stream->bridge = NULL;
        } else if (direct->bridge) {
            bridge = direct->bridge;
            direct->bridge = NULL;
        }
        if (!bridge) {
            return; // Disconnect already in progress.
        }
        direct->connect = PACKET_STREAM_DISCONNECTED;
        stream->connect = PACKET_STREAM_DISCONNECTED;
        direct->corked = true;
        stream->corked = true;
        install_bridge(stream, direct, bridge->bridge);
        direct->connection_closed_callback(direct->userdata, direct);
        stream->connection_closed_callback(stream->userdata, stream);
        bridge->broken = true; // See multiplex.c:direct_connect().
        if (!bridge->queue.count) {
            nn_multiplex_bridge_free(bridge);
        }
        return;
    }
    if (stream->connect == PACKET_STREAM_CONNECTING) {
        // wevent is always active when connect = CONNECTING.
        stop_write_internal(stream);
        stop_internal(stream);
    } else if (stream->corked) {
        stop_internal(stream);
    } else {
        // This relies on stream_read_callback() to resolve and call stop_internal().
        // Note that read() can still return data after shutdown(). This is why we have to
        // check if connect = DISCONNECTING in send_packet().
        stream->connect = PACKET_STREAM_DISCONNECTING;
        // This assumes select()/poll() will return after shutdown().
        nn_socket_shutdown(&stream->sock);
    }
}

void nn_packet_stream_free(struct nn_packet_stream *stream)
{
    al_assert(stream->connect == PACKET_STREAM_DISCONNECTED);
    stream->connect = PACKET_STREAM_DEINIT;
    al_assert(stream->in.packet);
    // Loosely assert that stop_internal() has run.
    al_assert(!stream->out.running);
    nn_packet_free(stream->in.packet);
    al_array_free(stream->out.queue);
}