// Copyright 2018 Citra Emulator Project
// Licensed under GPLv2 or any later version
// Refer to the license.txt file included.

#include "audio_core/hle/ffmpeg_decoder.h"
#include "common/dynamic_library/ffmpeg.h"

using namespace DynamicLibrary;

namespace AudioCore::HLE {

class FFMPEGDecoder::Impl {
public:
    explicit Impl(Memory::MemorySystem& memory);
    ~Impl();
    std::optional<BinaryMessage> ProcessRequest(const BinaryMessage& request);
    bool IsValid() const {
        return have_ffmpeg_dl;
    }

private:
    std::optional<BinaryMessage> Initalize(const BinaryMessage& request);

    void Clear();

    std::optional<BinaryMessage> Decode(const BinaryMessage& request);

    struct AVPacketDeleter {
        void operator()(AVPacket* packet) const {
            FFmpeg::av_packet_free(&packet);
        }
    };

    struct AVCodecContextDeleter {
        void operator()(AVCodecContext* context) const {
            FFmpeg::avcodec_free_context(&context);
        }
    };

    struct AVCodecParserContextDeleter {
        void operator()(AVCodecParserContext* parser) const {
            FFmpeg::av_parser_close(parser);
        }
    };

    struct AVFrameDeleter {
        void operator()(AVFrame* frame) const {
            FFmpeg::av_frame_free(&frame);
        }
    };

    bool initalized = false;
    bool have_ffmpeg_dl;

    Memory::MemorySystem& memory;

    const AVCodec* codec;
    std::unique_ptr<AVCodecContext, AVCodecContextDeleter> av_context;
    std::unique_ptr<AVCodecParserContext, AVCodecParserContextDeleter> parser;
    std::unique_ptr<AVPacket, AVPacketDeleter> av_packet;
    std::unique_ptr<AVFrame, AVFrameDeleter> decoded_frame;
};

FFMPEGDecoder::Impl::Impl(Memory::MemorySystem& memory) : memory(memory) {
    have_ffmpeg_dl = FFmpeg::LoadFFmpeg();
}

FFMPEGDecoder::Impl::~Impl() = default;

std::optional<BinaryMessage> FFMPEGDecoder::Impl::ProcessRequest(const BinaryMessage& request) {
    if (request.header.codec != DecoderCodec::DecodeAAC) {
        LOG_ERROR(Audio_DSP, "Got wrong codec {}", static_cast<u16>(request.header.codec));
        return {};
    }

    switch (request.header.cmd) {
    case DecoderCommand::Init: {
        return Initalize(request);
    }
    case DecoderCommand::EncodeDecode: {
        return Decode(request);
    }
    case DecoderCommand::Shutdown:
    case DecoderCommand::SaveState:
    case DecoderCommand::LoadState: {
        LOG_WARNING(Audio_DSP, "Got unimplemented binary request: {}",
                    static_cast<u16>(request.header.cmd));
        BinaryMessage response = request;
        response.header.result = ResultStatus::Success;
        return response;
    }
    default:
        LOG_ERROR(Audio_DSP, "Got unknown binary request: {}",
                  static_cast<u16>(request.header.cmd));
        return {};
    }
}

std::optional<BinaryMessage> FFMPEGDecoder::Impl::Initalize(const BinaryMessage& request) {
    if (initalized) {
        Clear();
    }

    BinaryMessage response = request;
    response.header.result = ResultStatus::Success;

    if (!have_ffmpeg_dl) {
        return response;
    }

    av_packet.reset(FFmpeg::av_packet_alloc());

    codec = FFmpeg::avcodec_find_decoder(AV_CODEC_ID_AAC);
    if (!codec) {
        LOG_ERROR(Audio_DSP, "Codec not found\n");
        return response;
    }

    parser.reset(FFmpeg::av_parser_init(codec->id));
    if (!parser) {
        LOG_ERROR(Audio_DSP, "Parser not found\n");
        return response;
    }

    av_context.reset(FFmpeg::avcodec_alloc_context3(codec));
    if (!av_context) {
        LOG_ERROR(Audio_DSP, "Could not allocate audio codec context\n");
        return response;
    }

    if (FFmpeg::avcodec_open2(av_context.get(), codec, nullptr) < 0) {
        LOG_ERROR(Audio_DSP, "Could not open codec\n");
        return response;
    }

    initalized = true;
    return response;
}

void FFMPEGDecoder::Impl::Clear() {
    if (!have_ffmpeg_dl) {
        return;
    }

    av_context.reset();
    parser.reset();
    decoded_frame.reset();
    av_packet.reset();
}

std::optional<BinaryMessage> FFMPEGDecoder::Impl::Decode(const BinaryMessage& request) {
    BinaryMessage response{};
    response.header.codec = request.header.codec;
    response.header.cmd = request.header.cmd;
    response.decode_aac_response.size = request.decode_aac_request.size;

    if (!initalized) {
        LOG_DEBUG(Audio_DSP, "Decoder not initalized");
        // This is a hack to continue games that are not compiled with the aac codec
        response.decode_aac_response.num_channels = 2;
        response.decode_aac_response.num_samples = 1024;
        return response;
    }

    if (request.decode_aac_request.src_addr < Memory::FCRAM_PADDR ||
        request.decode_aac_request.src_addr + request.decode_aac_request.size >
            Memory::FCRAM_PADDR + Memory::FCRAM_SIZE) {
        LOG_ERROR(Audio_DSP, "Got out of bounds src_addr {:08x}",
                  request.decode_aac_request.src_addr);
        return {};
    }
    u8* data = memory.GetFCRAMPointer(request.decode_aac_request.src_addr - Memory::FCRAM_PADDR);

    std::array<std::vector<u8>, 2> out_streams;

    std::size_t data_size = request.decode_aac_request.size;
    while (data_size > 0) {
        if (!decoded_frame) {
            decoded_frame.reset(FFmpeg::av_frame_alloc());
            if (!decoded_frame) {
                LOG_ERROR(Audio_DSP, "Could not allocate audio frame");
                return {};
            }
        }

        int ret = FFmpeg::av_parser_parse2(parser.get(), av_context.get(), &av_packet->data,
                                           &av_packet->size, data, static_cast<int>(data_size),
                                           AV_NOPTS_VALUE, AV_NOPTS_VALUE, 0);
        if (ret < 0) {
            LOG_ERROR(Audio_DSP, "Error while parsing");
            return {};
        }
        data += ret;
        data_size -= ret;

        ret = FFmpeg::avcodec_send_packet(av_context.get(), av_packet.get());
        if (ret < 0) {
            LOG_ERROR(Audio_DSP, "Error submitting the packet to the decoder");
            return {};
        }

        if (av_packet->size) {
            while (ret >= 0) {
                ret = FFmpeg::avcodec_receive_frame(av_context.get(), decoded_frame.get());
                if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF)
                    break;
                else if (ret < 0) {
                    LOG_ERROR(Audio_DSP, "Error during decoding");
                    return {};
                }
                int bytes_per_sample = FFmpeg::av_get_bytes_per_sample(av_context->sample_fmt);
                if (bytes_per_sample < 0) {
                    LOG_ERROR(Audio_DSP, "Failed to calculate data size");
                    return {};
                }

#if LIBAVCODEC_VERSION_INT >= AV_VERSION_INT(59, 24, 100)
                auto num_channels = static_cast<u32>(decoded_frame->ch_layout.nb_channels);
#else
                auto num_channels = static_cast<u32>(decoded_frame->channels);
#endif

                ASSERT(num_channels <= out_streams.size());

                std::size_t size = bytes_per_sample * (decoded_frame->nb_samples);

                response.decode_aac_response.sample_rate =
                    GetSampleRateEnum(decoded_frame->sample_rate);
                response.decode_aac_response.num_channels = num_channels;
                response.decode_aac_response.num_samples += decoded_frame->nb_samples;

                // FFmpeg converts to 32 signed floating point PCM, we need s16 PCM so we need to
                // convert it
                f32 val_float;
                for (std::size_t current_pos(0); current_pos < size;) {
                    for (std::size_t channel(0); channel < num_channels; channel++) {
                        std::memcpy(&val_float, decoded_frame->data[channel] + current_pos,
                                    sizeof(val_float));
                        val_float = std::clamp(val_float, -1.0f, 1.0f);
                        s16 val = static_cast<s16>(0x7FFF * val_float);
                        out_streams[channel].push_back(val & 0xFF);
                        out_streams[channel].push_back(val >> 8);
                    }
                    current_pos += sizeof(val_float);
                }
            }
        }
    }

