#include "vdm/io/write_buffer.hpp"

#include <atomic>
#include <cstdlib>
#include <cstring>
#include <new>
#include <string>
#include <vector>

#include "vtest.hpp"

using vdm::ConstByteSpan;
using vdm::Error;
using vdm::Result;
using vdm::io::WriteBuffer;

// --- global allocation counter, for the "no alloc in append()" test -----------------
namespace {
std::atomic<long> g_alloc_calls{0};
std::atomic<bool> g_count_allocs{false};
}  // namespace

void *operator new(std::size_t n) {
    if (g_count_allocs.load(std::memory_order_relaxed))
        g_alloc_calls.fetch_add(1, std::memory_order_relaxed);
    void *p = std::malloc(n ? n : 1);
    if (!p)
        throw std::bad_alloc();
    return p;
}
void operator delete(void *p) noexcept {
    std::free(p);
}
void operator delete(void *p, std::size_t) noexcept {
    std::free(p);
}
void *operator new[](std::size_t n) {
    return ::operator new(n);
}
void operator delete[](void *p) noexcept {
    std::free(p);
}
void operator delete[](void *p, std::size_t) noexcept {
    std::free(p);
}

namespace {

// A flush sink that records (offset, bytes) and never allocates after construction.
struct Sink {
    std::vector<std::byte> data;      // pre-reserved
    std::vector<std::uint64_t> offs;  // pre-reserved
    std::vector<std::size_t> lens;
    bool fail_next = false;

    WriteBuffer::FlushFn fn() {
        return [this](std::uint64_t off, ConstByteSpan s) -> Result<void> {
            if (fail_next) {
                fail_next = false;
                return vdm::Err{Error::io_error, "sink forced failure"};
            }
            offs.push_back(off);
            lens.push_back(s.size());
            data.insert(data.end(), s.begin(), s.end());
            return vdm::ok();
        };
    }
};

ConstByteSpan sv(const char *s) {
    return {reinterpret_cast<const std::byte *>(s), std::strlen(s)};
}

}  // namespace

VT_TEST(wb_accumulates_then_flushes_on_fill) {
    Sink sink;
    sink.data.reserve(1 << 16);
    sink.offs.reserve(64);
    sink.lens.reserve(64);

    WriteBuffer wb(0, 8, sink.fn());
    VT_CHECK(wb.append(sv("abc")).has_value());  // 3 buffered
    VT_CHECK_EQ(wb.pending(), 3u);
    VT_CHECK(sink.offs.empty());  // no flush yet

    VT_CHECK(wb.append(sv("defgh")).has_value());  // fills to 8 -> flush
    VT_REQUIRE(sink.offs.size() == 1);
    VT_CHECK_EQ(sink.offs[0], 0u);
    VT_CHECK_EQ(sink.lens[0], 8u);
    VT_CHECK_EQ(wb.pending(), 0u);
    VT_CHECK_EQ(wb.next_offset(), 8u);

    VT_CHECK(wb.append(sv("ij")).has_value());
    VT_CHECK(wb.flush().has_value());  // explicit tail flush
    VT_REQUIRE(sink.offs.size() == 2);
    VT_CHECK_EQ(sink.offs[1], 8u);
    VT_CHECK_EQ(sink.lens[1], 2u);

    VT_CHECK_EQ(std::string(reinterpret_cast<const char *>(sink.data.data()), sink.data.size()),
                std::string("abcdefghij"));
    VT_CHECK_EQ(wb.total_appended(), 10u);
}

VT_TEST(wb_flush_is_noop_when_empty) {
    Sink sink;
    sink.offs.reserve(4);
    WriteBuffer wb(100, 16, sink.fn());
    VT_CHECK(wb.flush().has_value());
    VT_CHECK(sink.offs.empty());
}

VT_TEST(wb_oversized_chunk_writes_through) {
    Sink sink;
    sink.data.reserve(1 << 16);
    sink.offs.reserve(16);
    sink.lens.reserve(16);

    WriteBuffer wb(0, 8, sink.fn());
    VT_CHECK(wb.append(sv("ab")).has_value());  // 2 buffered
    // 20 bytes arriving: buffer isn't empty, so first 6 top it off + flush(8), then the
    // remaining 14 (>= capacity, buffer now empty) write straight through.
    std::string big(20, 'x');
    VT_CHECK(wb.append(sv(big.c_str())).has_value());
    VT_CHECK(wb.flush().has_value());

    // reconstruct
    std::string got(reinterpret_cast<const char *>(sink.data.data()), sink.data.size());
    VT_CHECK_EQ(got, std::string("ab") + big);
    VT_CHECK_EQ(wb.total_appended(), 22u);
    // one full-buffer flush + one passthrough; order preserved
    VT_CHECK(sink.offs.size() >= 2);
    VT_CHECK_EQ(sink.offs.front(), 0u);
}

VT_TEST(wb_exact_capacity_chunk_from_empty_writes_through) {
    Sink sink;
    sink.data.reserve(64);
    sink.offs.reserve(4);
    sink.lens.reserve(4);
    WriteBuffer wb(0, 4, sink.fn());
    VT_CHECK(wb.append(sv("wxyz")).has_value());  // == capacity, empty -> passthrough
    VT_REQUIRE(sink.offs.size() == 1);
    VT_CHECK_EQ(sink.lens[0], 4u);
    VT_CHECK_EQ(wb.pending(), 0u);
}

VT_TEST(wb_flush_error_propagates_without_advancing_durable_offset) {
    Sink sink;
    sink.data.reserve(64);
    sink.offs.reserve(4);
    sink.lens.reserve(4);
    WriteBuffer wb(0, 8, sink.fn());

    VT_CHECK(wb.append(sv("abc")).has_value());  // 3 buffered, nothing durable yet
    VT_CHECK_EQ(wb.next_offset(), 0u);

    sink.fail_next = true;
    auto r = wb.flush();  // forced failure
    VT_REQUIRE(!r.has_value());
    VT_CHECK_EQ(r.error().code, Error::io_error);
    VT_CHECK_EQ(wb.next_offset(), 0u);  // durable offset did NOT move
    VT_CHECK(sink.offs.empty());

    // a retry flush succeeds and advances
    VT_CHECK(wb.flush().has_value());
    VT_CHECK_EQ(wb.next_offset(), 3u);
    VT_REQUIRE(sink.lens.size() == 1);
    VT_CHECK_EQ(sink.lens[0], 3u);
}

VT_TEST(wb_append_does_not_allocate) {
    // flush sink that never allocates: just sum sizes.
    std::atomic<std::uint64_t> total{0};
    auto flush = [&total](std::uint64_t, ConstByteSpan s) -> Result<void> {
        total.fetch_add(s.size());
        return vdm::ok();
    };
    WriteBuffer wb(0, 4096, flush);

    // Warm up (any first-call lazy init happens now, outside the measured window).
    std::string warm(100, 'w');
    (void)wb.append(sv(warm.c_str()));
    (void)wb.flush();

    std::vector<std::byte> chunk(512, std::byte{7});
    g_alloc_calls.store(0);
    g_count_allocs.store(true);
    for (int i = 0; i < 5000; ++i) {
        auto r = wb.append(ConstByteSpan(chunk.data(), 137 + (i % 200)));
        if (!r.has_value()) {
            g_count_allocs.store(false);
            VT_FAIL("append failed");
            return;
        }
    }
    (void)wb.flush();
    g_count_allocs.store(false);

    VT_CHECK_EQ(g_alloc_calls.load(), 0L);
    VT_CHECK(total.load() > 0);
}