#define AL_LOG_SECTION "ff_demuxer"
//#define AL_LOG_ENABLE_TRACE
#include <al/log.h>
#include <al/lib.h>

#include "demuxer.h"
#include "avio.h"

static void close_internal(struct camu_ff_demuxer *av)
{
    if (av->io_context) {
        av_free(av->io_context->buffer);
#if LIBAVFORMAT_VERSION_INT >= AV_VERSION_INT(57, 80, 100)
        avio_context_free(&av->io_context);
#else
        av_free(av->io_context);
        av->io_context = NULL;
#endif
    }
    if (av->format_context) {
        avformat_flush(av->format_context);
        avformat_close_input(&av->format_context);
        // avformat_close_input() should call avformat_free_context().
    }
}

static bool ff_demuxer_init(struct camu_demuxer *demux, struct cch_handle *handle)
{
    struct camu_ff_demuxer *av = (struct camu_ff_demuxer *)demux;

    av->format_context = avformat_alloc_context();
    if (!av->format_context) {
        log_error("Failed to create format context.");
        return false;
    }

    // This is meant to be the "block size" so, if I understand correctly, that would mean
    // the smallest amount for 1 read. Therefore, if reading from memory, it would be a page.
    u8 *buf = (u8 *)av_malloc(al_page_size);
    av->io_context = avio_alloc_context(buf, al_page_size, 0, handle, camu_avio_read, NULL, camu_avio_seek);
    if (!av->io_context) {
        log_error("Failed to create custom io context.");
        return false;
    }

    if (!cch_handle_can_seek(handle)) {
        // AVIO_SEEKABLE_NORMAL should always be unset if we know we cannot seek.
        // However, it should also never be unset if we *do* want to seek because it
        // disables long seeks (Not tested in a long time).
        // This will also make some resources (e.g. .opus) fail to determine a duration.
        av->io_context->seekable &= ~AVIO_SEEKABLE_NORMAL;
    } else {
        av->io_context->seekable |= AVIO_SEEKABLE_NORMAL;
    }

    av->format_context->pb = av->io_context;
    av->format_context->flags |= AVFMT_FLAG_CUSTOM_IO;

    AVDictionary *opts = NULL;
    av_dict_set(&opts, "analyzeduration", "10000000", 0);
    av_dict_set(&opts, "probesize", "20M", 0);

    // FFmpeg says the avio protocol configures "buffers and access patterns".
    // I think we already do enough buffering but, better access patterns for
    // network streams sounds like something we want.
    s32 ret = avformat_open_input(&av->format_context, "cache:", NULL, &opts);
    if (ret < 0) {
        log_error("Failed to open input (%s).", av_err2str(ret));
        av_dict_free(&opts);
        return false;
    }

    av_dict_free(&opts);

    ret = avformat_find_stream_info(av->format_context, NULL);
    if (ret < 0) {
        log_error("Failed to find stream info (%s).", av_err2str(ret));
        return false;
    }

    if (!av->format_context->nb_streams) {
        log_error("No streams found.");
        return false;
    }

    s64 default_duration = (av->format_context->duration < 0) ? 0 : av->format_context->duration;
    av->duration = 0;

    // @TODO:
    //  - 1 frame .gif.
    //  - No-duration video. (ex: data/test_resources/no_duration.mkv)
    //  - Incomplete video files with no duration. .temp.webm.
#define GUESS_STREAM_IS_IMAGE(stream) \
    ((stream->duration >= 0 && stream->nb_frames <= 1 && \
      (stream->avg_frame_rate.den == 0 && stream->r_frame_rate.den > 0)) || \
      (stream->duration < 0 && default_duration == 0))

    for (u32 i = 0; i < av->format_context->nb_streams; i++) {
        AVStream *stream = av->format_context->streams[i];
#ifdef CAMU_OLD_FFMPEG
        // This is probably not correct.
        if (stream->codecpar->channel_layout == 0) {
            stream->codecpar->channel_layout = AV_CH_LAYOUT_STEREO;
            stream->codecpar->channels = 2;
        }
#endif
        enum AVMediaType type = stream->codecpar->codec_type;
        u64 duration = 0;
        u8 mapped_type;
        if (type == AVMEDIA_TYPE_AUDIO || type == AVMEDIA_TYPE_VIDEO || type == AVMEDIA_TYPE_SUBTITLE) {
            // Try to detect attached images.
            if (type == AVMEDIA_TYPE_VIDEO && GUESS_STREAM_IS_IMAGE(stream)) {
                log_debug("Assuming stream #%u is an image.", i);
            } else {
                if (stream->duration < 0) {
                    stream->duration = av_rescale_q(default_duration, AV_TIME_BASE_Q, stream->time_base);
                    al_assert(stream->duration >= 0);
                    log_debug("Stream #%u has an invalid duration, defaulting to %.3fs.",
                        i, stream->duration * av_q2d(stream->time_base));
                }
                duration = (u64)av_rescale_q(stream->duration, stream->time_base, AV_TIME_BASE_Q);
                if (duration > av->duration) {
                    av->duration = duration;
                }
            }
            if (stream->start_time == AV_NOPTS_VALUE) {
                stream->start_time = 0;
            }
            mapped_type = type;
        } else if (type == AVMEDIA_TYPE_ATTACHMENT) {
            mapped_type = CAMU_STREAM_ATTACHMENT;
        } else {
            // CAMU_STREAM_UNKNOWN(7) is a different value than AVMEDIA_TYPE_UNKNOWN(-1) because
            // we don't support negative stream types.
            mapped_type = CAMU_STREAM_UNKNOWN;
        }
        al_array_push(av->demux.streams, ((struct camu_codec_stream){
            .mode = CAMU_FFMPEG_COMPAT,
            .type = mapped_type,
            .duration = duration,
            .index = stream->index,
            .av.stream = stream
        }));
    }

#if LIBAVFORMAT_VERSION_INT < AV_VERSION_INT(60, 30, 100)
    av_format_inject_global_side_data(av->format_context);
#endif

#ifdef AL_LOG_ENABLE_TRACE
    av_dump_format(av->format_context, 0, "", 0);
#endif

    av->eof = false;

    return true;
}

