#include "buffer.hh"

#include <algorithm>
#include <cassert>
#include <memory>
#include <limits>

namespace modxml {
namespace sax {

namespace {

class DynamicBuffer : public Buffer {
 public:
  DynamicBuffer(std::size_t default_size, std::size_t max_size)
      : default_size_(std::min(default_size, max_size)), max_size_(max_size),
        data_(std::make_unique_for_overwrite<uint8_t[]>(default_size_)),
        size_(default_size_) {}

  std::span<uint8_t> wspan(std::size_t need) override {
    auto avail = size_ - (offset_ + fill_);
    if (need > avail) {
      if (max_size_ - fill_ < need) // Early exit if need is never possible
        return {};
      if (offset_ > 0) {
        std::copy_n(data_.get() + offset_, fill_, data_.get());
        offset_ = 0;
      }
      avail = size_ - fill_;
      if (need > avail) {
        auto const max = std::numeric_limits<std::size_t>::max() / 2;
        std::size_t new_size = size_;
        while (true) {
          if (new_size <= max) {
            new_size *= 2;
          } else {
            new_size = std::numeric_limits<std::size_t>::max();
          }
          if (new_size >= max_size_) {
            new_size = max_size_;
            break;
          }
          if (new_size - fill_ >= need)
            break;
        }
        // Using new as it has std::nothrow which make_unique lacks.
        // Easy enought to keep track of the pointers here anyway.
        auto* tmp = new(std::nothrow) uint8_t[new_size];
        if (tmp == nullptr)
          return {};
        std::copy_n(data_.get(), fill_, tmp);
        size_ = new_size;
        data_.reset(tmp);
      }
    }
    return {data_.get() + offset_ + fill_, size_ - (offset_ + fill_)};
  }

  void commit(std::size_t size) override {
    assert(size_ - (offset_ + fill_) >= size);
    fill_ += size;
  }

  std::span<uint8_t const> rspan(std::size_t) override {
    return {data_.get() + offset_, fill_};
  }

  void consume(std::size_t size) override {
    if (size == 0)
      return;
    assert(fill_ >= size);
    fill_ -= size;
    if (fill_ == 0) {
      reset();
    } else {
      offset_ += size;
    }
  }

  std::span<uint8_t> mspan(std::size_t) override {
    return {data_.get() + offset_, fill_};
  }

  std::size_t uncommit(std::size_t size) override {
    auto ret = std::min(size, fill_);
    fill_ -= ret;
    if (fill_ == 0) {
      reset();
    }
    return ret;
  }

  bool empty() const override {
    return fill_ == 0;
  }

  bool full() const override {
    return fill_ >= max_size_;
  }

  void reset() override {
    if (size_ != default_size_)
      data_ = std::make_unique_for_overwrite<uint8_t[]>(size_ = default_size_);
    offset_ = 0;
    fill_ = 0;
  }

 private:
  std::size_t const default_size_;
  std::size_t const max_size_;
  std::unique_ptr<uint8_t[]> data_;
  std::size_t size_;
  std::size_t offset_{0};
  std::size_t fill_{0};
};

class FixedBuffer : public Buffer {
 public:
  explicit FixedBuffer(std::size_t size)
      : size_(size), data_(std::make_unique<uint8_t[]>(size_)) {}

  std::span<uint8_t> wspan(std::size_t need) override {
    auto avail = wavail();
    if (need > avail) {
      if (need > size_ - ravail())  // Early exit if need will never fit
        return {};
      if (rptr_ < wptr_ || (rptr_ == wptr_ && !full_)) {
        rotate();
        avail = wavail();
      } else {
        return {};
      }
    }
    return {data_.get() + wptr_, avail};
  }

  void commit(std::size_t size) override {
    if (size == 0)
      return;
    assert(wavail() >= size);
    wptr_ += size;
    if (wptr_ == size_)
      wptr_ = 0;
    if (rptr_ == wptr_)
      full_ = true;
  }

  std::span<uint8_t const> rspan(std::size_t want) override {
    return mspan(want);
  }

  void consume(std::size_t size) override {
    if (size == 0)
      return;
    assert(ravail() >= size);
    full_ = false;
    rptr_ += size;
    if (rptr_ == size_)
      rptr_ = 0;
    if (rptr_ == wptr_)
      reset();
  }

  std::span<uint8_t> mspan(std::size_t want) override {
    auto avail = ravail();
    if (want > avail) {
      if (rptr_ > wptr_ || (rptr_ == wptr_ && full_)) {
        rotate();
        avail = ravail();
      }
    }
    return {data_.get() + rptr_, avail};
  }

  std::size_t uncommit(std::size_t size) override {
    if (size == 0)
      return 0;
    auto ret = do_uncommit(size);
    if (ret < size) {
      ret += do_uncommit(size - ret);
    }
    return ret;
  }

  bool empty() const override {
    return rptr_ == wptr_ && !full_;
  }

  bool full() const override {
    return rptr_ == wptr_ && full_;
  }

  void reset() override {
    rptr_ = 0;
    wptr_ = 0;
    full_ = false;
  }

 private:
  std::size_t ravail() const {
    if (rptr_ < wptr_)
      return wptr_ - rptr_;
    if (rptr_ == wptr_ && !full_)
      return 0;
    return size_ - rptr_;
  }

