#include "screencast/codec/decoder.h"
#include "screencast/codec/encoder.h"

#include <algorithm>
#include <cassert>
#include <cstdint>
#include <cstdlib>
#include <iostream>
#include <map>
#include <vector>

namespace {

template <typename T> void expect_codec_result(const char* what, const sc::CodecResult<T>& result) {
    if (sc::is_codec_error(result)) {
        std::cerr << what << " failed: " << sc::codec_error(result).message << '\n';
        std::abort();
    }
}

constexpr int kWidth = 128;
constexpr int kHeight = 128;
constexpr int kFrames = 30;
constexpr int kFrameRate = 25;
constexpr uint64_t kNsPerFrame = 1'000'000'000ULL / kFrameRate;

sc::CapturedFrame make_frame(uint32_t index) {
    sc::CapturedFrame frame;
    frame.width = kWidth;
    frame.height = kHeight;
    frame.timestamp_ns = static_cast<uint64_t>(index) * kNsPerFrame;
    frame.pixel_format = sc::PixelFormat::Rgba;
    frame.stride = kWidth * 4;
    frame.pixels.resize(static_cast<std::size_t>(kWidth) * kHeight * 4);

    auto* pixels = reinterpret_cast<std::uint8_t*>(frame.pixels.data());
    for (int y = 0; y < kHeight; ++y) {
        for (int x = 0; x < kWidth; ++x) {
            const std::size_t offset = (static_cast<std::size_t>(y) * kWidth + x) * 4;
            pixels[offset + 0] = static_cast<std::uint8_t>((x + static_cast<int>(index) * 4) & 0xFF);
            pixels[offset + 1] = static_cast<std::uint8_t>((y + static_cast<int>(index) * 2) & 0xFF);
            pixels[offset + 2] = static_cast<std::uint8_t>(((x ^ y) + static_cast<int>(index)) & 0xFF);
            pixels[offset + 3] = 0xFF;
        }
    }

    return frame;
}

int max_channel_difference(const sc::DecodedFrame& decoded, const sc::CapturedFrame& expected) {
    const auto* decoded_pixels = reinterpret_cast<const std::uint8_t*>(decoded.rgba_pixels.data());
    const auto* expected_pixels = reinterpret_cast<const std::uint8_t*>(expected.pixels.data());
    const std::size_t count = std::min(decoded.rgba_pixels.size(), expected.pixels.size());

    int max_diff = 0;
    for (std::size_t i = 0; i < count; ++i) {
        const int diff = decoded_pixels[i] >= expected_pixels[i] ? decoded_pixels[i] - expected_pixels[i]
                                                                 : expected_pixels[i] - decoded_pixels[i];
        max_diff = std::max(max_diff, diff);
    }
    return max_diff;
}

bool starts_with_annex_b_prefix(const sc::EncodedFrame& frame) {
    if (frame.data.size() < 4) {
        return false;
    }
    return frame.data[0] == std::byte{0x00} && frame.data[1] == std::byte{0x00} && frame.data[2] == std::byte{0x00} &&
           frame.data[3] == std::byte{0x01};
}

} // namespace

int main() {
    sc::EncoderConfig encoder_config;
    encoder_config.codec_name = "h264";
    encoder_config.width = kWidth;
    encoder_config.height = kHeight;
    encoder_config.frame_rate_num = kFrameRate;
    encoder_config.frame_rate_den = 1;
    encoder_config.bitrate_kbps = 8000;
    encoder_config.hardware_accel = false;

    auto encoder_result = sc::EncoderFactory::create(encoder_config);
    expect_codec_result("encoder create", encoder_result);
    auto encoder = std::move(sc::codec_value(encoder_result));

    const auto extradata = encoder->get_extradata();
    assert(!extradata.empty());

    sc::DecoderConfig decoder_config;
    decoder_config.codec_name = "h264";
    decoder_config.width = kWidth;
    decoder_config.height = kHeight;
    decoder_config.extradata = extradata;

    auto decoder_result = sc::DecoderFactory::create(decoder_config);
    expect_codec_result("decoder create", decoder_result);
    auto decoder = std::move(sc::codec_value(decoder_result));

    std::map<uint64_t, sc::CapturedFrame> expected_by_timestamp;
    bool saw_keyframe = false;
    bool saw_annex_b_prefix = false;

    auto decode_packets = [&](const std::vector<sc::EncodedFrame>& packets) {
        for (const auto& packet : packets) {
            saw_keyframe = saw_keyframe || packet.is_keyframe;
            saw_annex_b_prefix = saw_annex_b_prefix || starts_with_annex_b_prefix(packet);

            auto decoded_result = decoder->decode(packet);
            expect_codec_result("decode", decoded_result);

            for (const auto& decoded : sc::codec_value(decoded_result)) {
                assert(decoded.width == kWidth);
                assert(decoded.height == kHeight);

                auto it = expected_by_timestamp.find(decoded.capture_timestamp_ns);
                assert(it != expected_by_timestamp.end());
                assert(max_channel_difference(decoded, it->second) <= 64);
            }
        }
    };

    for (uint32_t i = 0; i < kFrames; ++i) {
        const auto frame = make_frame(i);
        expected_by_timestamp.emplace(frame.timestamp_ns, frame);

        if (i == kFrames / 2) {
            encoder->request_keyframe();
        }

        auto encoded_result = encoder->encode(frame);
        expect_codec_result("encode", encoded_result);
        decode_packets(sc::codec_value(encoded_result));
    }

    auto flushed_result = encoder->flush();
    expect_codec_result("flush", flushed_result);
    decode_packets(sc::codec_value(flushed_result));

    assert(!expected_by_timestamp.empty());
    assert(saw_keyframe);
    assert(saw_annex_b_prefix);

    return 0;
}