#pragma once

#if !defined(MIJIN_LOGGING_BUFFER_SINK_HPP_INCLUDED)
#define MIJIN_LOGGING_BUFFER_SINK_HPP_INCLUDED 1

#include <array>
#include <cstdint>
#include <cstring>
#include <mutex>
#include <span>
#include "./logger.hpp"

namespace mijin
{
inline constexpr std::size_t BUFFER_SINK_DEFAULT_SIZE = 10 * 1024 * 1024; // default 10 MiB buffer

template<std::size_t VBufferSize = BUFFER_SINK_DEFAULT_SIZE, typename TChar = MIJIN_DEFAULT_CHAR_TYPE, typename TTraits = std::char_traits<TChar>>
class BaseBufferSink;

namespace impl
{
static constexpr std::uint32_t INVALID_BUFFER_INDEX = std::numeric_limits<std::uint32_t>::max();

template<typename TChar>
struct BufferedMessageHeader
{
    const BaseLogChannel<TChar>* channel;
    const BaseLogLevel<TChar>* level;
    std::source_location sourceLocation;
    std::uint32_t nextIdx = INVALID_BUFFER_INDEX;
    std::uint32_t prevIdx = INVALID_BUFFER_INDEX;
    std::uint32_t numChars = 0;
};

template<typename TChar, typename TTraits, std::size_t VBufferSize>
struct MessageBuffer
{
    using char_t = TChar;
    using traits_t = TTraits;
    using message_t = BaseLogMessage<char_t, traits_t>;
    using header_t = BufferedMessageHeader<TChar>;
    static constexpr std::size_t BUFFER_SIZE = VBufferSize;

    std::array<std::uint8_t, BUFFER_SIZE> bytes;
    mutable std::mutex mutex;

    [[nodiscard]]
    header_t& messageAt(std::uint32_t idx) MIJIN_NOEXCEPT
    {
        return *reinterpret_cast<header_t*>(bytes.data() + idx);
    }

    [[nodiscard]]
    std::span<char_t> messageText(header_t& header) MIJIN_NOEXCEPT
    {
        return {reinterpret_cast<char_t*>(reinterpret_cast<std::uint8_t*>(&header)) + sizeof(header_t), header.numChars};
    }

    [[nodiscard]]
    const header_t& messageAt(std::uint32_t idx) const MIJIN_NOEXCEPT
    {
        return *reinterpret_cast<const header_t*>(bytes.data() + idx);
    }

    [[nodiscard]]
    std::span<const char_t> messageText(const header_t& header) const MIJIN_NOEXCEPT
    {
        return {reinterpret_cast<const char_t*>(reinterpret_cast<const std::uint8_t*>(&header)) + sizeof(header_t), header.numChars};
    }

    static std::uint32_t messageBytes(std::uint32_t numChars) MIJIN_NOEXCEPT
    {
        return static_cast<std::uint32_t>(sizeof(header_t)) + (numChars * static_cast<std::uint32_t>(sizeof(char_t)));
    }
};

template<typename TChar, typename TTraits, std::size_t VBufferSize>
class BufferSinkRange;

template<typename TChar, typename TTraits, std::size_t VBufferSize>
class LockedMessageBuffer
{
public:
    using char_t = TChar;
    using traits_t = TTraits;
    
    static constexpr std::size_t BUFFER_SIZE = VBufferSize;

    using range_t = impl::BufferSinkRange<char_t, traits_t, BUFFER_SIZE>;
private:
    using buffer_t = MessageBuffer<char_t, traits_t, BUFFER_SIZE>;

    const buffer_t* buffer_ = nullptr;
    std::uint32_t firstIdx_ = INVALID_BUFFER_INDEX;
    std::uint32_t lastIdx_ = INVALID_BUFFER_INDEX;

    LockedMessageBuffer(const buffer_t& buffer, std::uint32_t firstIdx, std::uint32_t lastIdx) MIJIN_NOEXCEPT
        : buffer_(&buffer), firstIdx_(firstIdx), lastIdx_(lastIdx) {
        buffer_->mutex.lock();
    }
public:
    LockedMessageBuffer() noexcept = default;
    LockedMessageBuffer(const LockedMessageBuffer&) = delete;
    LockedMessageBuffer(LockedMessageBuffer&& other) noexcept
        : buffer_(std::exchange(other.buffer_, nullptr)), firstIdx_(other.firstIdx_), lastIdx_(other.lastIdx_) {}
    
        ~LockedMessageBuffer() noexcept
    {
        reset();
    }

    LockedMessageBuffer& operator=(const LockedMessageBuffer&) = delete;
    LockedMessageBuffer& operator=(LockedMessageBuffer&& other) noexcept
    {
        if (this != &other)
        {
            reset();
            buffer_ = std::exchange(other.buffer_, nullptr);
            firstIdx_ = other.firstIdx_;
            lastIdx_ = other.lastIdx_;
        }
        return *this;
    }

    [[nodiscard]]
    range_t getMessages() const MIJIN_NOEXCEPT;
private:
    void reset()
    {
        if (buffer_ != nullptr)
        {
            buffer_->mutex.unlock();
            buffer_ = nullptr;
        }
    }

