/*
 * Copyright (c) 2023 Fancy Code.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#include "FfmpegDecoderV2.h"
#include "ILog.h"
#ifdef __cplusplus
extern "C" {
#endif
#include <libavcodec/avcodec.h>
#include <libavcodec/codec.h>
#include <libavcodec/codec_id.h>
#include <libavcodec/packet.h>
#include <libavutil/avutil.h>
#include <libavutil/channel_layout.h>
#include <libavutil/error.h>
#include <libavutil/frame.h>
#include <libavutil/pixfmt.h>
#include <libavutil/samplefmt.h>
#ifdef __cplusplus
}
#endif
#include <cstddef>
#include <cstdlib>
#include <errno.h>
#include <functional>
#include <stdint.h>
FfmpegDecoderV2::FfmpegDecoderV2(const enum AVCodecID &codecId, const AVPixelFormat &decodePixelFormat,
                                 const int &width, const int &height)
    : mCodecId(codecId), mCodec(nullptr), mCodecCtx(nullptr), mFrame(nullptr), mPacket(nullptr), mParser(nullptr),
      mVideoWidth(width), mVideoHeight(height), mDecodePixelFormat(decodePixelFormat)
{
}
bool FfmpegDecoderV2::Init(void)
{
    int ret = 0;
    LogInfo("find decoder : %s\n", avcodec_get_name(mCodecId));
    mCodec = (AVCodec *)avcodec_find_decoder(mCodecId);
    // mCodec = (AVCodec *)avcodec_find_decoder_by_name("libfdk_aac");
    if (!(mCodec)) {
        LogError("decoder not found:%s\n", avcodec_get_name(mCodecId));
        return false;
    }
    mCodecCtx = avcodec_alloc_context3((const AVCodec *)(mCodec));
    if (!(mCodecCtx)) {
        LogError("Could not allocate codec context\n");
        return false;
    }
    if (AVMEDIA_TYPE_AUDIO == mCodec->type) {
        LogInfo("Audio decoder.\n");
        /* put sample parameters */
        mCodecCtx->bit_rate = 64000;
        // mCodecCtx->bit_rate = 352800;
        // mCodecCtx->sample_rate = 8000;

        /* check that the encoder supports s16 pcm input */
        mCodecCtx->sample_fmt = AV_SAMPLE_FMT_S16;
        if (!check_sample_fmt(mCodec, mCodecCtx->sample_fmt)) {
            LogError("decoder does not support sample format %s", av_get_sample_fmt_name(mCodecCtx->sample_fmt));
            return false;
        }

        /* select other audio parameters supported by the encoder */
        mCodecCtx->sample_rate = select_sample_rate(mCodec);
        LogInfo("decoder sample_rate:%d\n", mCodecCtx->sample_rate);
        // const AVChannelLayout src = (AVChannelLayout)AV_CHANNEL_LAYOUT_MONO;
        // av_channel_layout_copy(&mCodecCtx->ch_layout, &src);
        ret = select_channel_layout(mCodec, &(mCodecCtx->ch_layout));
        if (ret < 0) {
            LogError("Could not set channel layout\n");
            return false;
        }
    }
    else {
        mCodecCtx->pix_fmt = mDecodePixelFormat;
        mCodecCtx->width = mVideoWidth;
        mCodecCtx->height = mVideoHeight;
    }
    if ((ret = avcodec_open2(mCodecCtx, mCodec, nullptr)) < 0) {
        char error_str[AV_ERROR_MAX_STRING_SIZE] = {0};
        LogError("Could not open codec:%s\n", av_make_error_string(error_str, AV_ERROR_MAX_STRING_SIZE, ret));
        return false;
    }
    mFrame = av_frame_alloc();
    if (!mFrame) {
        LogError("Could not allocate video frame\n");
        return false;
    }
    mPacket = av_packet_alloc();
    if (!mPacket) {
        LogError("Could not allocate video frame\n");
        return false;
    }
    // mParser = av_parser_init(mCodec->id);
    // if (!mParser) {
    //     LogError("mParser not found : %s\n", avcodec_get_name(mCodec->id));
    //     return false;
    // }
    if (AVMEDIA_TYPE_AUDIO == mCodec->type) {
        // mFrame->nb_samples = mCodecCtx->frame_size;
        // mFrame->format = mCodecCtx->sample_fmt;
        ret = av_channel_layout_copy(&(mFrame->ch_layout), &(mCodecCtx->ch_layout));
        if (ret < 0) {
            LogError("Could not copy channel layout\n");
            return false;
        }
    }
    LogInfo("init success pix_fmt = %d\n", mCodecCtx->pix_fmt);
    return true;
}
bool FfmpegDecoderV2::UnInit(void)
{
    LogInfo("uninit %s\n", avcodec_get_name(mCodecId));
    av_packet_free(&mPacket);
    mPacket = nullptr;
    if (mParser) {
        av_parser_close(mParser);
        mParser = nullptr;
    }
    avcodec_free_context(&mCodecCtx);
    mCodecCtx = nullptr;
    if (mFrame) {
        av_frame_free(&mFrame);
        mFrame = nullptr;
    }
    return true;
}
void FfmpegDecoderV2::DecodeData(const void *data, const size_t &size, const unsigned long long &pts,
                                 std::function<void(AVFrame *frame)> callback)
{
    if (nullptr == mParser) {
        mPacket->data = (uint8_t *)data;
        mPacket->size = size;
        mPacket->pts = pts;
        mPacket->dts = mPacket->pts;
        // LogInfo("source data mPacket->pts:%d\n", mPacket->pts);
        AVDecodeData(mPacket, callback);
        return;
    }
    AVParseData(data, size, callback);
}
void inline FfmpegDecoderV2::AVParseData(const void *data, const size_t &size,
                                         std::function<void(AVFrame *frame)> callback)
{
    if (nullptr == data) {
        LogError("data is null\n");
        return;
    }
    uint8_t *frameData = (uint8_t *)data;
    size_t data_size = size;
    while (data_size > 0) {
        int ret = av_parser_parse2(mParser,
                                   mCodecCtx,
                                   &mPacket->data,
                                   &mPacket->size,
                                   frameData,
                                   data_size,
                                   AV_NOPTS_VALUE,
                                   AV_NOPTS_VALUE,
                                   0);
        if (ret < 0) {
            LogError("av_parse_frame failed\n");
            break;
        }
        frameData += ret;
        data_size -= ret;
        if (mPacket->size) {
            AVDecodeData(mPacket, callback);
        }
    }
}
// static void save_code_stream_file(const void *data, const size_t &size)
// {
//     char OutPath[128] = {0};
//     const void *pData = data;
//     FILE *file = NULL;
//     LogInfo("save_code_stream_file size = %d\n", size);
//     sprintf(OutPath, "./test.yuv");
//     file = fopen(OutPath, "a+");

