#pragma once

#include <cstdint>
#include <cstring>
#include <iostream>
#include <functional>

namespace Preserves {
    enum class BinaryTag {
        False = 0x80,
        True = 0x81,
        Float = 0x82,
        Double = 0x83,
        End = 0x84,
        Annotation = 0x85,
        Embedded = 0x86,
        SmallInteger_lo = 0x90,
        SmallInteger_hi = 0x9f,
        MediumInteger_lo = 0xa0,
        MediumInteger_hi = 0xaf,
        SignedInteger = 0xb0,
        String = 0xb1,
        ByteString = 0xb2,
        Symbol = 0xb3,
        Record = 0xb4,
        Sequence = 0xb5,
        Set = 0xb6,
        Dictionary = 0xb7,
    };

    class BinaryWriter {
        std::ostream& o;

    public:
        BinaryWriter(std::ostream& o) : o(o) {}

        std::ostream& stream() const { return o; }

        BinaryWriter& put(uint8_t b) {
            o.put(b);
            return *this;
        }

        BinaryWriter& operator<<(BinaryTag const& t) {
            return put(uint8_t(t));
        }

        BinaryWriter& write(void const* buf, size_t size) {
            o.write(reinterpret_cast<char const*>(buf), size);
            return *this;
        }

        BinaryWriter& varint(size_t n) {
            while (n >= 0x80) {
                put(uint8_t(0x80 | (n & 0x7f)));
                n >>= 7;
            }
            return put(uint8_t(n & 0x7f));
        }

        BinaryWriter& write(BinaryTag const& tag, void const* buf, size_t size) {
            (*this) << tag;
            varint(size);
            return write(buf, size);
        }

        BinaryWriter& bignum(std::vector<uint8_t> const& n) {
            if (n.size() == 0) return (*this) << BinaryTag::SmallInteger_lo;
            size_t startOffset = 0;
            if (n[0] & 0x80) {
                if (n.size() == 1 && int8_t(n[0]) >= -3) {
                    return put(uint8_t(int(BinaryTag::SmallInteger_lo) + 16 + int8_t(n[0])));
                } else {
                    while (true) {
                        if (startOffset == n.size() - 1) break;
                        if (n[startOffset] != 0xff) break;
                        startOffset++;
                        if ((n[startOffset] & 0x80) != 0x80) {
                            startOffset--;
                            break;
                        }
                    }
                }
            } else {
                if (n.size() == 1 && n[0] <= 12) {
                    return put(uint8_t(int(BinaryTag::SmallInteger_lo) + n[0]));
                } else {
                    while (true) {
                        if (startOffset == n.size() - 1) break;
                        if (n[startOffset] != 0x00) break;
                        startOffset++;
                        if ((n[startOffset] & 0x80) != 0x00) {
                            startOffset--;
                            break;
                        }
                    }
                }
            }
            size_t count = n.size() - startOffset;
            if (count > 16) {
                (*this) << BinaryTag::SignedInteger;
                varint(count);
            } else {
                put(uint8_t(int(BinaryTag::MediumInteger_lo) + count - 1));
            }
            return write(&n[startOffset], count);
        }

        BinaryWriter& string(std::string const& s) {
            return write(BinaryTag::String, &s[0], s.size());
        }

        BinaryWriter& symbol(std::string const& s) {
            return write(BinaryTag::Symbol, &s[0], s.size());
        }

        BinaryWriter& bytes(std::vector<uint8_t> const& v) {
            return bytes(v.data(), v.size());
        }

        BinaryWriter& bytes(void const* buf, size_t size) {
            return write(BinaryTag::ByteString, buf, size);
        }

        template <typename T = class GenericEmbedded>
        BinaryWriter& operator<<(Value<T> const& v) {
            return v->write(*this);
        }

        template <typename Iter>
        BinaryWriter& writeseq(Iter const& begin, Iter const& end) {
            for (Iter i = begin; i != end; ++i) (*this) << (*i);
            return *this;
        }

        template <typename Collection>
        BinaryWriter& operator<<(Collection const& c) {
            return writeseq(c.begin(), c.end());
        }

