// Phase 5 integration test: the local UDP sender -> receiver loopback without
// the interactive parts. Synthetic frames go through the real chain —
// encode, RTP packetization, a localhost UDP socket, depacketization, and
// decode (with in-band SPS/PPS, since there is no signaling channel).

#include "screencast/codec/decoder.h"
#include "screencast/codec/encoder.h"
#include "screencast/network/h264_packetizer.h"
#include "screencast/network/transport.h"

#include <atomic>
#include <chrono>
#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <memory>
#include <thread>
#include <vector>

namespace {

[[noreturn]] void fail(const char* what) {
    std::fprintf(stderr, "test_loopback: FAIL: %s\n", what);
    std::abort();
}

void check(bool condition, const char* what) {
    if (!condition) {
        fail(what);
    }
}

template <typename T> void expect_result(const char* what, const sc::CodecResult<T>& result) {
    if (sc::is_codec_error(result)) {
        std::fprintf(stderr, "test_loopback: FAIL: %s: %s\n", what, sc::codec_error(result).message.c_str());
        std::abort();
    }
}

constexpr int kWidth = 64;
constexpr int kHeight = 64;
constexpr int kFrameCount = 10;
constexpr uint64_t kNsPerFrame = 40'000'000; // 25 fps

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;
}

// Receiver state driven entirely by the transport's receive thread.
struct ReceiverSink {
    sc::H264Depacketizer depacketizer;
    std::unique_ptr<sc::Decoder> decoder;
    std::atomic<int> decoded_frames{0};
    bool saw_frame = false;
    int last_width = 0;
    int last_height = 0;

    void on_packet(sc::RtpPacket packet) {
        std::optional<std::vector<std::byte>> access_unit = depacketizer.depacketize(packet);
        if (!access_unit.has_value()) {
            return;
        }

        sc::EncodedFrame encoded;
        encoded.data = std::move(*access_unit);
        encoded.rtp_timestamp = packet.header.timestamp;

        auto decoded_result = decoder->decode(encoded);
        if (sc::is_codec_error(decoded_result)) {
            return;
        }
        for (const sc::DecodedFrame& decoded : sc::codec_value(decoded_result)) {
            last_width = decoded.width;
            last_height = decoded.height;
            saw_frame = decoded.width > 0 && !decoded.rgba_pixels.empty();
            decoded_frames.fetch_add(1);
        }
    }
};

} // namespace

int main() {
    // Encoder on the sender side.
    sc::EncoderConfig encoder_config;
    encoder_config.codec_name = "h264";
    encoder_config.width = kWidth;
    encoder_config.height = kHeight;
    encoder_config.frame_rate_num = 25;
    encoder_config.frame_rate_den = 1;
    encoder_config.bitrate_kbps = 2000;

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

    // The receiver decodes purely from the bitstream (in-band SPS/PPS).
    sc::DecoderConfig decoder_config;
    decoder_config.codec_name = "h264";
    auto decoder_result = sc::DecoderFactory::create(decoder_config);
    expect_result("decoder create", decoder_result);

    ReceiverSink sink;
    sink.decoder = std::move(sc::codec_value(decoder_result));

    // Bind the receiver to the first free port in a small range.
    auto receiver_transport = sc::RtpTransportFactory::create();
    uint16_t bound_port = 0;
    for (uint16_t port = 45904; port < 45934; ++port) {
        if (receiver_transport->start(sc::Endpoint{"0.0.0.0", port},
                                      [&sink](sc::RtpPacket packet) { sink.on_packet(std::move(packet)); })) {
            bound_port = port;
            break;
        }
    }
    check(bound_port != 0, "receiver transport bound");

    // Sender transport: no bind (OS picks the source port), peer is loopback.
    auto sender_transport = sc::RtpTransportFactory::create();
    check(sender_transport->start(sc::Endpoint{"", 0}, [](sc::RtpPacket) {}), "sender transport started");
    sender_transport->set_peer(sc::Endpoint{"127.0.0.1", bound_port});

    sc::H264Packetizer packetizer;
    for (uint32_t index = 0; index < kFrameCount; ++index) {
        const sc::CapturedFrame frame = make_frame(index);

        auto encoded_result = encoder->encode(frame);
        expect_result("encode", encoded_result);

        for (const sc::EncodedFrame& encoded : sc::codec_value(encoded_result)) {
            for (const sc::RtpPacket& packet : packetizer.packetize(encoded)) {
                check(sender_transport->send(packet), "send packet");
            }
        }
    }

    // Wait (bounded) for the receiver to decode the frames.
    const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5);
    while (sink.decoded_frames.load() < kFrameCount && std::chrono::steady_clock::now() < deadline) {
        std::this_thread::sleep_for(std::chrono::milliseconds(10));
    }

    check(sink.decoded_frames.load() >= kFrameCount / 2, "most frames decoded");
    check(sink.saw_frame, "decoded frame has pixels");
    check(sink.last_width == kWidth && sink.last_height == kHeight, "decoded dimensions");

    sender_transport->stop();
    receiver_transport->stop();

    std::printf("test_loopback: %d frames decoded over UDP loopback\n", sink.decoded_frames.load());
    return 0;
}