#define AL_LOG_SECTION "ff_decoder"
#include <al/log.h>
#include <al/lib.h>
#include <libavutil/cpu.h>

#include "decoder.h"

#ifdef CAMU_FF_DECODER_HWACCEL
#include "../../render/renderer.h"
static const char *hwdevices[] = {
#ifdef CAMU_HAVE_VULKAN_DECODE
    "vulkan",
#endif
#ifdef CAMU_HAVE_VAAPI
    "vaapi",
#endif
#ifdef CAMU_HAVE_D3D11VA
    "d3d11va",
#endif
#ifdef CAMU_HAVE_DXVA2
    "dxva2",
#endif
};
#endif

static bool alloc_codec_context_internal(struct camu_ff_decoder *av, const AVCodec *codec, AVCodecParameters *codecpar)
{
    if (!(av->codec_context = avcodec_alloc_context3(codec))) {
        log_error("Failed to alloc codec context.");
        return false;
    }

    if (avcodec_parameters_to_context(av->codec_context, codecpar) < 0) {
        log_error("Failed to copy codec parameters.");
        return false;
    }

    return true;
}

static bool open_avcodec_internal(struct camu_ff_decoder *av, const AVCodec *codec)
{
    AVDictionary *opts = NULL;
    s32 ret = avcodec_open2(av->codec_context, codec, &opts);
    if (ret < 0) {
        log_error("Failed to open codec %s (%s).", codec->name, av_err2str(ret));
        return false;
    }
    av_dict_free(&opts);
    return true;
}

#ifdef CAMU_FF_DECODER_HWACCEL
static s32 get_buffer2(AVCodecContext *context, AVFrame *pic, s32 flags)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)context->opaque;
    context->opaque = av->renderer->opaque;
    s32 ret = av->renderer->get_buffer2(context, pic, flags);
    context->opaque = av;
    return ret;
}

// This is currently untested.
static s32 init_hwframe_context(struct camu_ff_decoder *av, AVCodecContext *context, AVBufferRef *hw_device_context)
{
    AVBufferRef *hw_frames_ref;
    if (!(hw_frames_ref = av_hwframe_ctx_alloc(hw_device_context))) {
        log_error("Failed to create hardware frame context.");
        return -1;
    }

    AVHWFramesContext *frames_ctx = (AVHWFramesContext *)(hw_frames_ref->data);
    frames_ctx->format = av->hw_pix_fmt;
    // codec_context->sw_pix_fmt is only set in ffmpeg before calling get_hw_format().
    frames_ctx->sw_format = av->codec_context->sw_pix_fmt;
    frames_ctx->width = av->codec_context->width;
    frames_ctx->height = av->codec_context->height;
    frames_ctx->initial_pool_size = 12;

    s32 ret = av_hwframe_ctx_init(hw_frames_ref);
    if (ret < 0) {
        log_error("Failed to initialize hardware frame context (%s).", av_err2str(ret));
        av_buffer_unref(&hw_frames_ref);
        return ret;
    }

    context->hw_frames_ctx = hw_frames_ref;
    /*
    context->hw_frames_ctx = av_buffer_ref(hw_frames_ref);
    av_buffer_unref(&hw_frames_ref);
    */

    return ret;
}

static enum AVPixelFormat get_hw_format(AVCodecContext *context, const enum AVPixelFormat *pix_fmts)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)context->opaque;

    const AVCodec *codec = av->codec_context->codec;
    const char *hwdevice_name = av_hwdevice_get_type_name(av->hw_device_type);

    const enum AVPixelFormat *fmt = pix_fmts;
    for (; *fmt != AV_PIX_FMT_NONE; fmt++) {
        if (*fmt == av->hw_pix_fmt) {
            if (av->use_frames_context) {
                init_hwframe_context(av, av->codec_context, av->hw_context);
            }
            log_info("Using %s hardware decoding.", hwdevice_name);
            goto out;
        }
    }

    log_error("Failed to get %s HW surface format.", hwdevice_name);

    // Remake AVCodecContext as a software decoder.
    AVCodecParameters *codecpar = av->codecpar;
    av->errored_hw_context = av->codec_context;
    alloc_codec_context_internal(av, codec, codecpar);
    // Note about thread_type choice in ff_decoder_init().
    av->codec_context->thread_type = FF_THREAD_FRAME;
    av->codec_context->thread_count = av->thread_count;
    if (open_avcodec_internal(av, codec)) {
        log_info("Software fallback, using %i threads for decoder.", av->thread_count);
    }

