/*
 * Copyright (c) 2026, Gregory Bertilson <gregory@ladybird.org>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/ByteBuffer.h>
#include <LibCore/EventLoop.h>
#include <LibMedia/IncrementallyPopulatedStream.h>
#include <LibTest/TestCase.h>
#include <LibThreading/Thread.h>

static ByteBuffer make_test_data(size_t size)
{
    auto buffer = MUST(ByteBuffer::create_uninitialized(size));
    for (size_t i = 0; i < size; i++)
        buffer[i] = static_cast<u8>(i);
    return buffer;
}

TEST_CASE(create_empty)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    EXPECT(!stream->expected_size().has_value());

    stream->set_expected_size(500);
    EXPECT(stream->expected_size().has_value());
    EXPECT_EQ(stream->expected_size().value(), 500u);
}

TEST_CASE(create_from_data_and_buffer)
{
    auto data = make_test_data(256);

    auto stream = Media::IncrementallyPopulatedStream::create_from_buffer(data);
    EXPECT(stream->expected_size().has_value());
    EXPECT_EQ(stream->expected_size().value(), 256u);
    EXPECT_EQ(stream->size(), 256u);
}

TEST_CASE(cursor_seek_modes)
{
    auto data = make_test_data(100);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());
    auto cursor = stream->create_cursor();

    EXPECT_EQ(cursor->position(), 0u);
    EXPECT_EQ(cursor->size(), 100u);

    MUST(cursor->seek(50, SeekMode::SetPosition));
    EXPECT_EQ(cursor->position(), 50u);

    MUST(cursor->seek(10, SeekMode::FromCurrentPosition));
    EXPECT_EQ(cursor->position(), 60u);

    MUST(cursor->seek(-10, SeekMode::FromEndPosition));
    EXPECT_EQ(cursor->position(), 90u);
}

TEST_CASE(cursor_read_operations)
{
    auto data = make_test_data(100);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());
    auto cursor = stream->create_cursor();

    Array<u8, 10> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 10u);
    EXPECT_EQ(cursor->position(), 10u);
    for (size_t i = 0; i < 10; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));

    MUST(cursor->seek(50, SeekMode::SetPosition));
    MUST(cursor->read_into(buffer));
    for (size_t i = 0; i < 10; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(50 + i));

    MUST(cursor->seek(95, SeekMode::SetPosition));
    bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 5u);
    for (size_t i = 0; i < 5; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(95 + i));

    MUST(cursor->seek(100, SeekMode::SetPosition));
    auto result = cursor->read_into(buffer);
    EXPECT(result.is_error());
    EXPECT_EQ(result.error().category(), Media::DecoderErrorCategory::EndOfStream);

    MUST(cursor->seek(0, SeekMode::SetPosition));
    bytes_read = MUST(cursor->read_into(buffer.span().trim(0)));
    EXPECT_EQ(bytes_read, 0u);
    EXPECT_EQ(cursor->position(), 0u);
}

TEST_CASE(sequential_reads)
{
    auto data = make_test_data(256);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());
    auto cursor = stream->create_cursor();

    for (size_t i = 0; i < 256; i += 16) {
        Array<u8, 16> buffer;
        auto bytes_read = MUST(cursor->read_into(buffer));
        EXPECT_EQ(bytes_read, 16u);

        for (size_t j = 0; j < 16; j++)
            EXPECT_EQ(buffer[j], static_cast<u8>(i + j));
    }

    EXPECT_EQ(cursor->position(), 256u);
}

TEST_CASE(multiple_cursors_independent)
{
    auto data = make_test_data(100);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());
    auto cursor1 = stream->create_cursor();
    auto cursor2 = stream->create_cursor();

    MUST(cursor1->seek(10, SeekMode::SetPosition));
    MUST(cursor2->seek(50, SeekMode::SetPosition));

    EXPECT_EQ(cursor1->position(), 10u);
    EXPECT_EQ(cursor2->position(), 50u);

    Array<u8, 5> buffer1;
    Array<u8, 5> buffer2;
    MUST(cursor1->read_into(buffer1));
    MUST(cursor2->read_into(buffer2));

    for (size_t i = 0; i < 5; i++) {
        EXPECT_EQ(buffer1[i], static_cast<u8>(10 + i));
        EXPECT_EQ(buffer2[i], static_cast<u8>(50 + i));
    }
}

TEST_CASE(add_chunks_incrementally)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();

    constexpr size_t data_size = 100;
    auto data = make_test_data(data_size);
    stream->add_chunk_at(0, data.bytes().trim(50));

    stream->add_chunk_at(50, data.bytes().slice(50));
    stream->close();

    EXPECT(stream->expected_size().has_value());
    EXPECT_EQ(stream->expected_size().value(), data_size);

    auto cursor = stream->create_cursor();
    Array<u8, data_size> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));

    EXPECT_EQ(bytes_read, data_size);
    for (size_t i = 0; i < data_size; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));
}

TEST_CASE(add_overlapping_chunks)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();

    constexpr size_t data_size = 100;
    auto data = make_test_data(data_size);
    stream->add_chunk_at(0, data.bytes().trim(50));
    stream->add_chunk_at(40, data.bytes().slice(40));

    auto cursor = stream->create_cursor();
    Array<u8, data_size> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));

    EXPECT_EQ(bytes_read, data_size);
    for (size_t i = 0; i < data_size; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));
}

TEST_CASE(add_chunk_at_offset)
{
    Core::EventLoop loop;

    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    stream->set_expected_size(100);
    stream->set_data_request_callback([](u64) { });

    auto data = make_test_data(80);
    stream->add_chunk_at(0, data.bytes().trim(30));
    stream->add_chunk_at(50, data.bytes().slice(50));

    auto cursor = stream->create_cursor();
    MUST(cursor->seek(50, SeekMode::SetPosition));

    Array<u8, 30> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));

    EXPECT_EQ(bytes_read, 30u);
    for (size_t i = 0; i < 30; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(50 + i));
}

TEST_CASE(cursor_abort_and_reset)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    stream->set_expected_size(100);

    auto cursor = stream->create_cursor();

    EXPECT(!cursor->is_blocked());

    IGNORE_USE_IN_ESCAPING_LAMBDA Atomic<bool> read_completed { false };
    IGNORE_USE_IN_ESCAPING_LAMBDA Atomic<bool> was_aborted { false };

    auto thread = Threading::Thread::construct("TestAbort"sv, [&, cursor]() -> intptr_t {
        Array<u8, 10> buffer;
        auto result = cursor->read_into(buffer);
        read_completed = true;
        was_aborted = result.is_error() && result.error().category() == Media::DecoderErrorCategory::Aborted;
        return 0;
    });

    thread->start();

    while (!cursor->is_blocked())
        ;
    EXPECT(cursor->is_blocked());

    cursor->abort();
    MUST(thread->join());

    EXPECT_EQ(cursor->is_blocked(), false);
    EXPECT(read_completed.load());
    EXPECT(was_aborted.load());

    // After aborting a read, reset_abort() should allow us to read again.
    cursor->reset_abort();
    auto data = make_test_data(100);
    stream->add_chunk_at(0, data.bytes());

    Array<u8, 10> buffer;
    auto result = cursor->read_into(buffer);
    EXPECT(!result.is_error());
    EXPECT_EQ(result.value(), 10u);
}

TEST_CASE(redundant_chunk_within_existing_chunk_at_nonzero_offset)
{
    // Regression test: add_chunk_at used to compare chunk.size() (a relative byte count)
    // against new_chunk_end (an absolute file offset). When the existing chunk started at
    // a non-zero offset, chunk.size() < new_chunk_end even if the new data was fully
    // covered, causing the buffer to be shrunk and data beyond the new chunk's end to be lost.
    auto stream = Media::IncrementallyPopulatedStream::create_empty();

    constexpr size_t data_size = 200;
    auto data = make_test_data(data_size);

    // Add a chunk at a non-zero offset covering [100, 120).
    stream->add_chunk_at(100, data.bytes().slice(100, 20));

    // Add a redundant chunk fully within [100, 120), specifically [105, 115).
    // With the bug, this shrinks the existing chunk to [100, 115), losing bytes [115, 120).
    stream->add_chunk_at(105, data.bytes().slice(105, 10));

    stream->close();

    auto cursor = stream->create_cursor();
    MUST(cursor->seek(100, SeekMode::SetPosition));

    Array<u8, 20> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));

    EXPECT_EQ(bytes_read, 20u);
    for (size_t i = 0; i < 20; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(100 + i));
}

TEST_CASE(remove_byte_range_splits_chunk)
{
    auto data = make_test_data(100);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());

    stream->remove_byte_range(30, 70);

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 2u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 30u);
    EXPECT_EQ(ranges[1].start, 70u);
    EXPECT_EQ(ranges[1].end, 100u);

    auto cursor = stream->create_cursor();
    Array<u8, 30> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 30u);
    for (size_t i = 0; i < 30; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));

    MUST(cursor->seek(70, SeekMode::SetPosition));
    bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 30u);
    for (size_t i = 0; i < 30; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(70 + i));
}

TEST_CASE(remove_byte_range_trims_chunk)
{
    auto data = make_test_data(100);
    auto stream = Media::IncrementallyPopulatedStream::create_from_data(data.bytes());

    stream->remove_byte_range(0, 20);
    stream->remove_byte_range(80, 100);

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 20u);
    EXPECT_EQ(ranges[0].end, 80u);

    auto cursor = stream->create_cursor();
    MUST(cursor->seek(20, SeekMode::SetPosition));

    Array<u8, 60> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 60u);
    for (size_t i = 0; i < 60; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(20 + i));
}

TEST_CASE(remove_byte_range_across_disjoint_chunks)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(100);

    stream->add_chunk_at(0, data.bytes().trim(20));
    stream->add_chunk_at(40, data.bytes().slice(40, 20));
    stream->add_chunk_at(80, data.bytes().slice(80));
    stream->remove_byte_range(10, 90);

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 2u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 10u);
    EXPECT_EQ(ranges[1].start, 90u);
    EXPECT_EQ(ranges[1].end, 100u);

    auto cursor = stream->create_cursor();
    Array<u8, 10> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 10u);
    for (size_t i = 0; i < 10; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));

    MUST(cursor->seek(90, SeekMode::SetPosition));
    bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 10u);
    for (size_t i = 0; i < 10; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(90 + i));
}

TEST_CASE(add_touching_chunks_forward)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(100);

    stream->add_chunk_at(0, data.bytes().trim(50));
    stream->add_chunk_at(50, data.bytes().slice(50));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 100u);
}

TEST_CASE(add_touching_chunks_reverse)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(100);

    stream->add_chunk_at(50, data.bytes().slice(50));
    stream->add_chunk_at(0, data.bytes().trim(50));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 100u);

    // Verify the data is contiguous and correct.
    auto cursor = stream->create_cursor();
    Array<u8, 100> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 100u);
    for (size_t i = 0; i < 100; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));
}

TEST_CASE(add_disjoint_chunks)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(100);

    stream->add_chunk_at(0, data.bytes().trim(30));
    stream->add_chunk_at(70, data.bytes().slice(70));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 2u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 30u);
    EXPECT_EQ(ranges[1].start, 70u);
    EXPECT_EQ(ranges[1].end, 100u);
}

TEST_CASE(add_chunk_fills_gap)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(100);

    stream->add_chunk_at(0, data.bytes().trim(30));
    stream->add_chunk_at(70, data.bytes().slice(70));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 2u);

    // Fill the gap.
    stream->add_chunk_at(30, data.bytes().slice(30, 40));

    ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 100u);

    // Verify data integrity.
    auto cursor = stream->create_cursor();
    Array<u8, 100> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 100u);
    for (size_t i = 0; i < 100; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));
}

TEST_CASE(add_chunk_spans_multiple_existing_chunks)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(120);

    stream->add_chunk_at(0, data.bytes().trim(10));
    stream->add_chunk_at(20, data.bytes().slice(20, 10));
    stream->add_chunk_at(40, data.bytes().slice(40, 10));
    stream->add_chunk_at(10, data.bytes().slice(10, 60));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 70u);

    auto cursor = stream->create_cursor();
    Array<u8, 70> buffer;
    auto bytes_read = MUST(cursor->read_into(buffer));
    EXPECT_EQ(bytes_read, 70u);
    for (size_t i = 0; i < 70; i++)
        EXPECT_EQ(buffer[i], static_cast<u8>(i));
}

TEST_CASE(add_three_disjoint_then_connect)
{
    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    auto data = make_test_data(150);

    stream->add_chunk_at(0, data.bytes().trim(30));
    stream->add_chunk_at(60, data.bytes().slice(60, 30));
    stream->add_chunk_at(120, data.bytes().slice(120));

    auto ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 3u);

    // Connect first and second with a chunk that touches both.
    stream->add_chunk_at(30, data.bytes().slice(30, 30));
    ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 2u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 90u);

    // Connect second and third.
    stream->add_chunk_at(90, data.bytes().slice(90, 30));
    ranges = stream->available_byte_ranges();
    EXPECT_EQ(ranges.size(), 1u);
    EXPECT_EQ(ranges[0].start, 0u);
    EXPECT_EQ(ranges[0].end, 150u);
}

TEST_CASE(data_request_callback_invoked)
{
    Core::EventLoop loop;

    // Stream size must be larger than FORWARD_REQUEST_THRESHOLD (1 MiB) to test callback
    static constexpr u64 stream_size = 2 * MiB;
    static constexpr u64 initial_chunk_size = 100;
    static constexpr u64 seek_position = stream_size - 100;

    auto stream = Media::IncrementallyPopulatedStream::create_empty();
    stream->set_expected_size(stream_size);

    // Add initial chunk so the callback logic can be triggered
    auto initial_data = make_test_data(initial_chunk_size);
    stream->add_chunk_at(0, initial_data.bytes());

    bool callback_invoked { false };
    u64 requested_offset { 0 };

    stream->set_data_request_callback([&](u64 offset) {
        auto data = make_test_data(100);
        stream->add_chunk_at(seek_position, data.bytes());
        callback_invoked = true;
        requested_offset = offset;
    });

    auto cursor = stream->create_cursor();
    MUST(cursor->seek(seek_position, SeekMode::SetPosition));

    auto thread = Threading::Thread::construct("TestCallback"sv, [cursor]() -> intptr_t {
        Array<u8, 10> buffer;
        MUST(cursor->read_into(buffer));
        return 0;
    });
    thread->start();

    auto start_time = MonotonicTime::now_coarse();
    while (!callback_invoked) {
        loop.pump(Core::EventLoop::WaitMode::PollForEvents);
        if (MonotonicTime::now_coarse() - start_time > AK::Duration::from_seconds(1))
            break;
    }

    EXPECT(callback_invoked);
    EXPECT(requested_offset >= initial_chunk_size);

    MUST(thread->join());
}