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

#include "../server/common.h"
#ifdef CAMU_HAVE_FFMPEG
#include "../codec/ffmpeg/packet_ext.h"
#endif

#include "client.h"
#include "handlers.h"
#include "list.h"

enum {
    RECONNECT_NONE = 0,
    RECONNECT_RECOVER,
    RECONNECT_SEEK,
    RECONNECT_SIGNAL_CLIENT,
    RECONNECT_DISREGUARD,
    RECONNECT_DISCONNECTED
};

static void data_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct lia_client *client = (struct lia_client *)userdata;
    al_assert(client->vcr.data == stream);
    // The server might cut us off before sending any data, so don't attempt to
    // recover unless we received at least one data packet.
    if (!lia_vcr_push_packet(&client->vcr, packet)) {
        nn_packet_stream_return_packet(stream, packet);
    } else if (client->reconnect == RECONNECT_NONE) {
        client->reconnect = RECONNECT_RECOVER;
        lia_vcr_start(&client->vcr);
    }
}

static s32 stream_compare(const void *a, const void *b)
{
    struct camu_codec_stream *aa = (struct camu_codec_stream *)a;
    struct camu_codec_stream *bb = (struct camu_codec_stream *)b;
    if (aa->type == bb->type) {
        return aa->index - bb->index;
    }
    return aa->type - bb->type;
}

static void collect_streams(struct lia_client *client, struct nn_packet *packet)
{
    u32 count = nn_packet_read_u32(packet);
    for (u32 i = 0; i < count; i++) {
        // Very important zero-initialization.
        struct camu_codec_stream stream = { 0 };
        str codec;
        nn_packet_read_str(packet, &codec);
        stream.codec_info = camu_codec_info_by_id(&codec);
        u8 mode = nn_packet_read_u8(packet);
        u8 type = nn_packet_read_u8(packet);
        u64 duration = nn_packet_read_u64(packet);
        s32 index = nn_packet_read_s32(packet);
        al_assert(index < 64);
        switch (mode) {
        case CAMU_NORMAL: {
            if (type == CAMU_STREAM_AUDIO) {
                struct camu_audio_format *fmt = &stream.audio.fmt;
                fmt->format = nn_packet_read_s32(packet);
                fmt->sample_rate = nn_packet_read_s32(packet);
                fmt->channel_count = nn_packet_read_s32(packet);
#ifdef CAMU_HAVE_FFMPEG
                camu_default_channel_layout(&fmt->channel_layout, fmt->channel_count);
#endif
            } else if (type == CAMU_STREAM_VIDEO) {
                struct camu_video_format *fmt = &stream.video.fmt;
                fmt->width = nn_packet_read_u32(packet);
                fmt->height = nn_packet_read_u32(packet);
                fmt->format = nn_packet_read_s32(packet);
            }
            break;
        }
#ifdef CAMU_HAVE_FFMPEG
        case CAMU_FFMPEG_COMPAT: {
            enum AVCodecID codec_id = nn_packet_read_av_codec_id(packet);
            const AVCodec *av_codec = avcodec_find_decoder(codec_id);
            AVFormatContext *format_context = avformat_alloc_context();
            stream.av.stream = nn_packet_read_av_stream(format_context, av_codec, packet);
            switch (type) {
            case CAMU_STREAM_ATTACHMENT: {
                // Assume all the data we need is in the AVStream object.
                stream.type = CAMU_STREAM_ATTACHMENT;
                client->callback(client->userdata, LIANA_CLIENT_CONFIGURE, &stream, NULL);
                avformat_free_context(format_context);
                continue;
            }
            }
            stream.av.format_context = format_context;
            AVCodecParameters *codecpar = stream.av.stream->codecpar;
            switch (type) {
            case CAMU_STREAM_AUDIO: {
                struct camu_audio_format *fmt = &stream.audio.fmt;
                fmt->format = codecpar->format;
                fmt->sample_rate = codecpar->sample_rate;
                camu_copy_channel_layout(&fmt->channel_layout, &AV_CODECPAR_CHANNEL_LAYOUT(codecpar));
                fmt->channel_count = AV_CODECPAR_CHANNELS(codecpar);
                break;
            }
            case CAMU_STREAM_VIDEO: {
                struct camu_video_format *fmt = &stream.video.fmt;
                fmt->width = (u32)codecpar->width;
                fmt->height = (u32)codecpar->height;
                fmt->format = codecpar->format;
                break;
            }
            }
            break;
        }
#endif
        }
        stream.mode = mode;
        stream.type = type;
        stream.duration = duration;
        stream.index = index;
        if (!stream.codec_info) {
#ifdef CAMU_HAVE_FFMPEG
            if (stream.mode == CAMU_FFMPEG_COMPAT) {
                avformat_free_context(stream.av.format_context);
            }
#endif
            continue;
        }
        al_array_push(client->streams, stream);
    }
    // Video streams have to come before subtitle streams.
    al_array_sort(client->streams, struct camu_codec_stream, stream_compare);
}