    if (out_streams[0].size() != 0) {
        if (request.decode_aac_request.dst_addr_ch0 < Memory::FCRAM_PADDR ||
            request.decode_aac_request.dst_addr_ch0 + out_streams[0].size() >
                Memory::FCRAM_PADDR + Memory::FCRAM_SIZE) {
            LOG_ERROR(Audio_DSP, "Got out of bounds dst_addr_ch0 {:08x}",
                      request.decode_aac_request.dst_addr_ch0);
            return {};
        }
        std::memcpy(
            memory.GetFCRAMPointer(request.decode_aac_request.dst_addr_ch0 - Memory::FCRAM_PADDR),
            out_streams[0].data(), out_streams[0].size());
    }

    if (out_streams[1].size() != 0) {
        if (request.decode_aac_request.dst_addr_ch1 < Memory::FCRAM_PADDR ||
            request.decode_aac_request.dst_addr_ch1 + out_streams[1].size() >
                Memory::FCRAM_PADDR + Memory::FCRAM_SIZE) {
            LOG_ERROR(Audio_DSP, "Got out of bounds dst_addr_ch1 {:08x}",
                      request.decode_aac_request.dst_addr_ch1);
            return {};
        }
        std::memcpy(
            memory.GetFCRAMPointer(request.decode_aac_request.dst_addr_ch1 - Memory::FCRAM_PADDR),
            out_streams[1].data(), out_streams[1].size());
    }
    return response;
}

FFMPEGDecoder::FFMPEGDecoder(Memory::MemorySystem& memory) : impl(std::make_unique<Impl>(memory)) {}

FFMPEGDecoder::~FFMPEGDecoder() = default;

std::optional<BinaryMessage> FFMPEGDecoder::ProcessRequest(const BinaryMessage& request) {
    return impl->ProcessRequest(request);
}

bool FFMPEGDecoder::IsValid() const {
    return impl->IsValid();
}

} // namespace AudioCore::HLE