#include "VoxelStreamingAdmissionOrder.h"

#include <algorithm>
#include <cstdlib>
#include <iostream>
#include <vector>

struct FCandidate
{
    int Priority;
    int Id;
    int Level;
};

static bool LessCandidate(const FCandidate& A, const FCandidate& B)
{
    if (A.Priority != B.Priority) return A.Priority < B.Priority;
    return A.Id < B.Id;
}

static std::vector<int> RunPass(std::vector<FCandidate>& Candidates,
                                int InitiallyActive, int Capacity)
{
    int Active = InitiallyActive;
    bool FreedOneSlot = false;
    std::vector<int> Submitted;
    VoxelStreamingAdmissionOrder::AdmitInOrder(
        Candidates,
        [&]() { return Active >= Capacity; },
        [&](const FCandidate& Candidate)
        {
            if (Candidate.Level != 0)
            {
                std::cerr << "candidate escaped its level pass\n";
                std::exit(2);
            }
            Submitted.push_back(Candidate.Id);
            ++Active;
            if (!FreedOneSlot && Submitted.size() == 1)
            {
                --Active; // A worker finishes one task while this level pass is still running.
                FreedOneSlot = true;
            }
        },
        []() {});
    return Submitted;
}

int main()
{
    const std::vector<FCandidate> Input = {
        {4, 4, 0}, {0, 0, 0}, {3, 3, 0}, {1, 1, 0}, {2, 2, 0}};
    constexpr int InitiallyActive = 1;
    constexpr int Capacity = 2;
    const int PassStartFreeSlots = Capacity - InitiallyActive;

    std::vector<FCandidate> FullSort = Input;
    std::stable_sort(FullSort.begin(), FullSort.end(), LessCandidate);
    const std::vector<int> Expected = {FullSort[0].Id, FullSort[1].Id};

    std::vector<FCandidate> FixedCandidates = Input;
    VoxelStreamingAdmissionOrder::StableSortCandidates(
        FixedCandidates.begin(), FixedCandidates.end(), LessCandidate);
    const std::vector<int> FixedSubmitted =
        RunPass(FixedCandidates, InitiallyActive, Capacity);
    if (FixedSubmitted != Expected || FixedSubmitted.size() != 2
        || FixedCandidates[1].Level != FixedCandidates[0].Level)
    {
        std::cerr << "fixed submission order/refill mismatch\n";
        return 1;
    }

    // Negative control: the old pass-start partial_sort + SetNum truncation loses the
    // second candidate even though a slot frees after the first submission.
    std::vector<FCandidate> OldCandidates = Input;
    const std::size_t InitialCount = std::min(
        OldCandidates.size(), static_cast<std::size_t>(PassStartFreeSlots));
    std::partial_sort(OldCandidates.begin(), OldCandidates.begin() + InitialCount,
                      OldCandidates.end(), LessCandidate);
    OldCandidates.resize(InitialCount);
    const std::vector<int> OldSubmitted =
        RunPass(OldCandidates, InitiallyActive, Capacity);
    if (OldSubmitted == Expected)
    {
        std::cerr << "negative control did not reproduce the lost refill\n";
        return 1;
    }

    std::cout << "[StreamingAdmissionOrder] full_order=";
    for (const int Id : Expected) std::cout << Id;
    std::cout << " fixed_submissions=";
    for (const int Id : FixedSubmitted) std::cout << Id;
    std::cout << " old_submissions=";
    for (const int Id : OldSubmitted) std::cout << Id;
    std::cout << " same_level_refill=PASS negative_control=PASS status=PASS\n";
    return 0;
}