    template<std::size_t, typename, typename>
    friend class mijin::BaseBufferSink;
};

template<typename TChar, typename TTraits, std::size_t VBufferSize>
class BufferSinkIterator
{
public:
    using char_t = TChar;
    using traits_t = TTraits;
    using message_t = BaseLogMessage<char_t, traits_t>;
    static constexpr std::size_t BUFFER_SIZE = VBufferSize;

    using difference_type = std::ptrdiff_t;
    using value_type = const message_t;
    using pointer = void;
    using reference = value_type;
    using iterator_category = std::bidirectional_iterator_tag;
private:
    using buffer_t = MessageBuffer<char_t, traits_t, BUFFER_SIZE>;
    using header_t = BufferedMessageHeader<char_t>;
    using string_view_t = message_t::string_view_t;

    const buffer_t* buffer_ = nullptr;
    std::uint32_t idx_ = INVALID_BUFFER_INDEX;
    std::uint32_t prevIdx_ = INVALID_BUFFER_INDEX;

    BufferSinkIterator(const buffer_t& buffer, std::uint32_t idx, std::uint32_t prevIdx) MIJIN_NOEXCEPT : buffer_(&buffer), idx_(idx), prevIdx_(prevIdx) {}
public:
    BufferSinkIterator() MIJIN_NOEXCEPT = default;
    BufferSinkIterator(const BufferSinkIterator&) MIJIN_NOEXCEPT = default;

    BufferSinkIterator& operator=(const BufferSinkIterator&) MIJIN_NOEXCEPT = default;

    auto operator<=>(const BufferSinkIterator&) const noexcept = default;

    reference operator*() const MIJIN_NOEXCEPT
    {
        MIJIN_ASSERT(idx_ != INVALID_BUFFER_INDEX, "Attempting to dereference an invalid iterator.");
        const header_t& header = buffer_->messageAt(idx_);
        const string_view_t text(buffer_->messageText(header));
        return {
            .text = text,
            .channel = header.channel,
            .level = header.level,
            .sourceLocation = header.sourceLocation
        };
    }

    BufferSinkIterator& operator++() MIJIN_NOEXCEPT
    {
        MIJIN_ASSERT(idx_ != INVALID_BUFFER_INDEX, "Attempting to increment an invalid iterator.");
        prevIdx_ = idx_;
        idx_ = buffer_->messageAt(idx_).nextIdx;
        return *this;
    }

    BufferSinkIterator operator++(int) MIJIN_NOEXCEPT
    {
        BufferSinkIterator copy(*this);
        operator++();
        return copy;
    }

    BufferSinkIterator& operator--() MIJIN_NOEXCEPT
    {
        MIJIN_ASSERT(prevIdx_ != INVALID_BUFFER_INDEX, "Attempting to decrement an invalid iterator.");
        idx_ = prevIdx_;
        prevIdx_ = buffer_->messageAt(idx_).prevIdx;
        return *this;
    }

    BufferSinkIterator operator--(int) MIJIN_NOEXCEPT
    {
        BufferSinkIterator copy(*this);
        operator--();
        return copy;
    }
    
    template<typename, typename, std::size_t>
    friend class BufferSinkRange;
};
static_assert(std::forward_iterator<BufferSinkIterator<char, std::char_traits<char>, BUFFER_SINK_DEFAULT_SIZE>>);

template<typename TChar, typename TTraits, std::size_t VBufferSize>
class BufferSinkRange
{
public:
    using char_t = TChar;
    using traits_t = TTraits;
    static constexpr std::size_t BUFFER_SIZE = VBufferSize;
    
    using iterator = BufferSinkIterator<char_t, traits_t, BUFFER_SIZE>;
    using const_iterator = iterator;
private:
    using buffer_t = MessageBuffer<char_t, traits_t, BUFFER_SIZE>;

    const buffer_t* buffer_ = nullptr;
    std::uint32_t firstIdx_ = INVALID_BUFFER_INDEX;
    std::uint32_t lastIdx_ = INVALID_BUFFER_INDEX;

    BufferSinkRange(const buffer_t& buffer, std::uint32_t firstIdx, std::uint32_t lastIdx) MIJIN_NOEXCEPT
        : buffer_(&buffer), firstIdx_(firstIdx), lastIdx_(lastIdx) {
    }
public:
    BufferSinkRange() noexcept = default;
    BufferSinkRange(const BufferSinkRange&) MIJIN_NOEXCEPT = default;
    BufferSinkRange& operator=(const BufferSinkRange&) MIJIN_NOEXCEPT = default;

    [[nodiscard]]
    iterator begin() const MIJIN_NOEXCEPT { return {*buffer_, firstIdx_, INVALID_BUFFER_INDEX}; }

    [[nodiscard]]
    iterator end() const MIJIN_NOEXCEPT { return {*buffer_, INVALID_BUFFER_INDEX, lastIdx_}; }

    [[nodiscard]]
    const_iterator cbegin() const MIJIN_NOEXCEPT { return {*buffer_, firstIdx_, INVALID_BUFFER_INDEX}; }