out:
    return *fmt;
}

static s32 init_hwdevice_context(struct camu_ff_decoder *av, AVCodecContext *context)
{
    if (av->hw_device_type == AV_HWDEVICE_TYPE_D3D11VA) {
        log_info("av_hwdevice_ctx_create().");
    }
    s32 ret = av_hwdevice_ctx_create(&av->hw_context, av->hw_device_type, NULL, NULL, 0);
    if (av->hw_device_type == AV_HWDEVICE_TYPE_D3D11VA) {
        log_info("av_hwdevice_ctx_create() returns.");
    }
    if (ret < 0) {
        log_error("Failed to create specified HW device.");
        return ret;
    }

    context->hw_device_ctx = av_buffer_ref(av->hw_context);
    // Note that context->extra_hw_frames has the ability to cause corruption.
    context->extra_hw_frames = 12;

    return ret;
}

/*
static bool hw_device_supported_by_type(struct camu_ff_decoder *av, enum AVHWDeviceType type)
{
    enum AVHWDeviceType supported_type;
    al_array_foreach(av->supported_hw_devices, i, supported_type) {
        if (supported_type == type) return true;
    }
    return false;
}
*/

static enum AVHWDeviceType hw_device_supported_by_name(struct camu_ff_decoder *av, const char *name, size_t length)
{
    enum AVHWDeviceType supported_type;
    al_array_foreach(av->supported_hw_devices, i, supported_type) {
        if (al_strncmp(av_hwdevice_get_type_name(supported_type), name, length) == 0) {
            return supported_type;
        }
    }
    return AV_HWDEVICE_TYPE_NONE;
}

static bool hw_config_needs_frames_ctx(const AVCodecHWConfig *config)
{
    return (config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_FRAMES_CTX) &&
        !(config->methods & AV_CODEC_HW_CONFIG_METHOD_HW_DEVICE_CTX);
}

static bool get_hwdevice_config(struct camu_ff_decoder *av, const AVCodec *codec, enum AVHWDeviceType hw_device_type)
{
    s32 iter = 0;
    const AVCodecHWConfig *config = NULL;
    while ((config = avcodec_get_hw_config(codec, iter++))) {
        if (config->device_type == hw_device_type) {
            av->hw_pix_fmt = config->pix_fmt;
            av->use_frames_context = hw_config_needs_frames_ctx(config);
            return true;
        }
    }
    const char *hwdevice_name = av_hwdevice_get_type_name(hw_device_type);
    log_warn("No hw_config for device type %s (type: %d).", hwdevice_name, hw_device_type);
    return false;
}

static void collect_supported_hwaccels(struct camu_ff_decoder *av)
{
    const AVCodec *codec;
    void *iter = NULL;
    while ((codec = av_codec_iterate(&iter))) {
        if (!av_codec_is_decoder(codec)) continue;
        if (codec->capabilities & (AV_CODEC_CAP_HARDWARE | AV_CODEC_CAP_HYBRID)) {
            al_array_push(av->supported_hw_codecs, codec);
        }
    }
    enum AVHWDeviceType hw_device_type = AV_HWDEVICE_TYPE_NONE;
    while ((hw_device_type = av_hwdevice_iterate_types(hw_device_type)) != AV_HWDEVICE_TYPE_NONE) {
        al_array_push(av->supported_hw_devices, hw_device_type);
    }
}
#endif