  std::size_t wavail() const {
    if (rptr_ > wptr_)
      return rptr_ - wptr_;
    if (rptr_ == wptr_ && full_)
      return 0;
    return size_ - wptr_;
  }

  std::size_t do_uncommit(std::size_t size) {
    if (size == 0 || (rptr_ == wptr_ && !full_))
      return 0;

    full_ = false;

    if (wptr_ == 0)
      wptr_ = size_;

    auto avail = rptr_ < wptr_ ? wptr_ - rptr_ : wptr_;
    avail = std::min(avail, size);
    wptr_ -= avail;
    return avail;
  }

  void rotate() {
    assert(rptr_ > 0);

    if (rptr_ < wptr_) {
      std::copy(data_.get() + rptr_, data_.get() + wptr_, data_.get());
      wptr_ -= rptr_;
      rptr_ = 0;
    } else if (wptr_ < rptr_ || (wptr_ == rptr_ && full_)) {
      auto left = wptr_;
      auto right = size_ - rptr_;
      // TODO: Can we do this without allocations?
      if (left <= right) {
        auto tmp = std::make_unique<uint8_t[]>(left);
        std::copy_n(data_.get(), left, tmp.get());
        std::copy_n(data_.get() + rptr_, right, data_.get());
        std::copy_n(tmp.get(), left, data_.get() + right);
      } else {
        auto tmp = std::make_unique<uint8_t[]>(right);
        std::copy_n(data_.get() + rptr_, right, tmp.get());
        std::copy_backward(data_.get(), data_.get() + left,
                           data_.get() + left + right - 1);
        std::copy_n(tmp.get(), right, data_.get());
      }
      wptr_ = left + right;
      if (wptr_ == size_)
        wptr_ = 0;
      rptr_ = 0;
    } else {
      assert(false);
    }
  }

  std::size_t const size_;
  std::unique_ptr<uint8_t[]> data_;
  std::size_t rptr_{0};
  std::size_t wptr_{0};
  bool full_{false};
};

class ReadViewBufferImpl : public ReadViewBuffer {
 public:
  explicit ReadViewBufferImpl(std::unique_ptr<Buffer> buffer)
      : buffer_(std::move(buffer)) {}

  std::size_t consumed() const override {
    return offset_;
  }

  std::unique_ptr<Buffer> release() override {
    return std::move(buffer_);
  }

  std::span<uint8_t> wspan(std::size_t need) override {
    return buffer_->wspan(need);
  }

  void commit(std::size_t size) override {
    return buffer_->commit(size);
  }

  std::span<uint8_t const> rspan(std::size_t want) override {
    auto ret = buffer_->rspan(offset_ + want);
    if (ret.size() <= offset_)
      return ret.subspan(0, 0);
    return ret.subspan(offset_, ret.size() - offset_);
  }

  void consume(std::size_t size) override {
    offset_ += size;
  }

  std::span<uint8_t> mspan(std::size_t want) override {
    auto ret = buffer_->mspan(offset_ + want);
    if (ret.size() <= offset_)
      return ret.subspan(0, 0);
    return ret.subspan(offset_, ret.size() - offset_);
  }

  std::size_t uncommit(std::size_t size) override {
    return buffer_->uncommit(size);
  }

  bool empty() const override {
    if (buffer_->empty())
      return true;
    auto data = buffer_->rspan(offset_ + 1);
    return data.size() <= offset_;
  }

  bool full() const override {
    return buffer_->full();
  }

  void reset() override {
    offset_ = 0;
  }

 private:
  std::unique_ptr<Buffer> buffer_;
  std::size_t offset_{0};
};

}  // namespace

std::unique_ptr<Buffer> make_buffer(std::size_t default_size,
                                    std::size_t max_size) {
  if (default_size >= max_size)
    return std::make_unique<FixedBuffer>(max_size);

  return std::make_unique<DynamicBuffer>(default_size, max_size);
}

std::unique_ptr<ReadViewBuffer> make_read_view_buffer(
    std::unique_ptr<Buffer> buffer) {
  return std::make_unique<ReadViewBufferImpl>(std::move(buffer));
}

std::size_t Buffer::write(std::span<uint8_t const> data) {
  std::size_t offset = 0;
  while (offset < data.size()) {
    auto target = wspan();
    if (target.empty())
      break;
    auto size = std::min(data.size() - offset, target.size());
    std::copy_n(data.data() + offset, size, target.data());
    commit(size);
    offset += size;
  }
  return offset;
}

bool Buffer::write_all(std::span<uint8_t const> data) {
  if (data.empty())
    return true;
  auto target = wspan(data.size());
  if (target.empty())
    return false;
  std::copy(data.begin(), data.end(), target.begin());
  commit(data.size());
  return true;
}

std::size_t Buffer::read(std::span<uint8_t> data) {
  std::size_t offset = 0;
  while (offset < data.size()) {
    auto source = rspan();
    if (source.empty())
      break;
    auto size = std::min(data.size() - offset, source.size());
    std::copy_n(source.data(), size, data.data() + offset);
    consume(size);
    offset += size;
  }
  return offset;
}

bool Buffer::read_all(std::span<uint8_t> data) {
  auto source = rspan(data.size());
  if (source.size() < data.size())
    return false;
  std::copy_n(source.begin(), data.size(), data.begin());
  consume(data.size());
  return true;
}

}  // namespace sax
}  // namespace modxml