#include "prism/memory_lab.hpp"
#include <algorithm>
#include <array>
#include <limits>
#include <new>
#include <vector>

namespace prism {
void* CountingResource::do_allocate(std::size_t bytes, std::size_t alignment) {
    if (bytes > std::numeric_limits<std::size_t>::max() - live_ ||
        allocations_ == std::numeric_limits<std::size_t>::max()) throw std::bad_alloc{};
    // Commit accounting after upstream succeeds; an injected failure changes
    // neither live bytes nor successful-allocation count.
    void* pointer = upstream_.allocate(bytes, alignment);
    ++allocations_;
    live_ += bytes;
    peak_ = std::max(peak_, live_);
    return pointer;
}
void CountingResource::do_deallocate(void* pointer, std::size_t bytes, std::size_t alignment) {
    upstream_.deallocate(pointer, bytes, alignment);
    live_ -= bytes; // memory_resource requires the original size/alignment.
}

MemoryReport run_memory_lab() {
    alignas(std::max_align_t) std::array<std::byte, 512> storage{};
    // Null upstream is a hard allocation boundary: no fallback to the heap.
    std::pmr::monotonic_buffer_resource arena{storage.data(), storage.size(), std::pmr::null_memory_resource()};
    CountingResource counted{arena};
    MemoryReport report;
    {
        std::pmr::vector<int> samples{&counted}; // Borrowing container dies first.
        samples.reserve(16);
        for (int value = 0; value < 16; ++value) samples.push_back(value);
        std::pmr::string temporary{"A title promoted beyond the arena lifetime", &counted};
        report.promoted_title.assign(temporary.data(), temporary.size());
        const auto before = counted.live_bytes();
        try {
            samples.reserve(4096); // Deliberately exceeds the fixed storage.
        } catch (const std::bad_alloc&) {
            report.bounded_failure = samples.size() == 16 && samples[15] == 15 && counted.live_bytes() == before;
        }
    } // Both PMR containers release their logical allocations here.
    report.allocations = counted.allocations();
    report.peak_bytes = counted.peak_bytes();
    report.live_after_scope = counted.live_bytes();
    arena.release(); // Safe only after all arena-backed containers have died.
    // report.promoted_title owns separate storage; no stale arena view is read.
    return report;
}
} // namespace prism