//     if (file) { // TODO: Don't open very time.
//         fwrite(pData, 1, size, file);
//         fflush(file);
//     }

//     if (file)
//         fclose(file);
// }
// static void pgm_save(unsigned char *buf, int wrap, int xsize, int ysize, char *filename)
// {
//     FILE *f;
//     int i;

//     f = fopen(filename, "wb");
//     fprintf(f, "P5\n%d %d\n%d\n", xsize, ysize, 255);
//     for (i = 0; i < ysize; i++)
//         fwrite(buf + i * wrap, 1, xsize, f);
//     fclose(f);
// }
void inline FfmpegDecoderV2::AVDecodeData(AVPacket *pkt, std::function<void(AVFrame *frame)> callback)
{
    int ret = avcodec_send_packet(mCodecCtx, pkt);
    if (ret < 0) {
        LogError("Error sending a packet for decoding\n");
        av_packet_unref(pkt);
        return;
    }
    while (ret >= 0) {
        ret = avcodec_receive_frame(mCodecCtx, mFrame);
        if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) {
            break;
        }
        if (ret < 0) {
            LogError("Error during decoding\n");
            break;
        }
        if (callback) {
            // int i, ch, data_size;
            // data_size = av_get_bytes_per_sample(mCodecCtx->sample_fmt);
            // for (i = 0; i < mFrame->nb_samples; i++)
            //     for (ch = 0; ch < mCodecCtx->ch_layout.nb_channels; ch++)
            //         // fwrite(frame->data[ch] + data_size * i, 1, data_size, outfile);
            // save_code_stream_file(mFrame->data[ch] + data_size * i, data_size);
            // save_code_stream_file(mFrame->data[0], mFrame->linesize[0]);
            // if (mCodecId == AV_CODEC_ID_H264) {
            //     pgm_save(mFrame->data[0], mFrame->linesize[0], mFrame->width, mFrame->height, "./test.yuv");
            // }
            // LogInfo("decode frame pts = %llu, nb_samples = %d\n", mFrame->pts, mFrame->nb_samples);
            callback(mFrame);
        }
        break;
    }
    av_packet_unref(pkt);
}
/* just pick the highest supported samplerate */
int FfmpegDecoderV2::select_sample_rate(const AVCodec *codec)
{
    const int *p;
    int best_samplerate = 0;

    if (!codec->supported_samplerates)
        return 44100;

    p = codec->supported_samplerates;
    while (*p) {
        if (!best_samplerate || abs(44100 - *p) < abs(44100 - best_samplerate))
            best_samplerate = *p;
        p++;
    }
    return best_samplerate;
}
/* select layout with the highest channel count */
int FfmpegDecoderV2::select_channel_layout(const AVCodec *codec, AVChannelLayout *dst)
{
    const AVChannelLayout *p, *best_ch_layout = nullptr;
    int best_nb_channels = 0;
    AVChannelLayout channelLayout = AV_CHANNEL_LAYOUT_STEREO;
    if (!codec->ch_layouts)
        return av_channel_layout_copy(dst, &channelLayout);

    p = codec->ch_layouts;
    while (p->nb_channels) {
        int nb_channels = p->nb_channels;

        if (nb_channels > best_nb_channels) {
            best_ch_layout = p;
            best_nb_channels = nb_channels;
        }
        p++;
    }
    return av_channel_layout_copy(dst, best_ch_layout);
}
/* check that a given sample format is supported by the encoder */
int FfmpegDecoderV2::check_sample_fmt(const AVCodec *codec, enum AVSampleFormat sample_fmt)
{
    const enum AVSampleFormat *p = codec->sample_fmts;

    while (*p != AV_SAMPLE_FMT_NONE) {
        if (*p == sample_fmt)
            return 1;
        p++;
    }
    return 0;
}