#include "Encoder.hpp"

#include "Core.hpp"
#include "Utils.hpp"

#include <cstring>
#include <format>
#include <span>
#include <string_view>
#include <utility>

namespace CBOR
{
    class NotEnoughSpace: public EncodeError
    {
    public:
        NotEnoughSpace(std::size_t remaining, std::size_t needed)
            : EncodeError(std::format("the encoder needs at least more {} bytes, "
                                      "but only {} are left",
                                      needed, remaining))
        {}
    };

    class UnknownValue: public EncodeError
    {
    public:
        UnknownValue()
            : EncodeError("tried to encode an unknown CBOR value")
        {}
    };

    namespace
    {
        template <typename... Funcs>
        struct Overload: public Funcs...
        {
            using Funcs::operator()...;
        };

        std::uint8_t EncodeHeader(MajorType type, ArgumentPosition pos)
        {
            return std::to_underlying(type) | std::to_underlying(pos);
        }

        std::uint8_t EncodeHeader(MajorType type, MinorType minor)
        {
            return std::to_underlying(type) | std::to_underlying(minor);
        }

        std::uint8_t EncodeHeaderDirect(MajorType type, std::uint8_t value)
        {
            return std::to_underlying(type) | value;
        }

        void WriteHeader(EncoderBuffer &buffer, MajorType type, std::uint64_t argument)
        {
            if (argument <= 23) {
                buffer.Write(EncodeHeaderDirect(type, std::uint8_t(argument)));
            }
            else if (argument <= 255) {
                buffer.Write(EncodeHeader(type, ArgumentPosition::Next1B));
                buffer.Write(std::uint8_t(argument));
            }
            else if (argument <= 65535) {
                buffer.Write(EncodeHeader(type, ArgumentPosition::Next2B));
                buffer.Write(std::uint16_t(argument));
            }
            else if (argument <= 4294967295) {
                buffer.Write(EncodeHeader(type, ArgumentPosition::Next4B));
                buffer.Write(std::uint32_t(argument));
            }
            else {
                buffer.Write(EncodeHeader(type, ArgumentPosition::Next8B));
                buffer.Write(std::uint64_t(argument));
            }
        }
    }

    EncoderBuffer::EncoderBuffer(Buffer buffer)
        : mBuffer(FixedBuffer { .mCurrent = 0, .mBuffer = buffer, })
    {}

    EncoderBuffer::EncoderBuffer(std::vector<std::uint8_t> &buffer)
        : mBuffer(DynamicBuffer { .mBuffer = &buffer })
    {}

    EncoderBuffer::EncoderBuffer(EncoderBuffer &&other)
        : mBuffer(std::exchange(other.mBuffer, FixedBuffer { .mCurrent = 0, .mBuffer = {}, }))
    {}

