#pragma once

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

namespace Preserves {
    enum class BinaryTag {
        False = 0x80,
        True = 0x81,
        End = 0x84,
        Annotation = 0x85,
        Embedded = 0x86,
        Ieee754 = 0x87,
        SignedInteger = 0xb0,
        String = 0xb1,
        ByteString = 0xb2,
        Symbol = 0xb3,
        Record = 0xb4,
        Sequence = 0xb5,
        Set = 0xb6,
        Dictionary = 0xb7,
    };

    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;
    }

    class BinaryWriter {
        std::ostream& o;

        template <typename Collection>
        struct _iterable {
            Collection const& c;
        };

        template <typename Collection>
        _iterable<Collection> iterable(Collection const& c) { return _iterable<Collection>{c}; }

    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) {
            size_t startOffset = 0;
            if (n.size() > 0) {
                if (n[0] & 0x80) {
                    while (true) {
                        if (startOffset == n.size() - 1) break;
                        if (n[startOffset] != 0xff) break;
                        startOffset++;
                        if ((n[startOffset] & 0x80) != 0x80) {
                            startOffset--;
                            break;
                        }
                    }
                } 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;
            (*this) << BinaryTag::SignedInteger;
            varint(count);
            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<<(_iterable<Collection> const& c) {
            return writeseq(c.c.begin(), c.c.end());
        }

        BinaryWriter& operator<<(bool b) {
            return (*this) << (b ? BinaryTag::True : BinaryTag::False);
        }

        BinaryWriter& operator<<(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;
            (*this) << BinaryTag::Ieee754;
            put(uint8_t(sizeof(buf)));
            return write(buf, sizeof(buf));
        }

        BinaryWriter& _put_medium_int(uint64_t u, int shift, int extra) {
            (*this) << BinaryTag::SignedInteger;
            put(uint8_t((shift >> 3) + extra + 1));
            if (extra) put(0);
            while (shift >= 0) {
                put(uint8_t((u >> shift) & 0xff));
                shift -= 8;
            }
            return *this;
        }

        template <typename T>
        typename std::enable_if<std::is_integral<T>::value, BinaryWriter&>::type operator<<(T t) {
            if (t == 0) {
                (*this) << BinaryTag::SignedInteger;
                put(0);
                return *this;
            } else if (std::numeric_limits<T>::is_signed) {
                auto i = static_cast<int64_t>(t);
                uint64_t u;
                memcpy(&u, &i, sizeof(i));
                int shift = i < 0 ? _shift_for<0xff, 0x80>(u) : _shift_for<0x00, 0x00>(u);
                return _put_medium_int(u, shift, 0);
            } else {
                auto u = static_cast<uint64_t>(t);
                if ((u & 0x8000000000000000) != 0) {
                    return _put_medium_int(u, 56, 1);
                } else {
                    return _put_medium_int(u, _shift_for<0x00, 0x00>(u), 0);
                }
            }
        }

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

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

        template <typename T>
        BinaryWriter& set(std::set<Value<T>> const& vs) {
            return (*this) << BinaryTag::Set << iterable(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;
        }

        template <typename T>
        BinaryWriter& annotated(Value<T> const& underlying, std::vector<Value<T>> const& vs) {
            for (auto& v : vs) { (*this) << BinaryTag::Annotation << v; }
            return (*this) << underlying;
        }
    };
}