static void close_internal(struct camu_ff_decoder *av)
{
    if (av->codec_context) avcodec_free_context(&av->codec_context);
#ifdef CAMU_FF_DECODER_HWACCEL
    if (av->hw_context) av_buffer_unref(&av->hw_context);
    al_array_free(av->supported_hw_codecs);
    al_array_free(av->supported_hw_devices);
    if (av->errored_hw_context) avcodec_free_context(&av->errored_hw_context);
#endif
}

// If decoder_init() fails, we should still expect decoder_free() to by called which will run close_internal().
static bool ff_decoder_init(struct camu_decoder *dec, struct camu_renderer *renderer, struct camu_codec_stream *stream,
    void (*callback)(void *, struct camu_codec_frame *), void *userdata)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)dec;

    av->codec_context = NULL;
#ifdef CAMU_FF_DECODER_HWACCEL
    av->hw_context = NULL;
    al_array_init(av->supported_hw_codecs);
    al_array_init(av->supported_hw_devices);
    av->errored_hw_context = NULL;
#endif

    AVCodecParameters *codecpar = stream->av.stream->codecpar;
    av->codecpar = codecpar;

    const AVCodec *codec = avcodec_find_decoder(codecpar->codec_id);
    if (!codec) {
        log_error("Failed to find decoder.");
        return false;
    }

    bool is_video = codecpar->codec_type == AVMEDIA_TYPE_VIDEO && stream->duration > 0;

#ifdef CAMU_FF_DECODER_HWACCEL
    if (is_video) {
        collect_supported_hwaccels(av);
        av->hw_device_type = AV_HWDEVICE_TYPE_NONE;
        av->hw_pix_fmt = AV_PIX_FMT_NONE;
        for (u32 i = 0; i < ARRAY_SIZE(hwdevices); i++) {
            size_t length = sizeof(hwdevices[i]) - 1;
            av->hw_device_type = hw_device_supported_by_name(av, hwdevices[i], length);
            if (!get_hwdevice_config(av, codec, av->hw_device_type)) {
                av->hw_device_type = AV_HWDEVICE_TYPE_NONE;
            } else {
                break;
            }
        }
        if (av->hw_device_type == AV_HWDEVICE_TYPE_NONE) {
            const AVCodec *hw_codec;
            al_array_foreach(av->supported_hw_codecs, i, hw_codec) {
                if (hw_codec->id == codecpar->codec_id) {
                    // @TODO: hw_codec's without a hwdevice (v4l2m2m, cuvid, amf).
                    size_t length = al_strlen(hw_codec->wrapper_name);
                    av->hw_device_type = hw_device_supported_by_name(av, hw_codec->wrapper_name, length);
                    if (get_hwdevice_config(av, hw_codec, av->hw_device_type)) {
                        codec = hw_codec;
                        break;
                    }
                    av->hw_device_type = AV_HWDEVICE_TYPE_NONE;
                }
            }
        }
        if (av->hw_device_type == AV_HWDEVICE_TYPE_NONE) {
            log_warn("Hardware accelerated video decoding of %s not supported.", codec->name);
        }
    }
#endif

    if (!alloc_codec_context_internal(av, codec, codecpar)) {
        return false;
    }

#ifdef CAMU_FF_DECODER_HWACCEL
    if (is_video && av->hw_device_type != AV_HWDEVICE_TYPE_NONE) {
        //av->codec_context->hwaccel_flags |= AV_HWACCEL_FLAG_IGNORE_LEVEL | AV_HWACCEL_FLAG_ALLOW_PROFILE_MISMATCH;
        if (init_hwdevice_context(av, av->codec_context) == 0) {
            av->renderer = renderer;
            av->codec_context->opaque = av;
            av->codec_context->get_buffer2 = get_buffer2;
            av->codec_context->get_format = get_hw_format;
        } else {
            av->hw_device_type = AV_HWDEVICE_TYPE_NONE;
        }
    }