static const char *stream_type_to_str[] = {
    [CAMU_STREAM_AUDIO] = "Audio",
    [CAMU_STREAM_VIDEO] = "Video",
    [CAMU_STREAM_SUBTITLE] = "Subtitles",
    [CAMU_STREAM_ATTACHMENT] = "Attachment",
    [CAMU_STREAM_UNKNOWN] = "Unknown"
};

static void parse_info_packet(struct lia_client *client, struct nn_packet *packet)
{
    str handler;
    nn_packet_read_str(packet, &handler);
    client->duration = nn_packet_read_u64(packet);
    collect_streams(client, packet);
    if (!client->streams.count) {
        log_warn("Resource has no streams.");
        goto out;
    }
    struct lia_prefs *prefs = &client->prefs;
    u8 selected = 0;
    u8 accept_defaults = 0;
    for (; accept_defaults < 2; accept_defaults++) {
        struct camu_codec_stream *stream;
        al_array_foreach_ptr(client->streams, i, stream) {
            u8 type = stream->type;
            s32 index = stream->index;
            switch (type) {
            case CAMU_STREAM_AUDIO:
                if (index < prefs->index.audio_min) prefs->index.audio_min = index;
                if (index > prefs->index.audio_max) prefs->index.audio_max = index;
                break;
            case CAMU_STREAM_SUBTITLE:
                if (index < prefs->index.subtitles_min) prefs->index.subtitles_min = index;
                if (index > prefs->index.subtitles_max) prefs->index.subtitles_max = index;
                break;
            }
            if ((selected & (1 << type)) || !(prefs->enabled & (1 << type))) {
                continue;
            }
            const char *title = NULL;
#ifdef CAMU_HAVE_FFMPEG
            bool pref_is_index = false;
            if (!accept_defaults) {
                if (type == CAMU_STREAM_AUDIO && prefs->index.audio >= 0) {
                    if (index != prefs->index.audio) {
                        continue;
                    }
                    pref_is_index = true;
                }
                if (type == CAMU_STREAM_SUBTITLE && prefs->index.subtitles >= 0) {
                    if (index != prefs->index.subtitles) {
                        continue;
                    }
                    pref_is_index = true;
                }
            }
            if (stream->mode == CAMU_FFMPEG_COMPAT) {
                AVDictionary *metadata = stream->av.stream->metadata;
                const AVDictionaryEntry *title_entry = av_dict_get(metadata, "title", NULL, 0);
                if (title_entry) title = title_entry->value;
                if (!pref_is_index && !accept_defaults) {
                    const AVDictionaryEntry *lang_entry = av_dict_get(metadata, "language", NULL, 0);
                    if (lang_entry) {
                        u8 lang = CAMU_LANG_UNKNOWN;
                        str s = al_str_cr(lang_entry->value);
                        if      (al_str_eq(&s, &al_str_c("eng"))) lang = CAMU_LANG_ENGLISH;
                        else if (al_str_eq(&s, &al_str_c("jpn"))) lang = CAMU_LANG_JAPANESE;
                        if ((type == CAMU_STREAM_AUDIO    && lang != prefs->language.audio) ||
                            (type == CAMU_STREAM_SUBTITLE && lang != prefs->language.subtitles)) {
                            continue;
                        }
                    }
                }
            }
#endif
            switch (type) {
            case CAMU_STREAM_AUDIO:
                prefs->index.audio = index;
                break;
            case CAMU_STREAM_VIDEO:
                prefs->index.video = index;
                break;
            case CAMU_STREAM_SUBTITLE:
                prefs->index.subtitles = index;
                break;
            }
            if (title) {
                log_info("%s selected (index: %u, title: %s).", stream_type_to_str[type], index, title);
            } else {
                log_info("%s selected (index: %u).", stream_type_to_str[type], index);
            }
            selected |= (1 << type);
            struct lia_vcr_track *track = al_alloc_object(struct lia_vcr_track);
            track->stream = stream;
            track->handler = lia_handler_by_name(&handler)->create_client_handler();
            track->handler->callback = client->callback;
            track->handler->userdata = client->userdata;
            if (!track->handler->init(track->handler, client->renderer, track->stream)) {
                track->handler->free(&track->handler);
                al_free(track);
                // This will NOT attempt to select another stream.
                continue;
            }
            client->mask |= 1 << index;
            client->callback(client->userdata, LIANA_CLIENT_CONFIGURE, stream, track);
            lia_vcr_add_track(&client->vcr, track);
        }
    }
    if (client->mask == (1 << CAMU_STREAM_SUBTITLE)) {
        log_warn("Ignoring subtitle-only resource.");
        client->mask = 0;
    }
out:
    client->callback(client->userdata, LIANA_CLIENT_CONFIGURE_COMPLETE, NULL, NULL);
}

