#pragma once

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

#include "preserves.hpp"

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

    template <typename T = GenericEmbedded>
    class BinaryReader {
        std::istream& i;
        std::function<std::shared_ptr<T>(Value<>)> decodeEmbedded;

        bool next_chunk(void* p, size_t n) {
            i.read(static_cast<char *>(p), n);
            return i.good();
        }

    public:
        static boost::optional<uint64_t> varint(std::istream &i) {
            uint64_t n = 0;
            // Can read max 9 bytes, each with 7 bits of payload, for 9*7 = 63 bits.
            for (size_t count = 0; count < 9; count++) {
                int b = i.get();
                if (i.eof()) return boost::none;
                n |= (b & 0x7f) << (count * 7);
                if ((b & 0x80) == 0) {
                    return n;
                }
            }
            return boost::none;
        }

        BinaryReader(std::istream& i, std::function<std::shared_ptr<T>(Value<>)> decodeEmbedded) :
            i(i),
            decodeEmbedded(decodeEmbedded)
        {}

        boost::optional<float> next_float() {
            uint8_t buf[4];
            if (!next_chunk(buf, sizeof(buf))) return boost::none;
            uint32_t n = buf[0] << 24 | buf[1] << 16 | buf[2] << 8 | buf[3];
            float f;
            memcpy(&f, &n, sizeof(f));
            return f;
        }

        boost::optional<double> next_double() {
            uint8_t buf[8];
            if (!next_chunk(buf, sizeof(buf))) return boost::none;
            uint32_t n1 = buf[0] << 24 | buf[1] << 16 | buf[2] << 8 | buf[3];
            uint32_t n2 = buf[4] << 24 | buf[5] << 16 | buf[6] << 8 | buf[7];
            uint64_t n = uint64_t(n1) << 32 | n2;
            double d;
            memcpy(&d, &n, sizeof(d));
            return d;
        }

        boost::optional<uint64_t> next_unsigned(size_t n) {
            uint8_t buf[n];
            if (!next_chunk(buf, n)) return boost::none;
            uint64_t v = 0;
            for (size_t j = 0; j < n; j++) {
                v = (v << 8) | buf[j];
            }
            return v;
        }

        boost::optional<Value<T>> next_bignum(size_t n) {
            auto b = std::make_shared<BigNum<T>>(std::vector<uint8_t>());
            b->_value().resize(n);
            if (!next_chunk(&b->_value()[0], n)) return boost::none;
            return Value<T>(b);
        }

        boost::optional<Value<T>> next() {
            bool end_sentinel;
            return _next(end_sentinel);
        }

        boost::optional<Value<T>> _next(bool& end_sentinel) {
        more:
            end_sentinel = false;
            auto tag = BinaryTag(i.get());
            // std::cout << "tag " << std::hex << int(tag) << " pos " << i.tellg() - 1 << std::endl;
            if (i.eof()) return boost::none;
            switch (tag) {
                case BinaryTag::False: return Value<T>::from_bool(false);
                case BinaryTag::True: return Value<T>::from_bool(true);
                case BinaryTag::Float: return next_float().map(Value<T>::from_float);
                case BinaryTag::Double: return next_double().map(Value<T>::from_double);
                case BinaryTag::End: end_sentinel = true; return boost::none;
                case BinaryTag::Annotation:
                    if (!next()) return boost::none;
                    goto more;
                case BinaryTag::Embedded:
                    return BinaryReader<GenericEmbedded>(i, &GenericEmbedded::wrap)
                        .next().map(decodeEmbedded).map(Value<T>::from_embedded);
                case BinaryTag::SmallInteger_lo ... BinaryTag::SmallInteger_hi: {
                    int64_t n = int64_t(tag) - int64_t(BinaryTag::SmallInteger_lo);
                    return Value<T>::from_int(n <= 12 ? n : n - 16);
                }
                case BinaryTag::MediumInteger_lo ... BinaryTag::MediumInteger_hi: {
                    int n = int(tag) - int(BinaryTag::MediumInteger_lo) + 1;
                    if (n < 9) {
                        return next_unsigned(n).map([](uint64_t v) { return Value<T>::from_int(int64_t(v)); });
                    }
                    if (n == 9) {
                        // We can handle this with uint64_t if it's unsigned and the first byte is 0.
                        if (i.get() == 0) return next_unsigned(8).map(Value<T>::from_unsigned);
                    }
                    return next_bignum(n);
                }
                case BinaryTag::SignedInteger: return varint(i).flat_map([&](size_t len) {
                    return next_bignum(len);
                });
                case BinaryTag::String: return varint(i).flat_map([&](size_t len)-> boost::optional<Value<T>> {
                    auto s = std::make_shared<String<T>>(std::string());
                    s->_value().resize(len);
                    if (!next_chunk(&s->_value()[0], len)) return boost::none;
                    return Value<T>(s);
                });
                case BinaryTag::ByteString: return varint(i).flat_map([&](size_t len)-> boost::optional<Value<T>> {
                    auto s = std::make_shared<ByteString<T>>(std::vector<uint8_t>());
                    s->_value().resize(len);
                    if (!next_chunk(&s->_value()[0], len)) return boost::none;
                    return Value<T>(s);
                });
                case BinaryTag::Symbol: return varint(i).flat_map([&](size_t len)-> boost::optional<Value<T>> {
                    auto s = std::make_shared<Symbol<T>>(std::string());
                    s->_value().resize(len);
                    if (!next_chunk(&s->_value()[0], len)) return boost::none;
                    return Value<T>(s);
                });
                case BinaryTag::Record: return next().flat_map([&](auto label)-> boost::optional<Value<T>> {
                    auto r = std::make_shared<Record<T>>(label);
                    while (true) {
                        bool end_rec = false;
                        auto v = _next(end_rec);
                        if (end_rec) return Value<T>(r);
                        if (!v) return boost::none;
                        r->fields.push_back(*v);
                    }
                });
                case BinaryTag::Sequence: {
                    auto s = std::make_shared<Sequence<T>>();
                    while (true) {
                        bool end_rec = false;
                        auto v = _next(end_rec);
                        if (end_rec) return Value<T>(s);
                        if (!v) return boost::none;
                        s->values.push_back(*v);
                    }
                }
                case BinaryTag::Set: {
                    auto s = std::make_shared<Set<T>>();
                    while (true) {
                        bool end_rec = false;
                        auto v = _next(end_rec);
                        if (end_rec) return Value<T>(s);
                        if (!v) return boost::none;
                        s->values.insert(*v);
                    }
                }
                case BinaryTag::Dictionary: {
                    auto s = std::make_shared<Dictionary<T>>();
                    while (true) {
                        bool end_rec = false;
                        auto k = _next(end_rec);
                        if (end_rec) return Value<T>(s);
                        if (!k) return boost::none;
                        auto v = next();
                        if (!v) return boost::none;
                        s->values.emplace(*k, *v);
                    }
                }
                default:
                    return boost::none;
            }
        }
    };
}