#include <gtest/gtest.h>

#include <memory>
#include <vector>

#include "hw/hw_factory.h"

namespace rk3588 {
namespace {

static SimpleJson MakeDecoderConfig(const std::string& backend, const std::string& codec,
                                    const std::vector<uint8_t>& extradata = {}) {
    SimpleJson::Object obj;
    obj["backend"] = SimpleJson(backend);
    obj["codec"] = SimpleJson(codec);
    if (!extradata.empty()) {
        SimpleJson::Array arr;
        arr.reserve(extradata.size());
        for (auto b : extradata) {
            arr.emplace_back(static_cast<double>(b));
        }
        obj["extradata"] = SimpleJson(std::move(arr));
    }
    return SimpleJson(std::move(obj));
}

static SimpleJson MakeEncoderConfig(const std::string& backend, const std::string& codec) {
    SimpleJson::Object obj;
    obj["backend"] = SimpleJson(backend);
    obj["codec"] = SimpleJson(codec);
    obj["fps"] = SimpleJson(25.0);
    obj["gop"] = SimpleJson(50.0);
    obj["bitrate_kbps"] = SimpleJson(1000.0);
    obj["width"] = SimpleJson(64.0);
    obj["height"] = SimpleJson(64.0);
    obj["pixel_format"] = SimpleJson(std::string("nv12"));
    return SimpleJson(std::move(obj));
}

static std::shared_ptr<Frame> MakeTestFrame() {
    auto frame = std::make_shared<Frame>();
    frame->width = 64;
    frame->height = 64;
    frame->format = PixelFormat::NV12;
    frame->plane_count = 2;
    frame->stride = 64;
    const size_t y_size = static_cast<size_t>(frame->width * frame->height);
    const size_t total = y_size * 3 / 2;
    auto buffer = std::make_shared<std::vector<uint8_t>>(total, 0);
    frame->data = buffer->data();
    frame->data_size = buffer->size();
    frame->SetOwner(buffer);
    frame->planes[0] = {frame->data, frame->width, static_cast<int>(y_size), 0};
    frame->planes[1] = {frame->data + y_size, frame->width, static_cast<int>(y_size / 2), static_cast<int>(y_size)};
    frame->pts = 1000;
    frame->SyncBufferFromFrame();
    return frame;
}

TEST(CodecBackendTest, DecoderOpenCloseSmoke) {
#if !defined(RK3588_ENABLE_MPP) && !defined(RK3588_ENABLE_FFMPEG)
    GTEST_SKIP() << "Codec backends not enabled";
#endif
    std::string backend = "mpp";
#if !defined(RK3588_ENABLE_MPP)
    backend = "ffmpeg";
#endif
    auto cfg = MakeDecoderConfig(backend, "h264");
    auto decoder = HwFactory::CreateDecoder(cfg);
    ASSERT_NE(decoder, nullptr);
    if (decoder->Open(cfg).Failed()) {
        GTEST_SKIP() << "Decoder backend not available";
    }
    decoder->Close();
}

TEST(CodecBackendTest, EncoderOpenCloseSmoke) {
#if !defined(RK3588_ENABLE_MPP) && !defined(RK3588_ENABLE_FFMPEG)
    GTEST_SKIP() << "Codec backends not enabled";
#endif
    std::string backend = "mpp";
#if !defined(RK3588_ENABLE_MPP)
    backend = "ffmpeg";
#endif
    auto cfg = MakeEncoderConfig(backend, "h264");
    auto encoder = HwFactory::CreateEncoder(cfg);
    ASSERT_NE(encoder, nullptr);
    if (encoder->Open(cfg).Failed()) {
        GTEST_SKIP() << "Encoder backend not available";
    }
    encoder->Close();
}

TEST(CodecBackendTest, CodecPipelineEncodeDecodeOneFrame) {
#if !defined(RK3588_ENABLE_MPP)
    GTEST_SKIP() << "MPP backend not enabled";
#endif
    auto encoder_cfg = MakeEncoderConfig("mpp", "h264");
    auto encoder = HwFactory::CreateEncoder(encoder_cfg);
    ASSERT_NE(encoder, nullptr);
    if (encoder->Open(encoder_cfg).Failed()) {
        GTEST_SKIP() << "Encoder not available";
    }

    auto frame = MakeTestFrame();
    ASSERT_TRUE(encoder->Send(frame).Ok());

    auto pkt_res = encoder->Receive();
    if (!pkt_res.Ok() || pkt_res.Value().data.empty()) {
        GTEST_SKIP() << "Encoder did not produce packet";
    }
    auto pkt = std::move(pkt_res.Value());
    auto extradata = encoder->ExtraData();

    auto decoder_cfg = MakeDecoderConfig("mpp", "h264", extradata);
    auto decoder = HwFactory::CreateDecoder(decoder_cfg);
    ASSERT_NE(decoder, nullptr);
    if (decoder->Open(decoder_cfg).Failed()) {
        GTEST_SKIP() << "Decoder not available";
    }

    DecodePacket dp;
    dp.data = pkt.data.data();
    dp.size = pkt.data.size();
    dp.pts = pkt.pts;
    dp.keyframe = pkt.keyframe;
    ASSERT_TRUE(decoder->Send(dp).Ok());

    bool got_frame = false;
    for (int i = 0; i < 10; ++i) {
        auto out = decoder->Receive();
        if (out.Ok() && out.Value()) {
            got_frame = true;
            break;
        }
        if (out.ErrMessage() != "no_frame") break;
    }
    EXPECT_TRUE(got_frame);
    decoder->Close();
    encoder->Close();
}

}  // namespace
}  // namespace rk3588