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

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

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);
}

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.

    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);
    if (ret <= 0 || 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;
    }

    stream->in.index += ret;

    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);
    }

    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);
        }
    }

    if (!stream->out.packet) {
        if (stream->out.queue.count > 0) {
            al_array_pop_at(stream->out.queue, 0, 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);
    if (ret <= 0) {
        stop_write_internal(stream);
        return;
    }

    stream->out.index += ret;

    if (stream->out.index >= size) {
        stream->packet_sent_callback(stream->userdata, stream->out.packet);
        stream->out.packet = NULL;
    }
}

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->revent.data = stream;
    stream->wevent.data = stream;
    ev_init_n(&stream->revent, stream_read_callback);
    ev_init_n(&stream->wevent, stream_write_callback);
    stream->direct = NULL;
    stream->connection_callback = connection_callback;
    stream->connection_closed_callback = connection_closed_callback;
    stream->packet_callback = NULL;
    stream->packet_sent_callback = NULL;
    stream->packets_sent_callback = NULL;
    stream->userdata = userdata;
    stream->connect = PACKET_STREAM_DISCONNECTED;
}

void nn_packet_stream_set_nodelay(struct nn_packet_stream *stream, s32 nodelay)
{
    nn_socket_set_nodelay(&stream->sock, nodelay);
}

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);
    }
}

static void do_connect_internal(struct nn_packet_stream *stream, str *addr, u16 port)
{
    if (!nn_socket_init(&stream->sock, NNWT_SOCKET_NONBLOCKING)) {
        stream->connection_closed_callback(stream->userdata, stream);
        return;
    }
    if (!nn_socket_connect(&stream->sock, addr, port)) {
        nn_socket_close(&stream->sock);
        stream->connection_closed_callback(stream->userdata, stream);
        return;
    }
    s32 fd = nn_socket_get_fd(&stream->sock);
    ev_io_set(&stream->revent, fd, EV_READ);
    ev_io_set(&stream->wevent, fd, EV_WRITE);
    // Start non-blocking connection.
    stream->connect = PACKET_STREAM_CONNECTING;
    start_write_internal(stream);
}

void 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;
    do_connect_internal(stream, addr, port);
}

// Keep in mind this stream will have had it's state reset in stop_internal().
void 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);
    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->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->connect == PACKET_STREAM_DISCONNECTING) {
        return false;
    }
    al_assert(stream->connect == PACKET_STREAM_CONNECTED);
    nn_packet_write_size(packet);
    if (stream->direct) {
        struct nn_packet_stream *direct = stream->direct;
        direct->packet_callback(direct->userdata, direct, packet);
        return true;
    }
    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->direct) {
        struct nn_packet_stream *direct = stream->direct;
        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->direct) {
        struct nn_packet_stream *direct = stream->direct;
        if (direct->packets_sent_callback) {
            direct->packets_sent_callback(direct->userdata, packets, count);
        } else {
            for (u32 i = 0; i < count; i++) {
                if (packets[i]) {
                    direct->packet_sent_callback(direct->userdata, packets[i]);
                }
            }
        }
    } else {
        for (u32 i = 0; i < count; i++) {
            if (packets[i]) nn_packet_free(packets[i]);
        }
    }
}

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->direct) {
        struct nn_packet_stream *direct = stream->direct;
        direct->connect = PACKET_STREAM_DISCONNECTED;
        stream->connect = PACKET_STREAM_DISCONNECTED;
        direct->corked = true;
        stream->corked = true;
        struct nn_packet_stream *bridge = nn_multiplex_direct_get_bridge();
        bridge->connection_callback = NULL;
        bridge->connection_closed_callback = NULL;
        bridge->packet_sent_callback = direct->packet_sent_callback;
        bridge->packets_sent_callback = direct->packets_sent_callback;
        bridge->userdata = direct->userdata;
        direct->connection_closed_callback(direct->userdata, direct);
        stream->direct = bridge;
        stream->connection_closed_callback(stream->userdata, stream);
        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->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);
}