#include "prism_engineering/allocator_lab.hpp"
#include <algorithm>
#include <chrono>
#include <limits>
#include <new>
#include <stdexcept>
#include <vector>

namespace prism_engineering {
void* CountingResource::do_allocate(std::size_t bytes, std::size_t alignment) {
  if (bytes > std::numeric_limits<std::size_t>::max()-counts_.requested_bytes ||
      counts_.calls == std::numeric_limits<std::size_t>::max()) throw std::bad_alloc{};
  auto* result = upstream_.allocate(bytes, alignment);
  ++counts_.calls;
  counts_.requested_bytes += bytes;
  counts_.live_bytes += bytes;
  counts_.peak_live_bytes = std::max(counts_.peak_live_bytes,counts_.live_bytes);
  return result;
}
void CountingResource::do_deallocate(void* pointer, std::size_t bytes, std::size_t alignment) {
  upstream_.deallocate(pointer,bytes,alignment);
  counts_.live_bytes -= bytes; // A caller must satisfy memory_resource's deallocation contract.
}
namespace {
struct Accumulation { std::uint32_t r{},g{},b{}; };
constexpr std::uint32_t extent=64, tile_size=16, sample_count=4;
// Four positions in each pixel, on a 128x128 subpixel grid. An analytic diagonal
// divides a blue background from a red triangle; no image decoder or RNG needed.
Rgb sample_scene(std::uint32_t x, std::uint32_t y) noexcept {
  return x+y < 128U ? Rgb{240,64,32} : Rgb{24,48,160};
}
template<class MakeTile>
PreparedImage render_impl(MakeTile make_tile) {
  PreparedImage output{extent,extent,std::vector<Rgb>(extent*extent)};
  for (std::uint32_t top=0; top<extent; top+=tile_size) {
    for (std::uint32_t left=0; left<extent; left+=tile_size) {
      auto accumulation=make_tile();
      accumulation.resize(tile_size*tile_size);
      for (std::uint32_t y=0; y<tile_size; ++y) {
        for (std::uint32_t x=0; x<tile_size; ++x) {
          auto& sum=accumulation[static_cast<std::size_t>(y)*tile_size+x];
          for (std::uint32_t dy=0; dy<2; ++dy) for (std::uint32_t dx=0; dx<2; ++dx) {
            const auto sample=sample_scene(2U*(left+x)+dx,2U*(top+y)+dy);
            sum.r+=sample.r; sum.g+=sample.g; sum.b+=sample.b;
          }
          output.pixels[static_cast<std::size_t>(top+y)*extent+left+x]=Rgb{
            static_cast<std::uint8_t>(sum.r/sample_count),static_cast<std::uint8_t>(sum.g/sample_count),static_cast<std::uint8_t>(sum.b/sample_count)};
        }
      }
    } // Tile accumulation dies here; its resource remains alive.
  }
  return output;
}
template<class Resource, class Release>
RenderMeasurement measure(CountingResource& upstream, Resource& resource, Release release) {
  CountingResource requests{static_cast<std::pmr::memory_resource&>(resource)};
  const auto before=std::chrono::steady_clock::now();
  auto image=render_tiles(requests);
  const auto after=std::chrono::steady_clock::now();
  const auto retained=upstream.counts();
  release(); // No temporary PMR vector remains alive here.
  return {std::move(image),requests.counts(),retained,upstream.counts(),
          std::chrono::duration_cast<std::chrono::nanoseconds>(after-before).count()};
}
}
PreparedImage render_tiles(std::pmr::memory_resource& scratch) {
  return render_impl([&scratch] { return std::pmr::vector<Accumulation>{&scratch}; });
}
PreparedImage render_plain() {
  return render_impl([] { return std::vector<Accumulation>{}; });
}
RenderMeasurement measure_render(ResourceChoice choice) {
  CountingResource upstream{*std::pmr::new_delete_resource()};
  switch(choice) {
  case ResourceChoice::ordinary: return measure(upstream,upstream,[] {});
  case ResourceChoice::monotonic: {
    std::pmr::monotonic_buffer_resource arena{&upstream};
    return measure(upstream,arena,[&arena] { arena.release(); });
  }
  case ResourceChoice::pool: {
    std::pmr::unsynchronized_pool_resource pool{std::pmr::pool_options{4,4096},&upstream};
    return measure(upstream,pool,[&pool] { pool.release(); });
  }
  }
  throw std::invalid_argument{"unknown resource choice"};
}
std::uint64_t image_checksum(const PreparedImage& image) noexcept {
  std::uint64_t sum{};
  for (const auto pixel:image.pixels) sum+=pixel.r+pixel.g+pixel.b;
  return sum;
}
} // namespace prism_engineering