#include "vdm/segment/segmenter.hpp"

#include <atomic>
#include <cstdint>
#include <thread>
#include <vector>

#include "vtest.hpp"

using namespace vdm::segment;

namespace {
constexpr std::uint64_t MiB = 1u << 20;

// Assign `n` slots (bounded by what the segmenter hands out) and return the indices.
std::vector<std::uint32_t> fill(Segmenter &s, int n) {
    std::vector<std::uint32_t> idx;
    for (int i = 0; i < n; ++i) {
        auto a = s.assign_slot();
        if (!a)
            break;
        idx.push_back(*a);
    }
    return idx;
}

// Do the segments (by their [start,end] at this instant) tile [0,total) with no overlap?
bool tiles_exactly(const Segmenter &s) {
    auto snap = s.snapshot();
    std::vector<std::pair<std::uint64_t, std::uint64_t>> r;
    for (auto &v : snap)
        if (v.state != SegState::failed)
            r.emplace_back(v.start, v.end);
    std::sort(r.begin(), r.end());
    std::uint64_t cursor = 0;
    for (auto [a, b] : r) {
        if (a != cursor)
            return false;
        cursor = b + 1;
    }
    return cursor == s.total_size();
}
}  // namespace

VT_TEST(seg_target_count_clamps) {
    VT_CHECK_EQ(Segmenter(100 * MiB, 8, true).target_segment_count(), 8u);
    VT_CHECK_EQ(Segmenter(100 * MiB, 64, true).target_segment_count(), 32u);  // max 32
    VT_CHECK_EQ(Segmenter(100 * MiB, 0, true).target_segment_count(), 8u);    // default
    VT_CHECK_EQ(Segmenter(3 * MiB + 1, 8, true).target_segment_count(), 3u);  // total/min
    VT_CHECK_EQ(Segmenter(100 * MiB, 8, false).target_segment_count(), 1u);   // non-resumable
    VT_CHECK_EQ(Segmenter(0, 8, true).target_segment_count(), 1u);            // chunked
}

VT_TEST(seg_non_resumable_is_one_segment) {
    Segmenter s(50 * MiB, 8, false);
    auto idx = fill(s, 8);
    VT_REQUIRE(idx.size() == 1);
    auto snap = s.snapshot();
    VT_REQUIRE(snap.size() == 1);
    VT_CHECK_EQ(snap[0].start, 0u);
    VT_CHECK_EQ(snap[0].end, 50u * MiB - 1);
}

VT_TEST(seg_initial_split_covers_range) {
    Segmenter s(80 * MiB, 8, true);
    auto idx = fill(s, 8);
    VT_CHECK_EQ(idx.size(), 8u);
    VT_CHECK(tiles_exactly(s));
    // no segment below the 1 MiB floor
    for (auto &v : s.snapshot())
        VT_CHECK(v.length() >= MiB);
}

VT_TEST(seg_split_stops_at_min_floor) {
    // 5 MiB, floor 1 MiB, ask for 8: only ~5 splits possible (each half >= 1 MiB needs
    // the parent >= 2 MiB), so we get fewer than 8.
    Segmenter s(5 * MiB, 8, true);
    auto idx = fill(s, 8);
    VT_CHECK(idx.size() >= 1 && idx.size() <= 5);
    VT_CHECK(tiles_exactly(s));
}

VT_TEST(seg_steal_takes_second_half_of_largest_remaining) {
    Segmenter s(80 * MiB, 4, true);
    auto idx = fill(s, 4);
    VT_REQUIRE(idx.size() == 4);

    // Spread progress unevenly. Index != file position after splits, so identify the
    // largest-remaining segment by scanning the snapshot, not by index.
    s.advance(idx[1], 3 * MiB);
    s.advance(idx[2], 7 * MiB);
    s.advance(idx[3], 12 * MiB);

    SegmentView pre_victim{};
    std::uint64_t worst = 0;
    for (auto &v : s.snapshot())
        if (v.index != idx[0] && v.remaining() > worst) {
            worst = v.remaining();
            pre_victim = v;
        }

    auto cont = s.on_complete(idx[0], /*may_steal=*/true);
    VT_REQUIRE(cont.has_value());

    SegmentView victim{}, fresh{};
    for (auto &v : s.snapshot()) {
        if (v.index == pre_victim.index)
            victim = v;
        if (v.index == *cont)
            fresh = v;
    }
    VT_CHECK_EQ(fresh.end, pre_victim.end);    // fresh takes the tail of the victim's range
    VT_CHECK_EQ(victim.end + 1, fresh.start);  // contiguous, no gap / no overlap
    VT_CHECK(victim.end < pre_victim.end);     // the victim really did shrink
    VT_CHECK(fresh.length() >= MiB);
    VT_CHECK(victim.remaining() >= MiB);
    // fresh got roughly the back half of what was remaining
    VT_CHECK(fresh.length() >= worst / 2 - MiB && fresh.length() <= worst / 2 + MiB);
    VT_CHECK(tiles_exactly(s));
}

VT_TEST(seg_complete_without_steal_releases) {
    Segmenter s(4 * MiB, 2, true);
    auto idx = fill(s, 2);  // 2 x 2 MiB
    VT_REQUIRE(idx.size() == 2);
    s.advance(idx[1], 2 * MiB);
    s.set_segment_state(idx[1], SegState::complete);
    // idx[0] done, nothing left worth >= 1 MiB to steal -> release
    s.advance(idx[0], 2 * MiB);
    auto cont = s.on_complete(idx[0], true);
    VT_CHECK(!cont.has_value());
}