static bool subscribed_to_index(struct camu_demuxer *demux, s32 index)
{
    return (demux->subscribed >> index) & 1;
}

static s32 ff_demuxer_get_packet(struct camu_demuxer *demux, struct camu_codec_packet *packet)
{
    struct camu_ff_demuxer *av = (struct camu_ff_demuxer *)demux;
    packet->mode = CAMU_FFMPEG_COMPAT;
    AVPacket *pkt = packet->av.pkt;
    s32 ret = 0;
    for (;;) {
        ret = av_read_frame(av->format_context, pkt);
        if (ret == AVERROR_EOF) {
            av->eof = true;
            break;
        }
        if (ret < 0) {
            log_error("Failed to read frame (%s).", av_err2str(ret));
            break;
        }
        if (!subscribed_to_index(demux, pkt->stream_index)) {
            av_packet_unref(pkt);
            continue;
        }
        break;
    }
    return ret;
}

static u64 ff_demuxer_get_duration(struct camu_demuxer *demux)
{
    struct camu_ff_demuxer *av = (struct camu_ff_demuxer *)demux;
    return av->duration;
}

static bool ff_demuxer_seek(struct camu_demuxer *demux, u64 pos, bool bytes)
{
    struct camu_ff_demuxer *av = (struct camu_ff_demuxer *)demux;
    // avio_flush() here probably does nothing.
    avio_flush(av->io_context);
    avformat_flush(av->format_context);
    if (av->eof) av->eof = false;
    al_assert(pos <= INT64_MAX);
    al_assert(!(bytes && av->format_context->flags & AVFMT_NO_BYTE_SEEK));
    if (!bytes && pos > av->duration) pos = av->duration;
    s64 ts = (s64)pos;
    s32 flags = bytes ? AVSEEK_FLAG_BYTE : 0;
    s32 ret = avformat_seek_file(av->format_context, -1, INT64_MIN, ts, ts, flags);
    if (ret < 0) {
        log_error("Failed to seek resource (%s).", av_err2str(ret));
        return false;
    }
    return true;
}

static void ff_demuxer_free(struct camu_demuxer **demux)
{
    struct camu_ff_demuxer *av = (struct camu_ff_demuxer *)*demux;
    close_internal(av);
    al_array_free(av->demux.streams);
    al_free(av);
    *demux = NULL;
}

struct camu_demuxer *camu_ff_demuxer_create(void)
{
    struct camu_ff_demuxer *av = al_alloc_object(struct camu_ff_demuxer);
    av->demux.mode = CAMU_FFMPEG_COMPAT;
    al_array_init(av->demux.streams);
    av->demux.init = ff_demuxer_init;
    av->demux.get_packet = ff_demuxer_get_packet;
    av->demux.get_duration = ff_demuxer_get_duration;
    av->demux.seek = ff_demuxer_seek;
    av->demux.free = ff_demuxer_free;
    return (struct camu_demuxer *)av;
}