/*
 *******************************************************************************
 *                                   INCLUDE
 *******************************************************************************
 */
#include "Util.h"
#include <Poco/Exception.h>
#include <Poco/Format.h>
#include <array>

namespace Video
{
namespace FFmpeg
{
/*
 *******************************************************************************
 *                                   TYPEDEFS
 *******************************************************************************
 */
struct CodecIDItem
{
    Core::CodecID cid;
    AVCodecID     avid;
};

struct SampleFormatItem
{
    Core::SampleFormat cfmt;
    AVSampleFormat     avfmt;
};

/*
 *******************************************************************************
 *                                  LOCAL VARIABLES
 *******************************************************************************
 */
static const std::array<CodecIDItem, 8> codecMap = {{
    {Core::CodecID::H264, AV_CODEC_ID_H264},
    {Core::CodecID::H265, AV_CODEC_ID_H265},
    {Core::CodecID::MJPEG, AV_CODEC_ID_MJPEG},
    {Core::CodecID::MP3, AV_CODEC_ID_MP3},
    {Core::CodecID::AAC, AV_CODEC_ID_AAC},
    {Core::CodecID::OPUS, AV_CODEC_ID_OPUS},
    {Core::CodecID::G711A, AV_CODEC_ID_PCM_ALAW},
    {Core::CodecID::G711U, AV_CODEC_ID_PCM_MULAW},
}};

static const std::array<SampleFormatItem, 12> formatMap = {{
    {
        Core::SampleFormat::U8,
        AV_SAMPLE_FMT_U8,
    },
    {
        Core::SampleFormat::S16,
        AV_SAMPLE_FMT_S16,
    },
    {
        Core::SampleFormat::S32,
        AV_SAMPLE_FMT_S32,
    },
    {
        Core::SampleFormat::S64,
        AV_SAMPLE_FMT_S64,
    },
    {
        Core::SampleFormat::FLT,
        AV_SAMPLE_FMT_FLT,
    },
    {
        Core::SampleFormat::DBL,
        AV_SAMPLE_FMT_DBL,
    },
    {
        Core::SampleFormat::U8P,
        AV_SAMPLE_FMT_U8P,
    },
    {
        Core::SampleFormat::S16P,
        AV_SAMPLE_FMT_S16P,
    },
    {
        Core::SampleFormat::S32P,
        AV_SAMPLE_FMT_S32P,
    },
    {
        Core::SampleFormat::S64P,
        AV_SAMPLE_FMT_S64P,
    },
    {
        Core::SampleFormat::FLTP,
        AV_SAMPLE_FMT_FLTP,
    },
    {
        Core::SampleFormat::DBLP,
        AV_SAMPLE_FMT_DBLP,
    },
}};

/*
 *******************************************************************************
 *                                   FUNCTIONS
 *******************************************************************************
 */
static Core::CodecID findCodecID(AVCodecID id)
{
    for (const auto& m : codecMap)
    {
        if (id == m.avid)
            return m.cid;
    }

    throw Poco::NotFoundException(
        Poco::format("Video::FFmpeg: Can't find %s in codec-map", Util::getCodecIDString(id)));
}

static AVCodecID findCodecID(Core::CodecID id)
{
    for (const auto& m : codecMap)
    {
        if (id == m.cid)
            return m.avid;
    }

    throw Poco::NotFoundException(Poco::format("Video::FFmpeg: Can't find %s in codec-map", Core::codecIDToString(id)));
}

/*
 *******************************************************************************
 *                                CLASS DEFINE
 *******************************************************************************
 */
std::string Util::getMediaTypeString(AVMediaType mediaType)
{
    return av_get_media_type_string(mediaType);
}

std::string Util::getErrorString(int err)
{
    char buf[AV_ERROR_MAX_STRING_SIZE];
    return av_make_error_string(buf, AV_ERROR_MAX_STRING_SIZE, err);
}

std::string Util::getCodecIDString(AVCodecID id)
{
    const char* name = avcodec_get_name(id);
    if (name != nullptr)
        return name;

    return Poco::format("UNKNOWN(%u)", (uint32_t)id);
}

std::string Util::getSampleFormatString(AVSampleFormat fmt)
{
    const char* name = av_get_sample_fmt_name(fmt);
    if (name != nullptr)
        return name;

    return Poco::format("UNKNOWN(%u)", (uint32_t)fmt);
}

const AVCodec* Util::getEncoder(Core::CodecID codec)
{
    return avcodec_find_encoder(findCodecID(codec));
}

const AVCodec* Util::getDecoder(Core::CodecID codec)
{
    return avcodec_find_decoder(findCodecID(codec));
}

const Core::CodecID Util::getCodecID(AVCodecID id)
{
    return findCodecID(id);
}

AVChannelLayout Util::getChannelLayout(Core::ChannelLayout layout)
{
    switch (layout)
    {
    case Core::ChannelLayout::MONO:
    case Core::ChannelLayout::L:
    case Core::ChannelLayout::R:
        return AV_CHANNEL_LAYOUT_MONO;

    default:
    case Core::ChannelLayout::STEREO:
        return AV_CHANNEL_LAYOUT_STEREO;
    }
}

Core::ChannelLayout Util::getChannelLayout(AVChannelLayout layout)
{
    AVChannelLayout other = AV_CHANNEL_LAYOUT_STEREO;
    if (av_channel_layout_compare(&layout, &other) == 0)
        return Core::ChannelLayout::STEREO;

    return Core::ChannelLayout::MONO;
}

AVSampleFormat Util::getSampleFormat(Core::SampleFormat format)
{
    for (const auto& m : formatMap)
    {
        if (format == m.cfmt)
            return m.avfmt;
    }

    throw Poco::NotFoundException(
        Poco::format("Video::FFmpeg: Can't find %s in format-map", Core::sampleFormatName(format)));
}

Core::SampleFormat Util::getSampleFormat(AVSampleFormat format)
{
    for (const auto& m : formatMap)
    {
        if (format == m.avfmt)
            return m.cfmt;
    }

    throw Poco::NotFoundException(
        Poco::format("Video::FFmpeg: Can't find %s in format-map", getSampleFormatString(format)));
}

int Util::convertJPG2NV16(std::string srcPath, std::string dstPath)
{
#define TARGET_WIDTH 1920
#define TARGET_HEIGHT 1080
    const char* input_filename  = srcPath.c_str();
    const char* output_filename = dstPath.c_str();

    AVFormatContext* fmt_ctx   = NULL;
    AVCodecContext*  codec_ctx = NULL;
    const AVCodec*   codec     = NULL;
    AVFrame *        frame = NULL, *scaled_frame = NULL;
    AVPacket*        pkt              = NULL;
    SwsContext*      sws_ctx          = NULL;
    FILE*            outfile          = NULL;
    int              ret              = 0;
    int              video_stream_idx = -1;

    if ((ret = avformat_open_input(&fmt_ctx, input_filename, NULL, NULL)) < 0)
    {
        fprintf(stderr, "Cannot open input file\n");
        return -1;
    }

    if ((ret = avformat_find_stream_info(fmt_ctx, NULL)) < 0)
    {
        fprintf(stderr, "Cannot find stream information\n");
        return -1;
    }

    for (int i = 0; i < fmt_ctx->nb_streams; i++)
    {
        if (fmt_ctx->streams[i]->codecpar->codec_type == AVMEDIA_TYPE_VIDEO)
        {
            video_stream_idx = i;
            break;
        }
    }

    if (video_stream_idx == -1)
    {
        fprintf(stderr, "No video stream found\n");
        ret = -1;
        return -1;
    }

    AVCodecParameters* codecpar = fmt_ctx->streams[video_stream_idx]->codecpar;
    codec                       = avcodec_find_decoder(codecpar->codec_id);
    if (!codec)
    {
        fprintf(stderr, "Unsupported codec\n");
        ret = -1;
        return -1;
    }

    codec_ctx = avcodec_alloc_context3(codec);
    if (!codec_ctx)
    {
        fprintf(stderr, "Cannot allocate codec context\n");
        ret = -1;
        return -1;
    }

    if ((ret = avcodec_parameters_to_context(codec_ctx, codecpar)) < 0)
    {
        fprintf(stderr, "Cannot copy codec parameters\n");
        return -1;
    }

    if ((ret = avcodec_open2(codec_ctx, codec, NULL)) < 0)
    {
        fprintf(stderr, "Cannot open codec\n");
        return -1;
    }

    frame = av_frame_alloc();
    if (!frame)
    {
        fprintf(stderr, "Cannot allocate frame\n");
        ret = -1;
        return -1;
    }

    scaled_frame = av_frame_alloc();
    if (!scaled_frame)
    {
        fprintf(stderr, "Cannot allocate scaled frame\n");
        ret = -1;
        return -1;
    }

    scaled_frame->width  = TARGET_WIDTH;
    scaled_frame->height = TARGET_HEIGHT;
    scaled_frame->format = AV_PIX_FMT_YUV422P;

    if ((ret = av_frame_get_buffer(scaled_frame, 32)) < 0)
    {
        fprintf(stderr, "Cannot allocate scaled frame buffer\n");
        return -1;
    }

    sws_ctx = sws_getContext(codec_ctx->width, codec_ctx->height, codec_ctx->pix_fmt, TARGET_WIDTH, TARGET_HEIGHT,
                             AV_PIX_FMT_YUV422P, SWS_BILINEAR, NULL, NULL, NULL);
    if (!sws_ctx)
    {
        fprintf(stderr, "Cannot create sws context\n");
        ret = -1;
        return -1;
    }

    outfile = fopen(output_filename, "wb");
    if (!outfile)
    {
        fprintf(stderr, "Cannot open output file\n");
        ret = -1;
        return -1;
    }

    pkt = av_packet_alloc();
    if (!pkt)
    {
        fprintf(stderr, "Cannot allocate packet\n");
        ret = -1;
        return -1;
    }

    while (av_read_frame(fmt_ctx, pkt) >= 0)
    {
        if (pkt->stream_index == video_stream_idx)
        {
            ret = avcodec_send_packet(codec_ctx, pkt);
            if (ret < 0)
            {
                fprintf(stderr, "Error sending packet\n");
                break;
            }

            while (ret >= 0)
            {
                ret = avcodec_receive_frame(codec_ctx, frame);
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    break;
                else if (ret < 0)
                {
                    fprintf(stderr, "Error receiving frame\n");
                    break;
                }

                sws_scale(sws_ctx, (const uint8_t* const*)frame->data, frame->linesize, 0, frame->height,
                          scaled_frame->data, scaled_frame->linesize);

                int width  = scaled_frame->width;
                int height = scaled_frame->height;

                for (int i = 0; i < height; i++)
                {
                    fwrite(scaled_frame->data[0] + i * scaled_frame->linesize[0], 1, width, outfile);
                }

                for (int i = 0; i < height; i++)
                {
                    uint8_t* u_line = scaled_frame->data[1] + i * scaled_frame->linesize[1];
                    uint8_t* v_line = scaled_frame->data[2] + i * scaled_frame->linesize[2];

                    for (int j = 0; j < width / 2; j++)
                    {
                        fwrite(&v_line[j], 1, 1, outfile);
                        fwrite(&u_line[j], 1, 1, outfile);
                    }
                }

                printf("Image decoded and converted: %dx%d -> %dx%d\n", frame->width, frame->height, TARGET_WIDTH,
                       TARGET_HEIGHT);
            }
        }
        av_packet_unref(pkt);

        break;
    }

    printf("Conversion completed. Output saved to %s\n", output_filename);

    if (outfile)
        fclose(outfile);
    if (sws_ctx)
        sws_freeContext(sws_ctx);
    if (pkt)
        av_packet_free(&pkt);
    if (frame)
        av_frame_free(&frame);
    if (scaled_frame)
        av_frame_free(&scaled_frame);
    if (codec_ctx)
        avcodec_free_context(&codec_ctx);
    if (fmt_ctx)
        avformat_close_input(&fmt_ctx);

    return ret < 0 ? -1 : 0;
}
}
}