    [[nodiscard]]
    const_iterator cend() const MIJIN_NOEXCEPT { return {*buffer_, INVALID_BUFFER_INDEX, lastIdx_}; }

    friend class LockedMessageBuffer<char_t, traits_t, BUFFER_SIZE>;
};

static_assert(std::bidirectional_iterator<BufferSinkIterator<char, std::char_traits<char>, 100>>);
static_assert(std::ranges::bidirectional_range<BufferSinkRange<char, std::char_traits<char>, 100>>);

template<typename TChar, typename TTraits, std::size_t VBufferSize>
auto LockedMessageBuffer<TChar, TTraits, VBufferSize>::getMessages() const MIJIN_NOEXCEPT -> range_t
{
    return range_t(*buffer_, firstIdx_, lastIdx_);
}
}

using impl::LockedMessageBuffer;

template<std::size_t VBufferSize, typename TChar, typename TTraits>
class BaseBufferSink : public BaseLogSink<TChar, TTraits>
{
public:
    using base_t = BaseLogSink<TChar>;
    using typename base_t::char_t;
    using typename base_t::traits_t;
    using typename base_t::message_t;

    static constexpr std::size_t BUFFER_SIZE = VBufferSize;
    
    using locked_buffer_t = LockedMessageBuffer<char_t, traits_t, BUFFER_SIZE>;
private:
    using buffer_t = impl::MessageBuffer<char_t, traits_t, BUFFER_SIZE>;
    using header_t = impl::BufferedMessageHeader<char_t>;

    buffer_t buffer_;
    std::uint32_t firstIdx_ = impl::INVALID_BUFFER_INDEX;
    std::uint32_t lastIdx_ = impl::INVALID_BUFFER_INDEX;
public:
    void handleMessage(const message_t& message) MIJIN_NOEXCEPT override
    {
        const std::uint32_t numChars = static_cast<std::uint32_t>(message.text.size());
        const std::uint32_t totalBytes = buffer_t::messageBytes(numChars);

        std::scoped_lock _(buffer_.mutex); // TODO: use a message queue and try_lock here
        if (lastIdx_ == impl::INVALID_BUFFER_INDEX)
        {
            // no message yet
            insertMessageAt(0, message);
            firstIdx_ = 0;
            return;
        }

        header_t& lastHeader = buffer_.messageAt(lastIdx_);
        const std::uint32_t newIdx = lastIdx_ + buffer_t::messageBytes(lastHeader.numChars);
        if (newIdx + totalBytes < BUFFER_SIZE)
        {
            // enough space in the buffer, can append
            insertMessageAt(newIdx, message);
        }
        else
        {
            // not enough space, put at front
            insertMessageAt(0, message);
        }
    }

    [[nodiscard]]
    locked_buffer_t lockBuffer() const MIJIN_NOEXCEPT { return locked_buffer_t(buffer_, firstIdx_, lastIdx_); }
private:
    void insertMessageAt(std::uint32_t idx, const message_t& message)
    {
        const std::uint32_t numChars = static_cast<std::uint32_t>(message.text.size());
        freeSpace(idx, idx + buffer_t::messageBytes(numChars));

        if (lastIdx_ != impl::INVALID_BUFFER_INDEX) {
            buffer_.messageAt(lastIdx_).nextIdx = idx;
        }

        header_t& newHeader = buffer_.messageAt(idx);
        ::new (&newHeader) header_t({
            .channel = message.channel,
            .level = message.level,
            .sourceLocation = message.sourceLocation,
            .nextIdx = impl::INVALID_BUFFER_INDEX,
            .prevIdx = lastIdx_,
            .numChars = numChars
        });
        lastIdx_ = idx;
        std::ranges::copy(message.text, buffer_.messageText(newHeader).begin());
    }

    void freeSpace(std::uint32_t startIdx, std::uint32_t endIdx)
    {
        while (firstIdx_ != impl::INVALID_BUFFER_INDEX && firstIdx_ >= startIdx && firstIdx_ < endIdx)
        {
            header_t& message = buffer_.messageAt(firstIdx_);
            firstIdx_ = message.nextIdx;
            message.~header_t();
        }

        // cleared everything?
        if (firstIdx_ == impl::INVALID_BUFFER_INDEX) {
            lastIdx_ = impl::INVALID_BUFFER_INDEX;
        }
        else {
            buffer_.messageAt(firstIdx_).prevIdx = impl::INVALID_BUFFER_INDEX;
        }
    }
};

#define SINK_COMMON_ARGS(chr_type) std::size_t BUFFER_SIZE
#define SINK_SET_ARGS(chr_type) BUFFER_SIZE, chr_type

MIJIN_DEFINE_CHAR_VERSIONS_TMPL(BufferSink, SINK_COMMON_ARGS, SINK_SET_ARGS)

#undef SINK_COMMON_ARGS
#undef SINK_SET_ARGS
} // namespace mijin


#endif // !defined(MIJIN_LOGGING_STDIO_SINK_HPP_INCLUDED)