        template <typename T>
        BinaryWriter& record(Value<T> const& label, std::vector<Value<T>> const& fields) {
            return (*this) << BinaryTag::Record << label << fields << BinaryTag::End;
        }

        template <typename T>
        BinaryWriter& sequence(std::vector<Value<T>> const& vs) {
            return (*this) << BinaryTag::Sequence << vs << BinaryTag::End;
        }

        template <typename T>
        BinaryWriter& set(std::set<Value<T>> const& vs) {
            return (*this) << BinaryTag::Set << vs << BinaryTag::End;
        }

        template <typename T>
        BinaryWriter& dictionary(std::map<Value<T>, Value<T>> const& vs) {
            (*this) << BinaryTag::Dictionary;
            for (auto& i : vs) { (*this) << i.first << i.second; }
            return (*this) << BinaryTag::End;
        }
    };

    inline BinaryWriter& operator<<(BinaryWriter& w, bool b) {
        return w << (b ? BinaryTag::True : BinaryTag::False);
    }

    inline BinaryWriter& operator<<(BinaryWriter& w, float f) {
        uint32_t n;
        memcpy(&n, &f, sizeof(f));
        uint8_t buf[4];
        buf[0] = (n >> 24) & 0xff;
        buf[1] = (n >> 16) & 0xff;
        buf[2] = (n >> 8) & 0xff;
        buf[3] = (n) & 0xff;
        w << BinaryTag::Float;
        return w.write(buf, sizeof(buf));
    }

    inline BinaryWriter& operator<<(BinaryWriter& w, double d) {
        uint64_t n;
        memcpy(&n, &d, sizeof(d));
        uint8_t buf[8];
        buf[0] = (n >> 56) & 0xff;
        buf[1] = (n >> 48) & 0xff;
        buf[2] = (n >> 40) & 0xff;
        buf[3] = (n >> 32) & 0xff;
        buf[4] = (n >> 24) & 0xff;
        buf[5] = (n >> 16) & 0xff;
        buf[6] = (n >> 8) & 0xff;
        buf[7] = (n) & 0xff;
        w << BinaryTag::Double;
        return w.write(buf, sizeof(buf));
    }

    template <uint64_t wholeTest, uint64_t bitTest>
    int _shift_for(uint64_t u) {
        int shift = 56;
        while (true) {
            if (shift == 0) break;
            if (((u >> shift) & 0xff) != wholeTest) break;
            shift -= 8;
            if (((u >> shift) & 0x80) != bitTest) {
                shift += 8;
                break;
            }
        }
        return shift;
    }

    inline BinaryWriter& _put_medium_int(BinaryWriter& w, uint64_t u, int shift, int extra) {
        w.put(uint8_t(int(BinaryTag::MediumInteger_lo) + (shift >> 3) + extra));
        if (extra) w.put(0);
        while (shift >= 0) {
            w.put(uint8_t((u >> shift) & 0xff));
            shift -= 8;
        }
        return w;
    }

    inline BinaryWriter& operator<<(BinaryWriter& w, int64_t i) {
        if (i < 0) {
            if (i >= -3) {
                return w.put(uint8_t(int(BinaryTag::SmallInteger_lo) + 16 + i));
            } else {
                uint64_t u;
                memcpy(&u, &i, sizeof(i));
                return _put_medium_int(w, u, _shift_for<0xff, 0x80>(u), 0);
            }
        } else {
            if (i <= 12) {
                return w.put(uint8_t(int(BinaryTag::SmallInteger_lo) + i));
            } else {
                uint64_t u;
                memcpy(&u, &i, sizeof(i));
                return _put_medium_int(w, u, _shift_for<0x00, 0x00>(u), 0);
            }
        }
    }

    inline BinaryWriter& operator<<(BinaryWriter& w, uint64_t u) {
        if (u <= 12) {
            return w.put(uint8_t(int(BinaryTag::SmallInteger_lo) + u));
        } else if ((u & 0x8000000000000000) != 0) {
            return _put_medium_int(w, u, 56, 1);
        } else {
            return _put_medium_int(w, u, _shift_for<0x00, 0x00>(u), 0);
        }
    }
}