static void info_packet_callback(void *userdata, struct nn_packet_stream *stream, struct nn_packet *packet)
{
    struct lia_client *client = (struct lia_client *)userdata;
    client->connection_id = nn_packet_read_u32(packet);
    parse_info_packet(client, packet);
    nn_packet_stream_return_packet(stream, packet);
    if (client->mask == 0 || lia_vcr_is_empty(&client->vcr)) {
        log_warn("Discarding resource with no applicable streams.");
        al_assert(client->reconnect == RECONNECT_NONE);
        client->reconnect = RECONNECT_DISREGUARD;
        nn_packet_stream_disconnect(&client->data);
        return;
    }
    stream->packet_callback = data_packet_callback;
    struct nn_packet *rpacket = nn_packet_create();
    nn_packet_write_u64(rpacket, client->mask);
    nn_packet_stream_send_packet(stream, rpacket);
}

static void packet_dequeued_callback(void *userdata, struct nn_packet *packet)
{
    (void)userdata;
    nn_packet_write_size(packet);
}

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

static bool connection_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_client *client = (struct lia_client *)userdata;

    struct nn_socket *sock = &client->data.sock;
    if (sock->type == NNWT_SOCKET_TCP) {
#ifdef AL_LOG_ENABLE_TRACE
        u32 rcvbuf = nn_socket_get_recv_buf(sock);
        u32 sndbuf = nn_socket_get_send_buf(sock);
#endif
        nn_socket_set_recv_buf(sock, MB(2));
        nn_socket_set_send_buf(sock, KB(8));
#ifdef AL_LOG_ENABLE_TRACE
        log_trace("rcvbuf: %u -> %u", rcvbuf, nn_socket_get_recv_buf(sock));
        log_trace("sndbuf: %u -> %u", sndbuf, nn_socket_get_send_buf(sock));
#endif
    }

    if (client->reconnect == RECONNECT_SIGNAL_CLIENT) {
        client->reconnect = RECONNECT_NONE;
        // Even if client_seek() was called before the initial connection_callback(),
        // we still want to call RESUME_AT here.
        struct lia_timing time = {
            .at = client->at,
            .pos = client->pos,
            .pause = LIANA_PAUSE_NONE
        };
        client->at = LIANA_TIMESTAMP_INVALID;
        client->callback(client->userdata, LIANA_CLIENT_RESUME_AT, NULL, &time);
        // The value of client->mask will not have changed since connection_closed_callback().
        if (client->rec.unconfigured) {
            al_assert(client->connection_id == 0);
            log_warn("Handling reconnect on unconfigured client.");
        }
        client->callback(client->userdata, LIANA_CLIENT_RECONNECTED, NULL, &client->rec);
    } else {
        al_assert(client->connection_id == 0 && client->reconnect == RECONNECT_NONE);
    }

    struct nn_packet *packet = nn_packet_create();
    nn_packet_write_u32(packet, client->node_id);
    nn_packet_write_u32(packet, client->connection_id);
    nn_packet_write_u64(packet, client->mask);
    nn_packet_write_u64(packet, client->pos);

    stream->packet_callback = (client->mask == 0) ? info_packet_callback : data_packet_callback;
    stream->packet_dequeued_callback = packet_dequeued_callback;
    stream->packet_sent_callback = packet_sent_callback;

    nn_packet_stream_send_packet(stream, packet);

    return true;
}