#else
    (void)renderer;
#endif

    if (is_video) {
        av->thread_count = MIN(8, MAX(1, av_cpu_count() / 2));
#ifdef CAMU_FF_DECODER_HWACCEL
        if (av->hw_device_type == AV_HWDEVICE_TYPE_NONE) {
#endif
            // FF_THREAD_FRAME or FF_THREAD_SLICE.
            // Both modes can get choked on certain videos, higher thread_count with
            // FRAME seems like the safest default.
            // @TODO: Check ffmpeg source, SLICE + FRAME?
            //  - https://ffmpeg.org/ffmpeg-codecs.html#Codec-Options
            av->codec_context->thread_type = FF_THREAD_FRAME;
            av->codec_context->thread_count = av->thread_count;
            log_info("Using %i threads for decoder.", av->thread_count);
#ifdef CAMU_FF_DECODER_HWACCEL
        }
#endif
    }

    //av->codec_context->flags |= AV_CODEC_FLAG_BITEXACT;
    //av->codec_context->flags2 |= AV_CODEC_FLAG2_FAST;

    if (!open_avcodec_internal(av, codec)) {
        return false;
    }

    log_info("Codec: %s (%s) %lldkbps.", codec->name,
        codec->long_name ? codec->long_name : codec->name,
        (av->codec_context->bit_rate > 0) ? av->codec_context->bit_rate / 1000 : 0);

    av->callback = callback;
    av->userdata = userdata;

    return true;
}

static s32 send_packet(struct camu_ff_decoder *av, AVPacket *pkt)
{
    s32 ret = avcodec_send_packet(av->codec_context, pkt);
    if (ret < 0 && !(ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)) {
        log_error("Error sending packet to the decoder (%s).", av_err2str(ret));
    }
    return ret;
}

static s32 ff_decoder_push_av_packet(struct camu_decoder *dec, AVPacket *pkt)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)dec;
    return send_packet(av, pkt);
}

static void ff_decoder_flush(struct camu_decoder *dec)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)dec;
    avcodec_flush_buffers(av->codec_context);
}

static s32 receive_frames(struct camu_ff_decoder *av)
{
    s32 ret;
    for (;;) {
        struct camu_codec_frame *frame = al_alloc_object(struct camu_codec_frame);
        frame->mode = CAMU_FFMPEG_COMPAT;
        frame->av.frame = av_frame_alloc();
        AVFrame *avframe = frame->av.frame;
        ret = avcodec_receive_frame(av->codec_context, avframe);
        if (ret < 0) {
            av_frame_free(&avframe);
            al_free(frame);
            // Checking for EAGAIN should prevent an infinite loop.
            if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) break;
            log_error("Error receiving packet from the decoder (%s).", av_err2str(ret));
            continue;
        }
        if (avframe->best_effort_timestamp == AV_NOPTS_VALUE) {
            // This trips under what I think are normal conditions, more testing is needed.
            //al_assert(avframe->duration == 0);
            avframe->best_effort_timestamp = 0;
        }
        av->callback(av->userdata, frame);
    }
    return ret;
}

static s32 ff_decoder_process(struct camu_decoder *dec)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)dec;
    return receive_frames(av);
}

static void ff_decoder_free(struct camu_decoder **dec)
{
    struct camu_ff_decoder *av = (struct camu_ff_decoder *)*dec;
    close_internal(av);
    al_free(av);
    *dec = NULL;
}

struct camu_decoder *camu_ff_decoder_create(void)
{
    struct camu_ff_decoder *av = al_alloc_object(struct camu_ff_decoder);
    av->dec.mode = CAMU_FFMPEG_COMPAT;
    av->dec.init = ff_decoder_init;
    av->dec.push_av_packet = ff_decoder_push_av_packet;
    av->dec.process = ff_decoder_process;
    av->dec.flush = ff_decoder_flush;
    av->dec.free = ff_decoder_free;
    return (struct camu_decoder *)av;
}