    void EncoderBuffer::Write(std::uint8_t value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    void EncoderBuffer::Write(std::uint16_t value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    void EncoderBuffer::Write(std::uint32_t value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    void EncoderBuffer::Write(std::uint64_t value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    void EncoderBuffer::Write(ConstBuffer value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    void EncoderBuffer::Write(std::string_view value)
    {
        std::visit(Overload {
            [value] (FixedBuffer &buffer) { buffer.Write(value); },
            [value] (DynamicBuffer &buffer) { buffer.Write(value); },
        }, mBuffer);
    }

    std::size_t EncoderBuffer::Size() const
    {
        return std::visit(Overload {
            [] (const FixedBuffer &buffer) { return buffer.Size(); },
            [] (const DynamicBuffer &buffer) { return buffer.Size(); },
        }, mBuffer);
    }

    void EncoderBuffer::FixedBuffer::Write(std::uint8_t value)
    {
        EnsureSpace(1);
        mBuffer[mCurrent++] = value;
    }

    void EncoderBuffer::FixedBuffer::Write(std::uint16_t value)
    {
        EnsureSpace(2);
        static constexpr std::uint16_t mask = 0x00'FF;
        std::uint16_t network = HostToNetwork(value);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network     ) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 8) & mask);
    }

    void EncoderBuffer::FixedBuffer::Write(std::uint32_t value)
    {
        EnsureSpace(4);
        static constexpr std::uint32_t mask = 0x00'00'00'FF;
        std::uint32_t network = HostToNetwork(value);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network      ) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >>  8) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 16) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 24) & mask);
    }

    void EncoderBuffer::FixedBuffer::Write(std::uint64_t value)
    {
        EnsureSpace(8);
        static constexpr std::uint64_t mask = 0x00'00'00'00'00'00'00'FF;
        std::uint64_t network = HostToNetwork(value);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network      ) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >>  8) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 16) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 24) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 32) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 40) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 48) & mask);
        mBuffer[mCurrent++] = static_cast<std::uint8_t>((network >> 56) & mask);
    }

    void EncoderBuffer::FixedBuffer::Write(ConstBuffer value)
    {
        EnsureSpace(value.size());
        std::memcpy(mBuffer.data() + mCurrent, value.data(), value.size());
        mCurrent += value.size();
    }

    void EncoderBuffer::FixedBuffer::Write(std::string_view value)
    {
        EnsureSpace(value.size());
        std::memcpy(mBuffer.data() + mCurrent, value.data(), value.size());
        mCurrent += value.size();
    }

    void EncoderBuffer::FixedBuffer::EnsureSpace(std::size_t size) const
    {
        std::size_t spaceLeft = mCurrent > mBuffer.size() ? 0 : mBuffer.size() - mCurrent;

        if (size > spaceLeft) {
            throw NotEnoughSpace(spaceLeft, size);
        }
    }

    std::size_t EncoderBuffer::FixedBuffer::Size() const
    {
        return mCurrent;
    }

    void EncoderBuffer::DynamicBuffer::Write(std::uint8_t value)
    {
        mBuffer->push_back(value);
    }

    void EncoderBuffer::DynamicBuffer::Write(std::uint16_t value)
    {
        static constexpr std::uint16_t mask = 0x00'FF;
        std::uint16_t network = HostToNetwork(value);
        mBuffer->push_back(static_cast<std::uint8_t>((network     ) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 8) & mask));
    }

    void EncoderBuffer::DynamicBuffer::Write(std::uint32_t value)
    {
        static constexpr std::uint32_t mask = 0x00'00'00'FF;
        std::uint32_t network = HostToNetwork(value);
        mBuffer->push_back(static_cast<std::uint8_t>((network      ) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >>  8) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 16) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 24) & mask));
    }

    void EncoderBuffer::DynamicBuffer::Write(std::uint64_t value)
    {
        static constexpr std::uint64_t mask = 0x00'00'00'00'00'00'00'FF;
        std::uint64_t network = HostToNetwork(value);
        mBuffer->push_back(static_cast<std::uint8_t>((network      ) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >>  8) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 16) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 24) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 32) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 40) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 48) & mask));
        mBuffer->push_back(static_cast<std::uint8_t>((network >> 56) & mask));
    }

    void EncoderBuffer::DynamicBuffer::Write(ConstBuffer value)
    {
        mBuffer->append_range(value);
    }

    void EncoderBuffer::DynamicBuffer::Write(std::string_view value)
    {
        mBuffer->append_range(value);
    }

    std::size_t EncoderBuffer::DynamicBuffer::Size() const
    {
        return mBuffer->size();
    }

    BasicEncoder::BasicEncoder(EncoderBuffer buffer)
        : mBuffer(std::move(buffer))
    {}

    void BasicEncoder::Encode(bool value)
    {
        if (value) {
            mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::True));
        }
        else {
            mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::False));
        }
    }

    void BasicEncoder::Encode(Special value)
    {
        switch (value) {
        case Special::Null:
            mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::Null));
            break;
        case Special::Undefined:
            mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::Undefined));
            break;
        default:
            throw UnknownValue();
        }
    }

    void BasicEncoder::Encode(std::int8_t value)
    {
        if (value >= 0) {
            if (value <= 23) {
                mBuffer.Write(EncodeHeaderDirect(MajorType::Unsigned, std::uint8_t(value)));
            }
            else {
                mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next1B));
                mBuffer.Write(static_cast<std::uint8_t>(value));
            }
        }
        else {
            std::int8_t actual = -(value + 1);
            if (actual <= 23) {
                mBuffer.Write(EncodeHeaderDirect(MajorType::Negative, std::uint8_t(actual)));
            }
            else {
                mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next1B));
                mBuffer.Write(static_cast<std::uint8_t>(actual));
            }
        }
    }

    void BasicEncoder::Encode(std::int16_t value)
    {
        if (value >= -128 && value <= 127) {
            Encode(static_cast<std::int8_t>(value));
        }
        else if (value >= 0 && value <= 255) {
            Encode(static_cast<std::uint8_t>(value));
        }
        else if (value >= -256 && value <= -1) {
            std::uint8_t actual = static_cast<std::uint8_t>(-(value + 1));
            mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next1B));
            mBuffer.Write(actual);
        }
        else if (value >= 0) {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next2B));
            mBuffer.Write(static_cast<std::uint16_t>(value));
        }
        else {
            std::int16_t actual = -(value + 1);
            mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next2B));
            mBuffer.Write(static_cast<std::uint16_t>(actual));
        }
    }

    void BasicEncoder::Encode(std::int32_t value)
    {
        if (value >= -32768 && value <= 32767){
            Encode(static_cast<std::int16_t>(value));
        }
        else if (value >= 0 && value <= 65535) {
            Encode(static_cast<std::uint16_t>(value));
        }
        else if (value >= -65536 && value <= -1) {
            std::uint16_t actual = static_cast<std::uint16_t>(std::abs(value + 1));
            mBuffer.Write(actual);
        }
        else if (value >= 0) {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next4B));
            mBuffer.Write(static_cast<std::uint32_t>(value));
        }
        else {
            std::int32_t actual = std::abs(value + 1);
            mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next4B));
            mBuffer.Write(static_cast<std::uint32_t>(actual));
        }
    }

    void BasicEncoder::Encode(std::int64_t value)
    {
        if (value >= -2147483648 && value <= 2147483647){
            Encode(static_cast<std::int32_t>(value));
        }
        else if (value >= 0 && value <= 4294967295) {
            Encode(static_cast<std::uint32_t>(value));
        }
        else if (value >= -2147483648 && value <= -1) {
            std::uint32_t actual = static_cast<std::uint32_t>(std::abs(value + 1));
            mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next4B));
            mBuffer.Write(actual);
        }
        else if (value >= 0) {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next8B));
            mBuffer.Write(static_cast<std::uint64_t>(value));
        }
        else {
            std::int64_t actual = std::abs(value + 1);
            mBuffer.Write(EncodeHeader(MajorType::Negative, ArgumentPosition::Next8B));
            mBuffer.Write(static_cast<std::uint64_t>(actual));
        }
    }

    void BasicEncoder::Encode(std::uint8_t value)
    {
        if (value <= 23) {
            mBuffer.Write(EncodeHeaderDirect(MajorType::Unsigned, value));
        }
        else {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next1B));
            mBuffer.Write(value);
        }
    }

    void BasicEncoder::Encode(std::uint16_t value)
    {
        if (value <= 255) {
            Encode(static_cast<std::uint8_t>(value));
        }
        else {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next2B));
            mBuffer.Write(value);
        }
    }

    void BasicEncoder::Encode(std::uint32_t value)
    {
        if (value <= 65535) {
            Encode(static_cast<std::uint16_t>(value));
        }
        else {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next4B));
            mBuffer.Write(value);
        }
    }

    void BasicEncoder::Encode(std::uint64_t value)
    {
        if (value <= 4294967295) {
            Encode(static_cast<std::uint32_t>(value));
        }
        else {
            mBuffer.Write(EncodeHeader(MajorType::Unsigned, ArgumentPosition::Next8B));
            mBuffer.Write(value);
        }
    }

    void BasicEncoder::Encode(float value)
    {
        // Note(3011): This is suboptimal, but dealing with IEEE 754 numbers correctly
        // is difficult, and has been left for later.
        mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::Float));
        mBuffer.Write(std::bit_cast<std::uint32_t>(value));
    }

    void BasicEncoder::Encode(double value)
    {
        // Note(3011): This is suboptimal, but dealing with IEEE 754 numbers correctly
        // is difficult, and has been left for later.
        mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::Double));
        mBuffer.Write(std::bit_cast<std::uint64_t>(value));
    }

    void BasicEncoder::Encode(ConstBuffer value)
    {
        WriteHeader(mBuffer, MajorType::Binary, value.size());
        mBuffer.Write(value);
    }

    void BasicEncoder::Encode(const char *value)
    {
        Encode(std::string_view(value));
    }

    void BasicEncoder::Encode(std::string_view value)
    {
        WriteHeader(mBuffer, MajorType::String, value.size());
        mBuffer.Write(value);
    }

    void BasicEncoder::BeginArray(std::size_t size)
    {
        WriteHeader(mBuffer, MajorType::Array, size);
    }

    void BasicEncoder::BeginMap(std::size_t size)
    {
        WriteHeader(mBuffer, MajorType::Map, size);
    }

    void BasicEncoder::BeginIndefiniteBinary()
    {
        mBuffer.Write(EncodeHeader(MajorType::Binary, ArgumentPosition::Indefinite));
    }

    void BasicEncoder::BeginIndefiniteString()
    {
        mBuffer.Write(EncodeHeader(MajorType::String, ArgumentPosition::Indefinite));
    }

    void BasicEncoder::BeginIndefiniteArray()
    {
        mBuffer.Write(EncodeHeader(MajorType::Array, ArgumentPosition::Indefinite));
    }

    void BasicEncoder::BeginIndefiniteMap()
    {
        mBuffer.Write(EncodeHeader(MajorType::Map, ArgumentPosition::Indefinite));
    }

    void BasicEncoder::End()
    {
        mBuffer.Write(EncodeHeader(MajorType::Other, MinorType::Break));
    }

    void BasicEncoder::EncodeTag(std::uint64_t value)
    {
        WriteHeader(mBuffer, MajorType::Tag, value);
    }

    std::size_t BasicEncoder::Size() const
    {
        return mBuffer.Size();
    }

    Encoder::Encoder(EncoderBuffer buffer)
        : mEncoder(std::move(buffer))
    {}

    std::size_t Encoder::Size() const
    {
        return mEncoder.Size();
    }
}