static void connection_closed_callback(void *userdata, struct nn_packet_stream *stream)
{
    struct lia_client *client = (struct lia_client *)userdata;

    bool reconnect = client->reconnect == RECONNECT_RECOVER || client->reconnect == RECONNECT_SEEK;
    if (client->reconnect == RECONNECT_RECOVER) {
        al_assert(client->at == LIANA_TIMESTAMP_INVALID);
        struct lia_timing time;
        client->callback(client->userdata, LIANA_CLIENT_RECOVER_TO, NULL, &time);
        client->pos = time.pos;
        client->at = time.at;
    }

    if (reconnect) {
        lia_vcr_flush(&client->vcr);
    } else {
        lia_vcr_close_all(&client->vcr);
    }

    // If reconnect = SIGNAL_CLIENT, we can guarantee that CLIENT_REMOVE_BUFFERS
    // has been called after the most recent connection_callback().
    if (client->reconnect != RECONNECT_SIGNAL_CLIENT) {
        client->rec = (struct lia_reconnect_info){
            .reconnect = reconnect,
            .unconfigured = client->mask == 0,
            .mask = client->mask
        };
        al_array_init(client->rec.detached);
        // CLIENT_REMOVE_BUFFERS should be allowed to run the event loop to wait.
        // It should also attempt to maintain the same state if called consecutively.
        client->callback(client->userdata, LIANA_CLIENT_REMOVE_BUFFERS, NULL, &client->rec);
        client->mask = client->rec.mask;
        struct camu_codec_stream *detached;
        al_array_foreach(client->rec.detached, i, detached) {
            client->mask &= ~(1 << detached->index);
            // We have to remove the track or it will erroneously receive a NULL packet on PACKET_EOF.
            // It would also be wasteful to spin up a track_thread() for a removed track anyway.
            bool removed = lia_vcr_remove_track_by_stream(&client->vcr, detached);
            al_assert(removed);
        }
        al_array_free(client->rec.detached);
    }

    if (reconnect) {
        // If stream_reconnect() errors or is aborted, the client will be closed on recursion.
        client->reconnect = RECONNECT_SIGNAL_CLIENT;
        // It's possible the client was not configured yet. For example, if it was seeked
        // before an info packet. In that case, we don't want to forcefully close it.
        bool was_configured = !client->rec.unconfigured;
        if (was_configured && client->mask == 0) {
            // If mask was emptied by CLIENT_REMOVE_BUFFERS, close the client immediately.
            connection_closed_callback(userdata, stream);
        } else {
#ifdef CAMU_DIRECT_MODE
            nn_multiplex_direct_reconnect(stream);
#else
            nn_packet_stream_reconnect(stream, &client->addr, client->port);
#endif
        }
    } else {
        client->callback(client->userdata, LIANA_CLIENT_CLOSED, NULL, &client->rec);
    }
}

void lia_client_connect(struct lia_client *client, struct nn_event_loop *loop,
    u8 type, str *addr, u16 port, u32 node_id, u64 pos)
{
    client->loop = loop;
    client->node_id = node_id;
    client->duration = LIANA_TIMESTAMP_INVALID;
    lia_vcr_init(&client->vcr, client->loop, &client->data, node_id);
    al_array_init(client->streams);
    client->mask = 0;
    client->pos = pos;
    client->at = LIANA_TIMESTAMP_INVALID;
    client->reconnect = RECONNECT_NONE;
    al_str_clone(&client->addr, addr);
    client->port = port;
    client->connection_id = 0;
    nn_packet_stream_init(&client->data, connection_callback, connection_closed_callback, client);
#ifdef CAMU_DIRECT_MODE
    (void)type;
    nn_multiplex_direct_connect(&client->data, CAMU_MULTIPLEX_LIANA);
#else
    nn_packet_stream_connect(&client->data, client->loop, CAMU_MULTIPLEX_LIANA, type, &client->addr, client->port);
#endif
}

void lia_client_seek(struct lia_client *client, u64 pos, u64 at)
{
    u8 reconnect = client->reconnect;
    if (reconnect == RECONNECT_DISCONNECTED || reconnect == RECONNECT_DISREGUARD) {
        return;
    }
    // If reconnect = SEEK, RECOVER or SIGNAL_CLIENT, we are safe to edit pos
    // and at in-place because they aren't evaluated until connection_callback().
    client->pos = pos;
    client->at = at;
    if (reconnect == RECONNECT_NONE || reconnect == RECONNECT_RECOVER) {
        client->reconnect = RECONNECT_SEEK;
        nn_packet_stream_disconnect(&client->data);
    }
}

void lia_client_disconnect(struct lia_client *client)
{
    u8 reconnect = client->reconnect;
    al_assert(reconnect != RECONNECT_DISCONNECTED);
    client->reconnect = RECONNECT_DISCONNECTED;
    // If reconnect == SIGNAL_CLIENT, stream_disconnect() needs to ensure
    // connection_callback() is never called.
    if (reconnect != RECONNECT_SEEK && reconnect != RECONNECT_DISREGUARD) {
        nn_packet_stream_disconnect(&client->data);
    }
}

void lia_client_free(struct lia_client *client)
{
    lia_vcr_free(&client->vcr);
    nn_packet_stream_free(&client->data);
#ifdef CAMU_HAVE_FFMPEG
    struct camu_codec_stream *stream;
    al_array_foreach_ptr(client->streams, i, stream) {
        if (stream->mode == CAMU_FFMPEG_COMPAT) {
            avformat_free_context(stream->av.format_context);
        }
    }
#endif
    al_array_free(client->streams);
    al_str_free(&client->addr);
}