VT_TEST(seg_yield_returns_nullopt) {
    Segmenter s(80 * MiB, 4, true);
    auto idx = fill(s, 4);
    s.advance(idx[0], 20 * MiB);
    auto cont = s.on_complete(idx[0], /*may_steal=*/false);  // yielding
    VT_CHECK(!cont.has_value());
}

VT_TEST(seg_third_connection_failure_with_mirror_requeues) {
    Segmenter s(40 * MiB, 2, true);
    auto idx = fill(s, 2);
    s.advance(idx[0], 4 * MiB);

    VT_CHECK(s.on_failed(idx[0], /*conn=*/true, /*mirror=*/true) == FailAction::retry);
    VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::retry);
    VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::requeue);  // 3rd

    VT_CHECK_EQ(s.segment_state(idx[0]), SegState::failed);
    // the orphaned tail is now assignable again
    auto again = s.assign_slot();
    VT_REQUIRE(again.has_value());
    auto snap = s.snapshot();
    SegmentView reborn{};
    for (auto &v : snap)
        if (v.index == *again)
            reborn = v;
    VT_CHECK_EQ(reborn.start, 4u * MiB);     // resumes where the failed one stopped
    VT_CHECK_EQ(reborn.end, 20u * MiB - 1);  // its half of the file
}

VT_TEST(seg_failure_without_mirror_always_retries) {
    Segmenter s(40 * MiB, 2, true);
    auto idx = fill(s, 2);
    for (int i = 0; i < 6; ++i)
        VT_CHECK(s.on_failed(idx[0], true, /*mirror=*/false) == FailAction::retry);
    // a non-connection error also retries regardless of count
    VT_CHECK(s.on_failed(idx[1], /*conn=*/false, /*mirror=*/true) == FailAction::retry);
}

VT_TEST(seg_note_connected_resets_failure_count) {
    Segmenter s(40 * MiB, 2, true);
    auto idx = fill(s, 2);
    s.on_failed(idx[0], true, true);
    s.on_failed(idx[0], true, true);
    s.note_connected(idx[0]);
    VT_CHECK(s.on_failed(idx[0], true, true) == FailAction::retry);  // count restarted
}

VT_TEST(seg_resume_from_meta_table) {
    std::vector<ResumedRange> table = {
        {0, 9 * MiB - 1, 9 * MiB},         // fully done
        {9 * MiB, 19 * MiB - 1, 3 * MiB},  // partial
        {19 * MiB, 40 * MiB - 1, 0},       // untouched
    };
    Segmenter s(40 * MiB, 8, table, true);
    auto snap = s.snapshot();
    VT_REQUIRE(snap.size() == 3);
    VT_CHECK_EQ(snap[0].state, SegState::complete);
    VT_CHECK_EQ(snap[1].completed, 3u * MiB);
    VT_CHECK_EQ(s.downloaded(), 12u * MiB);
    VT_CHECK(tiles_exactly(s));

    // assign hands out the two incomplete ranges before splitting
    auto a = s.assign_slot();
    auto b = s.assign_slot();
    VT_REQUIRE(a && b);
}

VT_TEST(seg_resume_from_bad_table_falls_back) {
    std::vector<ResumedRange> gappy = {{0, 4 * MiB - 1, 0}, {8 * MiB, 40 * MiB - 1, 0}};
    Segmenter s(40 * MiB, 8, gappy, true);
    VT_CHECK(s.snapshot().empty());  // lazy fresh layout
    auto idx = fill(s, 8);
    VT_CHECK(idx.size() >= 1);
    VT_CHECK(tiles_exactly(s));
}

VT_TEST(seg_all_complete_and_downloaded) {
    Segmenter s(8 * MiB, 4, true);
    auto idx = fill(s, 4);
    VT_CHECK(!s.all_complete());
    for (auto i : idx) {
        std::uint64_t len = s.segment_end(i) - s.segment_start(i) + 1;
        s.advance(i, len);
        s.set_segment_state(i, SegState::complete);
    }
    VT_CHECK(s.all_complete());
    VT_CHECK_EQ(s.downloaded(), 8u * MiB);
}

// --- the steal path under the sanitizers -----------------------------------------------

VT_TEST(seg_concurrent_steal_and_advance) {
    constexpr std::uint64_t total = 64 * MiB;
    Segmenter s(total, 8, true);
    auto idx = fill(s, 8);
    VT_REQUIRE(idx.size() == 8);

    std::vector<std::jthread> workers;
    for (std::uint32_t w = 0; w < 8; ++w) {
        workers.emplace_back([&s, seg = idx[w]]() mutable {
            std::uint32_t cur = seg;
            for (int guard = 0; guard < 200000; ++guard) {
                const std::uint64_t start = s.segment_start(cur);
                const std::uint64_t end = s.segment_end(cur);  // may shrink under a steal
                const std::uint64_t len = end - start + 1;
                const std::uint64_t done = s.segment_completed(cur);
                if (done >= len) {
                    auto nxt = s.on_complete(cur, /*may_steal=*/true);
                    if (!nxt)
                        return;  // nothing left to steal — this worker is finished
                    cur = *nxt;
                    continue;
                }
                s.advance(cur, std::min(done + 128 * 1024, len));
            }
        });
    }
    workers.clear();  // join

    VT_CHECK(s.all_complete());
    VT_CHECK_EQ(s.downloaded(), total);
    VT_CHECK(tiles_exactly(s));
}