#define AL_LOG_SECTION "liana"
#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_ON_CONNECTION_CLOSED,
    RECONNECT_SIGNAL_CLIENT,
    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);
    lia_vcr_push_packet(&client->vcr, packet);
}

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_name(&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 < 32);
        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] = "subtitle",
    [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 == 0) {
        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;
            if (type == CAMU_STREAM_SUBTITLE && client->mask == 0) {
                log_warn("Ignoring subtitle-only resource.");
                goto out;
            }
            if ((selected & (1 << type)) || !(prefs->enabled & (1 << type))) {
                continue;
            }
            const char *title = NULL;
#ifdef CAMU_HAVE_FFMPEG
            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;
                }
                const AVDictionaryEntry *lang_entry = av_dict_get(metadata, "language", NULL, 0);
                if (lang_entry && !accept_defaults) {
                    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;
                    }
                    break;
                }
            }
#endif
            if (title) {
                log_info("Selected %s stream (index: %u, title: %s).", stream_type_to_str[type], stream->index, title);
            } else {
                log_info("Selected %s stream (index: %u).", stream_type_to_str[type], stream->index);
            }
            selected |= (1 << type);
            struct lia_vcr_track *track = al_alloc_object(struct lia_vcr_track);
            track->stream = stream;
            track->client = lia_handler_by_name(&handler)->create_client_handler();
            track->client->callback = client->callback;
            track->client->userdata = client->userdata;
            if (!track->client->init(track->client, client->renderer, track->stream)) {
                track->client->free(&track->client);
                al_free(track);
                // This will NOT attempt to select another stream.
                continue;
            }
            client->mask |= 1 << stream->index;
            client->callback(client->userdata, LIANA_CLIENT_CONFIGURE, stream, track);
            lia_vcr_add_track(&client->vcr, track);
        }
    }
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);
        nn_packet_stream_disconnect(&client->data);
        return;
    }
    stream->packet_callback = data_packet_callback;
    struct nn_packet *rpacket = nn_packet_create();
    nn_packet_write_u32(rpacket, client->mask);
    nn_packet_stream_send_packet(stream, rpacket);
    lia_vcr_start(&client->vcr);
}

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;
    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->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);
    }
    stream->packet_sent_callback = packet_sent_callback;
    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_u32(packet, client->mask);
    nn_packet_write_u64(packet, client->pos);
    if (client->mask == 0) {
        stream->packet_callback = info_packet_callback;
    } else {
        stream->packet_callback = data_packet_callback;
        lia_vcr_start(&client->vcr);
    }
    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;
    if (client->reconnect == RECONNECT_ON_CONNECTION_CLOSED) {
        lia_vcr_flush(&client->vcr);
    } else {
        lia_vcr_close_all(&client->vcr);
    }
    // If reconnect = SIGNAL_CLIENT, we either never connected or recursed at the reconnect step.
    if (client->reconnect != RECONNECT_SIGNAL_CLIENT) {
        client->rec = (struct lia_reconnect_info){
            .reconnect = client->reconnect == RECONNECT_ON_CONNECTION_CLOSED,
            .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 and should
        // 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 (client->reconnect == RECONNECT_ON_CONNECTION_CLOSED) {
        // If stream_reconnect() errors, the client will be closed on recursion.
        client->reconnect = RECONNECT_SIGNAL_CLIENT;
        if (!client->mask) {
            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->pos = pos;
    client->mask = 0;
    al_array_init(client->streams);
    client->reconnect = RECONNECT_NONE;
    lia_vcr_init(&client->vcr, client->loop, &client->data, node_id);
    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)
{
    if (client->reconnect == RECONNECT_DISCONNECTED) return;
    // If reconnect = ON_CONNECTION_CLOSED 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 (client->reconnect == RECONNECT_NONE) {
        client->reconnect = RECONNECT_ON_CONNECTION_CLOSED;
        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 != RECONNECT_ON_CONNECTION_CLOSED) {
        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);
}