#pragma once
#include "binary_stream.h"
#include "util.h"
#include "log.h"

template<typename T> BinaryStreamWriter& operator<<(BinaryStreamWriter& w, const T& obj)
{
    obj.write(w);
    return w;
}

template<typename T> BinaryStreamReader& operator>>(BinaryStreamReader& r, T& obj)
{
    obj.read(r);
    return r;
}

class Serializer
{
public:
    struct Type
    {
        using Vec = std::vector<std::shared_ptr<Type>>;
        using Map = std::map<std::string, std::shared_ptr<Type>>;
        using Ref = std::shared_ptr<Type>;
        virtual std::string str(int indent, const std::string& prefix) const
        {
            return "type";
        }
        virtual bool read(BinaryStreamReader& r) { return false; }
        virtual void write(BinaryStreamWriter& w) const { }
        virtual std::string type_key() const { return ""; }
    };
    struct Class : public Type { };
    struct Property : public Type { };
    struct Reference : public Type { };
    struct EnumRef : public Type { };
    struct Offset : public Type { };
    struct Identifier : public Type { };
    struct Index : public Type { };
    struct Name : public Type { };
    struct LargeInteger : public Type { };
    struct Alias : public Type { };
    struct List : public Type
    {
        Type::Vec items;
        virtual std::string type_key() const override { return "VlLs"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            auto ret = std::string(indent, '-') + fmt::format("list: {} items:", items.size());
            for (int i = 0; i < items.size(); i++)
                ret += "\n" + items[i]->str(indent + 1, fmt::format("{}) ", i));
            return ret;
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            auto count = r.ru32();
            for (int i = 0; i < count; i++)
            {
                auto type = r.rstring(4);
                auto item = instanciate(type);
                if (!item || !item->read(r))
                    return false;
                items.push_back(item);
            }
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wi32((int)items.size());
            for (auto& i : items)
            {
                w.wstring_raw(i->type_key());
                i->write(w);
            }
        }
        template <typename T>
        std::shared_ptr<T> add()
        {
            auto ptr = std::make_shared<T>();
            items.emplace_back(ptr);
            return ptr;
        }
    };
    struct Double : public Type
    {
        using native_type = double;
        double value;
        Double() = default;
        Double(double value) : value(value) { }
        virtual std::string type_key() const override { return "doub"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("double: {}", value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.rdbl();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wdbl(value);
        }
    };
    struct Float : public Type
    {
        using native_type = float;
        float value;
        Float() = default;
        Float(float value) : value(value) { }
        virtual std::string type_key() const override { return "flt "; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("float: {}", value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.rflt();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wflt(value);
        }
    };
    struct UnitFloat : public Type
    {
        using native_type = double;
        static constexpr const char* UnitAngle   = "#Ang"; // angle: base degrees
        static constexpr const char* UnitDensity = "#Rsl"; // density: base per inch
        static constexpr const char* UnitDist    = "#Rlt"; // distance: base 72ppi
        static constexpr const char* UnitNone    = "#Nne"; // none: coerced.
        static constexpr const char* UnitPercent = "#Prc"; // percent: unit value
        static constexpr const char* UnitPixel   = "#Pxl"; // pixels: tagged unit value
        std::string unit;
        double value;
        UnitFloat() = default;
        UnitFloat(const std::string& unit, double value) : unit(unit), value(value) { }
        virtual std::string type_key() const override { return "UntF"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("float: {} ({})", value, unit);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            unit = r.rstring(4);
            value = r.rdbl();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wstring_raw(unit);
            w.wdbl(value);
        }
    };
    struct String : public Type
    {
        using native_type = std::wstring;
        std::wstring value;
        String() = default;
        String(const std::wstring& s) : value(s) { }
        String(const std::string& s) : value(str2wstr(s)) { }
        virtual std::string type_key() const override { return "TEXT"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("string: {}", wstr2str(value));
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.rwstring();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wwstring(value);
        }
    };
    struct CString : public Type
    {
        using native_type = std::string;
        std::string value;
        CString() = default;
        CString(const std::string& s) : value(s) { }
        virtual std::string type_key() const override { return "cstr"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("cstring: {}", value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.rstring();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wstring(value);
        }
    };
    struct Vec2 : public Type
    {
        using native_type = glm::vec2;
        glm::vec2 value;
        Vec2() = default;
        Vec2(glm::vec2 v) : value(v) { }
        virtual std::string type_key() const override { return "vec2"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("vec2: {} {}", value.x, value.y);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.rflt();
            value.y = r.rflt();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wflt(value.x);
            w.wflt(value.y);
        }
    };
    struct Vec3 : public Type
    {
        using native_type = glm::vec3;
        glm::vec3 value;
        Vec3() = default;
        Vec3(glm::vec3 v) : value(v) { }
        virtual std::string type_key() const override { return "vec3"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("vec3: {} {} {}", value.x, value.y, value.z);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.rflt();
            value.y = r.rflt();
            value.z = r.rflt();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wflt(value.x);
            w.wflt(value.y);
            w.wflt(value.z);
        }
    };
    struct Vec4 : public Type
    {
        using native_type = glm::vec4;
        glm::vec4 value;
        Vec4() = default;
        Vec4(glm::vec4 v) : value(v) { }
        virtual std::string type_key() const override { return "vec4"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("vec4: {} {} {} {}", value.x, value.y, value.z, value.w);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.rflt();
            value.y = r.rflt();
            value.z = r.rflt();
            value.w = r.rflt();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wflt(value.x);
            w.wflt(value.y);
            w.wflt(value.z);
            w.wflt(value.w);
        }
    };
    struct IVec2 : public Type
    {
        using native_type = glm::ivec2;
        glm::ivec2 value;
        IVec2() = default;
        IVec2(glm::ivec2 v) : value(v) { }
        virtual std::string type_key() const override { return "ive2"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("ivec2: {} {}", value.x, value.y);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.ri32();
            value.y = r.ri32();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wi32(value.x);
            w.wi32(value.y);
        }
    };
    struct IVec3 : public Type
    {
        using native_type = glm::ivec3;
        glm::ivec3 value;
        IVec3() = default;
        IVec3(glm::ivec3 v) : value(v) { }
        virtual std::string type_key() const override { return "ive3"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("ivec3: {} {} {}", value.x, value.y, value.z);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.ri32();
            value.y = r.ri32();
            value.z = r.ri32();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wu32(value.x);
            w.wu32(value.y);
            w.wu32(value.z);
        }
    };
    struct IVec4 : public Type
    {
        using native_type = glm::ivec4;
        glm::ivec4 value;
        IVec4() = default;
        IVec4(glm::ivec4 v) : value(v) { }
        virtual std::string type_key() const override { return "ive4"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("ivec4: {} {} {} {}", value.x, value.y, value.z, value.w);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value.x = r.ri32();
            value.y = r.ri32();
            value.z = r.ri32();
            value.w = r.ri32();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wi32(value.x);
            w.wi32(value.y);
            w.wi32(value.z);
            w.wi32(value.w);
        }
    };
    struct Enum : public Type
    {
        std::string type;
        std::string value;
        Enum() = default;
        Enum(const std::string& type, const std::string& value) : type(type), value(value) { }
        virtual std::string type_key() const override { return "enum"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("enum {}: {}", type, value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            type = r.rkey_or_string();
            value = r.rkey_or_string();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wkey_or_string(type);
            w.wkey_or_string(value);
        }
    };
    struct Integer : public Type
    {
        using native_type = int32_t;
        int32_t value;
        Integer() = default;
        Integer(int32_t value) : value(value) { }
        virtual std::string type_key() const override { return "long"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("int: {}", value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.ri32();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wi32(value);
        }
    };
    struct Boolean : public Type
    {
        using native_type = bool;
        bool value;
        Boolean() = default;
        Boolean(bool value) : value(value) { }
        virtual std::string type_key() const override { return "bool"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("bool: {}", value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            value = r.ru8();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wu8(value);
        }
    };
    struct RawData : public Type
    {
        std::vector<uint8_t> data;
        RawData() = default;
        //RawData(const RawData&) = delete;
        RawData(std::vector<uint8_t> data) : data(data) { }
        virtual std::string type_key() const override { return "tdta"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("raw: {} bytes", data.size());
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            data = r.rraw();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wraw(data);
        }
    };
    struct Descriptor : public Type
    {
        std::wstring name;
        std::string class_id = "desc";
        Type::Map props;
        Descriptor() = default;
        virtual std::string type_key() const override { return "Objc"; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            auto ret = std::string(indent, '-') + prefix +
                fmt::format("objc {} ({}): {} props:", wstr2str(name), class_id, props.size());
            for (const auto& p : props)
                ret += "\n" + p.second->str(indent + 1, fmt::format("'{}' ", p.first));
            return ret;
        }
        bool has(const std::string& key) const
        {
            return props.find(key) != props.end();
        }
        template<typename T> std::shared_ptr<T> get(const std::string& key) const
        {
            return has(key) ? std::dynamic_pointer_cast<T>(props.at(key)) : nullptr;
        }
        template<typename T> std::shared_ptr<T> set(const std::string& key, const T value)
        {
            if (auto ptr = get<T>(key))
            {
                *ptr = value;
                return ptr;
            }
            else
            {
                ptr = std::make_shared<T>(value);
                props[key] = ptr;
                return ptr;
            }
        }
        template<typename T> auto value(const std::string& key) const
        {
            if (auto v = get<T>(key))
                return v->value;
            return decltype(T::value){};
        }
        template<typename T, typename D = decltype(T::value)> auto value_or(const std::string& key, const D val) const
        {
            if (auto v = get<T>(key))
                return v->value;
            return val;
        }
        template<typename T, typename D = decltype(T::value)> void value(const std::string& key, D& dest) const
        {
            if (auto v = get<T>(key))
                dest = static_cast<D>(v->value);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            name = r.rwstring();
            class_id = r.rkey_or_string();
            auto count = r.ru32();
            for (int i = 0; i < count; i++)
            {
                auto key = r.rkey_or_string();
                auto type = r.rstring(4);
                //printf("prop %s\n", t.c_str());
                auto p = instanciate(type);
                if (!p || !p->read(r))
                    return false;
                if (props.find(key) != props.end())
                    LOG("DUPLICATE prop %s\n", key.c_str());
                props[key] = p;
            }
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wwstring(name);
            w.wkey_or_string(class_id);
            w.wu32((int)props.size());
            for (auto& p : props)
            {
                w.wkey_or_string(p.first);
                w.wstring_raw(p.second->type_key());
                p.second->write(w);
            }
        }
    };
    struct Rectangle : public Type
    {
        uint32_t top;
        uint32_t left;
        uint32_t bottom;
        uint32_t right;
        Rectangle() = default;
        Rectangle(uint32_t top, uint32_t left, uint32_t bottom, uint32_t right) :
            top(top), left(left), bottom(bottom), right(right) { }
        uint32_t area() const { return (right - left) * (bottom - top); }
        uint32_t width() const { return right - left; }
        uint32_t height() const { return bottom - top; }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("rect: [{}, {}, {}, {}]", top, left, bottom, right);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            top = r.ru32();
            left = r.ru32();
            bottom = r.ru32();
            right = r.ru32();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wu32(top);
            w.wu32(left);
            w.wu32(bottom);
            w.wu32(right);
        }
    };
    struct Point : public Type
    {
        uint32_t x;
        uint32_t y;
        Point() = default;
        Point(uint32_t x, uint32_t y) : x(x), y(y) { }
        virtual std::string str(int indent, const std::string& prefix) const override
        {
            return std::string(indent, '-') + prefix + fmt::format("point: [{}, {}]", x, y);
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            x = r.ru16();
            y = r.ru16();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
            w.wu16(x);
            w.wu16(y);
        }
    };
    struct Channel : public Type
    {
        uint32_t depth;
        Rectangle rect;
        uint8_t compression;
        std::vector<uint8_t> data;
        Channel() = default;
        Channel(uint32_t depth, Rectangle rect, uint8_t compression, std::vector<uint8_t> data) :
            depth(depth), rect(rect), compression(compression), data(std::move(data)) { } 
        virtual bool read(BinaryStreamReader& r) override
        {
            auto length = r.ru32(); // skip if 0, length is from the next field to the end
            if (length == 0)
                return true;
            depth = r.ru32(); // Pixel depth: 1, 8, 16 or 32
            rect.read(r);
            auto depth2 = r.ru16(); // again?
            compression = r.ru8(); // 1 = zip
            if (depth != 8)
            {
                LOG("unsupported depth %d bits\n", depth);
                r.skip(length - 23);
                return true;
            }
            if (compression == 0)
            {
                int data_size = (depth >> 3) * rect.area();
                data = r.rraw(data_size);
            }
            else if (compression == 1)
            {
                auto start = r.pos();
                auto height = rect.height();
                // contain the compressed length of each scanline
                std::vector<uint16_t> scanlines;
                scanlines.reserve(height);
                for (int i = 0; i < height; i++)
                    scanlines.push_back(r.ru16());
                for (auto sl : scanlines)
                {
                    auto decoded = r.rrle(sl);
                    data.insert(data.end(), decoded.begin(), decoded.end());
                }
                //auto len = r.pos() - start;
            }
            else
            {
                LOG("unsupported compression mode %d\n", compression);
                r.skip(length - 23);
                return true;
            }
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
        }
    };
/*
    struct ImagePNG : public Type
    {
        using native_type = Image;
        Image data;
        
        bool read(BinaryStreamReader& r) override
        {
            Descriptor d;
            if (d.class_id != "image_png")
                return false;
            r >> d;
            d.value<Integer>("width", data.width);
            d.value<Integer>("height", data.height);
            d.value<Integer>("comp", data.comp);
            data.file_base = wstr2str(d.value<Serializer::String>("file_base"));
            data.file_name = wstr2str(d.value<Serializer::String>("file_name"));
            data.file_ext = wstr2str(d.value<Serializer::String>("file_ext"));
            auto img_raw = d.get<RawData>("data");
            int png_width, png_height, png_comp;
            data.m_data = std::unique_ptr<uint8_t[]>(stbi_load_from_memory(
                img_raw->data.data(), img_raw->data.size(), &png_width, &png_height, &png_comp, 4));
            return true;
        }

        void write(BinaryStreamWriter& w) const override
        {
            Descriptor d;
            d.class_id = "image_png";
            d.name = L"Image class";

            d.props["width"] = std::make_shared<Integer>(data.width);
            d.props["height"] = std::make_shared<Integer>(data.height);
            d.props["comp"] = std::make_shared<Integer>(data.comp);
            d.props["file_base"] = std::make_shared<Serializer::String>(data.file_base);
            d.props["file_name"] = std::make_shared<Serializer::String>(data.file_name);
            d.props["file_ext"] = std::make_shared<Serializer::String>(data.file_ext);

            // really ugly way to compress the png and store it with a lambda
            stbi_write_png_to_func([](void* context, void* data, int size) {
                Descriptor& d = *static_cast<Descriptor*>(context);
                d.props["data"] = std::make_shared<RawData>(std::vector<uint8_t>((uint8_t*)data, (uint8_t*)data + size));
            }, &d, data.width, data.height, data.comp, data.m_data.get(), 0);

            w << d;
        }
    };
*/
    struct VMArray : public Type
    {
        struct ImageData
        {
            std::unique_ptr<uint8_t[]> data;
            size_t size = 0;
            int width = 0;
            int height = 0;
            int comp = 0;
        };
        uint32_t version; // = 3
        Rectangle rect;
        std::vector<Channel> channels;
        VMArray() = default;
        VMArray(uint32_t version, const Rectangle& rect) : version(version), rect(rect) { }
        ImageData image(bool grayscale, bool invert) const
        {
            int nc = (int)channels.size();
            auto pixels = (channels[0].depth >> 3) * rect.area();
            if (nc == 1 || nc >= 3)
            {
                ImageData img;
                img.comp = 4;
                img.width = rect.width();
                img.height = rect.height();
                img.size = pixels * 4;
                img.data = std::make_unique<uint8_t[]>(pixels * 4);
                auto out = reinterpret_cast<glm::u8vec4*>(img.data.get());
                if (grayscale)
                {
                    auto const& raw = channels[0].data;
                    if (invert)
                    {
                        for (int i = 0; i < raw.size(); i++)
                            out[i] = glm::u8vec4(glm::u8vec3(255 - raw[i]), 255);
                    }
                    else
                    {
                        for (int i = 0; i < raw.size(); i++)
                            out[i] = glm::u8vec4(glm::u8vec3(raw[i]), 255);
                    }
                }
                else
                {
                    std::fill_n(out, pixels, glm::u8vec4(255));
                    for (int ch = 0; ch < std::min(nc, 3); ch++)
                    {
                        auto const& raw = channels[ch].data;
                        if (invert)
                        {
                            for (int i = 0; i < raw.size(); i++)
                                out[i][ch] = 255 - raw[i];
                        }
                        else
                        {
                            for (int i = 0; i < raw.size(); i++)
                                out[i][ch] = raw[i];
                        }
                    }
                }
                return img;
                //stbi_write_png(fmt::format("x64/out/{}.png", uid).c_str(), 
                //    image->rect.width(), image->rect.height(), 4, out.data(), 0);
            }
            else
            {
                LOG("Error image with %ld channels\n", channels.size());
            }
            return {};
        }
        virtual bool read(BinaryStreamReader& r) override
        {
            // Virtual Memory Array List
            version = r.ru32(); // = 3
            assert(version == 3);
            auto vmem_length = r.ru32();
            // TODO: check if at the end there's good data
            // check if the bounds are within the parent's size
            rect.read(r);
            auto vmem_channels = r.ru32();
            // The following is a virtual memory array,
            // repeated for the number of channels 
            // + one for a user mask + one for a sheet mask.
            vmem_channels += 2; // user and sheet mask
            for (int ch = 0; ch < vmem_channels; ch++)
            {
                auto array_written = r.ru32(); // skip if 0
                if (array_written == 0)
                    continue;
                channels.emplace_back();
                if (!channels.back().read(r))
                    return false;
            }
            r.snap();
            return true;
        }
        virtual void write(BinaryStreamWriter& w) const override
        {
        }
    };
    template<typename T> static T parse(BinaryStreamReader& r)
    {
        T ret{};
        ret.read(r);
        return ret;
    }
protected:
    static Type::Ref instanciate(const std::string& key);
    static std::map<std::string /*key*/, std::function<Type::Ref()>> m_ctor_table;
};