#pragma once
#include "asset.h"
#include <codecvt>
#include <fmt/core.h>
#include "util.h"
#include "image.h"
#include "brush.h"
#include "log.h"

class BinaryStream
{
public:
    enum class ByteOrder : uint8_t { BigEndian, LittleEndian, Host };
    static ByteOrder sys_order()
    {
        union {
            uint32_t i;
            char c[4];
        } bint { 0x01020304 };

        return bint.c[0] == 1 ? ByteOrder::BigEndian : ByteOrder::LittleEndian;
    }
    template<typename T> T swap(T x)
    {
#if _MSC_VER >= 1400

        if (sizeof(T) == 2)
        {
            auto y = _byteswap_ushort(*reinterpret_cast<uint16_t*>(&x));
            return *reinterpret_cast<T*>(&y);
        }
        else if (sizeof(T) == 4)
        {
            auto y = _byteswap_ulong(*reinterpret_cast<uint32_t*>(&x));
            return *reinterpret_cast<T*>(&y);
        }
        else if (sizeof(T) == 8)
        {
            auto y = _byteswap_uint64(*reinterpret_cast<uint64_t*>(&x));
            return *reinterpret_cast<T*>(&y);
        }
#else
        auto p = reinterpret_cast<uint8_t*>(&x);
        if (sizeof(T) == 2)
        {
            std::swap(p[0], p[1]);
        }
        else if (sizeof(T) == 4)
        {
            std::swap(p[0], p[3]);
            std::swap(p[1], p[2]);
        }
        else if (sizeof(T) == 8)
        {
            std::swap(p[0], p[7]);
            std::swap(p[1], p[6]);
            std::swap(p[2], p[5]);
            std::swap(p[3], p[4]);
        }
        else
        {
            static_assert(true, "Should not reach here");
        }
        return x;
#endif
    }
protected:
    bool m_swap = false;
    ByteOrder m_byte_order = ByteOrder::Host;
};

class BinaryStreamReader : public BinaryStream
{
public:
    BinaryStreamReader(const BinaryStreamReader&) = delete;
    BinaryStreamReader() = default;
    ~BinaryStreamReader()
    {
        m_ptr = m_cur = nullptr;
        m_size = 0;
    }
    void init(uint8_t* data_ptr, size_t size, ByteOrder byte_order = ByteOrder::Host)
    {
        m_ptr = m_cur = data_ptr;
        m_size = size;
        m_byte_order = byte_order;
        m_swap = byte_order == ByteOrder::Host ? false : byte_order != sys_order();
    }
    size_t pos() { return std::distance(m_ptr, m_cur); }
    void skip(size_t bytes) { m_cur += bytes; }
    // snap to the next 4-alignment
    void snap() { if ((size_t)m_cur % 4 != 0) m_cur += 4 - (size_t)m_cur % 4; }
    bool eof() { return std::distance(m_ptr, m_cur) >= m_size; }
    bool has_data(size_t sz) { return std::distance(m_ptr, m_cur + sz) <= m_size; }
    uint8_t  ru8()  { return read<uint8_t>(); }
    uint16_t ru16() { return m_swap ? swap(read<uint16_t>()) : read<uint16_t>(); }
    uint32_t ru32() { return m_swap ? swap(read<uint32_t>()) : read<uint32_t>(); }
    uint64_t ru64() { return m_swap ? swap(read<uint64_t>()) : read<uint64_t>(); }
    int8_t   ri8()  { return read<int8_t>(); }
    int16_t  ri16() { return m_swap ? swap(read<int16_t>()) : read<int16_t>(); }
    int32_t  ri32() { return m_swap ? swap(read<int32_t>()) : read<int32_t>(); }
    int64_t  ri64() { return m_swap ? swap(read<int64_t>()) : read<int64_t>(); }
    float    rflt() { return m_swap ? swap(read<float>()) : read<float>(); }
    double   rdbl() { return m_swap ? swap(read<double>()) : read<double>(); }
    std::string pick(size_t chars) { return { (char*)m_cur, chars }; }
    std::string rstring() 
    {
        auto len = ru32();
        return { advance<char>(len), len };
    }
    std::string rstring(size_t len) { return { advance<char>(len), len }; }
    std::wstring rwstring() 
    { 
        auto len = ru32();
        return rwstring(len); 
    }
    std::wstring rwstring(size_t len) 
    { 
        auto ptr = advance<char>(len * 2);
        //for (int i = 0; i < len; i++)
        //    std::swap(ptr[i * 2], ptr[i * 2 + 1]);

        // right trim trailing zeroes
        auto wptr = (uint16_t*)ptr;
        for (int i = len - 1; i >= 0; i--)
            if (wptr[i] == 0) len--;

        // wide to UTF-16le
        std::wstring_convert<std::codecvt_utf16<wchar_t, 0x10ffff>> converter;
        return converter.from_bytes(ptr, ptr + len * 2);
    }
    std::string rpascal()
    {
        auto len = ru8();
        return rstring(len);
    }
    std::vector<uint8_t> rraw()
    {
        auto size = ru32();
        return rraw(size);
    }
    std::vector<uint8_t> rraw(size_t bytes)
    {
        std::vector<uint8_t> ret(bytes);
        std::copy_n(m_cur, bytes, ret.data());
        skip(bytes);
        return ret;
    }
    std::vector<uint8_t> rrle(size_t encoded_bytes)
    {
        std::vector<uint8_t> data;
        int32_t n;
        for (int j = 0; j < encoded_bytes;)
        {
            n = ri8();
            j++;
            // copy the following char -n + 1 times
            if (n < 0)
            {
                // NOP
                if (n == -128)
                    continue;
                n = -n + 1;

                j++;
                auto ch = ru8();
                for (int c = 0; c < n; c++)
                    data.push_back(ch);
            }
            else
            {
                // read the following n + 1 chars (no compr)
                for (int c = 0; c < n + 1; c++, j++)
                    data.push_back(ru8());
            }
        }
        return data;
    }
protected:
    template<typename T> inline T align4(T x) { return ((x - T{1}) & (~(T{3}))) + T{4}; }
    template<typename T> T read()
    {
        T ret{};
        if (m_ptr)
        {
            if (has_data(sizeof(T)))
                ret = *reinterpret_cast<T*>(m_cur);
            m_cur += sizeof(T);
        }
        return ret;
    }
    template<typename T> T* advance(size_t bytes)
    {
        T* ret = nullptr;
        if (m_ptr)
        {
            if (has_data(bytes))
                ret = reinterpret_cast<T*>(m_cur);
            m_cur += bytes;
        }
        return ret;
    }
private:
    uint8_t *m_ptr = nullptr;
    uint8_t *m_cur = nullptr;
    size_t m_size = 0;
};

class BinaryStreamWriter : public BinaryStream
{
public:
    std::vector<uint8_t> m_data;
    BinaryStreamWriter(const BinaryStreamWriter&) = delete;
    BinaryStreamWriter() = default;
    ~BinaryStreamWriter() = default;
    void init(ByteOrder byte_order = ByteOrder::Host)
    {
        m_byte_order = byte_order;
        m_swap = byte_order == ByteOrder::Host ? false : byte_order != sys_order();
    }
    //size_t pos() { return std::distance(m_ptr, m_cur); }
    void skip(size_t bytes, uint8_t fill = 0) { m_data.resize(m_data.size() + bytes); }
    // snap to the next 4-alignment
    void snap() { if ((size_t)m_data.size() % 4 != 0) m_data.resize(m_data.size() + 4 - (size_t)m_data.size() % 4); }
    //bool eof() { return std::distance(m_ptr, m_cur) >= m_size; }
    //bool has_data(size_t sz) { return std::distance(m_ptr, m_cur + sz) < m_size; }
    void wu8 (uint8_t v ) { return write<uint8_t>(v); }
    void wu16(uint16_t v) { return m_swap ? write(swap(v)) : write(v); }
    void wu32(uint32_t v) { return m_swap ? write(swap(v)) : write(v); }
    void wu64(uint64_t v) { return m_swap ? write(swap(v)) : write(v); }
    void wi8 (int8_t v  ) { return write<int8_t>(v); }
    void wi16(int16_t v ) { return m_swap ? write(swap(v)) : write(v); }
    void wi32(int32_t v ) { return m_swap ? write(swap(v)) : write(v); }
    void wi64(int64_t v ) { return m_swap ? write(swap(v)) : write(v); }
    void wflt(float v   ) { return m_swap ? write(swap(v)) : write(v); }
    void wdbl(double v  ) { return m_swap ? write(swap(v)) : write(v); }
    //std::string pick(size_t chars) { return { (char*)m_cur, chars }; }
    void wstring(std::string s)
    {
        wu32(s.size());
        write(s.data(), s.size());
    }
    void wwstring(std::wstring s)
    {
        wu32(s.size());
        std::wstring_convert<std::codecvt_utf16<wchar_t, 0x10ffff>> converter;
        wstring(converter.to_bytes(s));
    }
    void wpascal(std::string s)
    {
        wu8(s.size());
        write(s.data(), s.size());
    }
    void wraw(std::vector<uint8_t> raw)
    {
        wu32(raw.size());
        write(raw.data(), raw.size());
    }
    void wrle(std::vector<uint8_t> data)
    {
        int idx = 0;
        auto ptr = data.data();
        
        // find the next sequence and return the position
        auto next_seq = [&] () -> std::pair<size_t, size_t> {
            size_t i = idx;
            while (i < data.size() - 1)
            {
                size_t count = 0;
                for (size_t j = i + 1; j < data.size() && data[i] == data[j]; j++)
                    count++;
                if (count > 0)
                    return { i, count + 1 };
                i++;
            }
            return { i + 1, 0 };
        };

        while (idx < data.size())
        {
            auto seq = next_seq();
            // non sequence chars
            if (seq.first - idx > 0)
            {
                wi8(seq.first - idx - 1);
                write(ptr + idx, seq.first - idx);
            }
            if (seq.second > 0)
            {
                wi8(-(int8_t)(seq.second - 1));
                write(ptr[seq.first]);
            }
            idx = seq.first + seq.second;
        }
    }
protected:
    template<typename T> void write(T x)
    {
        uint8_t* bytes = reinterpret_cast<uint8_t*>(&x);
        m_data.insert(m_data.end(), bytes, bytes + sizeof(T));
    }
    template<typename T> void write(const T* ptr, size_t size)
    {
        const uint8_t* bytes = reinterpret_cast<const uint8_t*>(ptr);
        m_data.insert(m_data.end(), bytes, bytes + size * sizeof(T));
    }
};

class SerializedStream
{
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";
        }
    };
    struct Class : public Type { };
    struct Property : public Type { };
    struct Reference : public Type { };
    struct List : public Type
    {
        Type::Vec items;
        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;
        }
    };
    struct Double : public Type
    {
        using native_type = double;
        double value;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("double: {}", value); }
    };
    struct UnitFloat : public Type
    {
        using native_type = double;
        std::string unit;
        double value;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("float: {} ({})", value, unit); }
    };
    struct String : public Type 
    {
        using native_type = std::wstring;
        std::wstring value;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("string: {}", wstr2str(value)); }
    };
    struct Enum : public Type
    {
        std::string type;
        std::string value;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("enum {}: {}", type, value); }
    };
    struct EnumRef : public Type { };
    struct Offset : public Type { };
    struct Identifier : public Type { };
    struct Index : public Type { };
    struct Name : public Type { };
    struct Integer : public Type 
    { 
        using native_type = int32_t;
        int32_t value;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("int: {}", value); }
    };
    struct LargeInteger : public Type { };
    struct Boolean : public Type 
    {
        using native_type = bool;
        bool value; 
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("bool: {}", value); }
    };
    struct Alias : public Type { };
    struct RawData : public Type
    {
        std::vector<uint8_t> data;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("raw: {} bytes", data.size()); }
    };
    struct Descriptor : public Type 
    {
        std::wstring name;
        std::string class_id;
        Type::Map props;
        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> 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> 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> void value(const std::string& key, D& dest) const
        {
            if (auto v = get<T>(key))
                dest = static_cast<D>(v->value);
        }
    };
    struct Rectangle : public Type
    {
        uint32_t top;
        uint32_t left;
        uint32_t bottom;
        uint32_t 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); }
    };
    struct Point : public Type
    {
        uint32_t x;
        uint32_t y;
        virtual std::string str(int indent, const std::string& prefix) const override
        { return std::string(indent, '-') + prefix + fmt::format("point: [{}, {}]", x, 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)) { } 
    };
    struct VMArray : public Type
    {
        uint32_t version; // = 3
        Rectangle rect;
        std::vector<Channel> channels;
        VMArray() = default;
        VMArray(uint32_t version, const Rectangle& rect) : version(version), rect(rect) { }
        std::shared_ptr<Image> image(bool grayscale, bool invert) const
        {
            int nc = channels.size();
            auto pixels = (channels[0].depth >> 3) * rect.area();
            if (nc == 1 || nc >= 3)
            {
                auto img = std::make_shared<Image>();
                img->comp = 4;
                img->width = rect.width();
                img->height = rect.height();
                img->m_data = std::make_unique<uint8_t[]>(pixels * 4);
                auto out = reinterpret_cast<glm::u8vec4*>(img->m_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 %d channels\n", channels.size());
            }
            return nullptr;
        }
    };
};

class SerializedStreamReader : public SerializedStream, public BinaryStreamReader
{
public:
    std::shared_ptr<VMArray>    parse_vmem(); // Parse Virtual Memory Array List
    std::shared_ptr<List>       parse_vlls();
    std::shared_ptr<String>     parse_text();
    std::shared_ptr<Descriptor> parse_objc();
    std::shared_ptr<UnitFloat>  parse_untf();
    std::shared_ptr<Boolean>    parse_bool();
    std::shared_ptr<Integer>    parse_long();
    std::shared_ptr<Double>     parse_doub();
    std::shared_ptr<Enum>       parse_enum();
    std::shared_ptr<RawData>    parse_tdta();

    std::map<std::string /*key*/, std::function<Type::Ref()>> m_parser_table;
    SerializedStreamReader();

    std::string rkey_or_string()
    {
        auto len = ru32();
        if (len == 0)
            len = 4;
        return rstring(len);
    }
    Rectangle rrect()
    {
        Rectangle ret;
        ret.top = ru32();
        ret.left = ru32();
        ret.bottom = ru32();
        ret.right = ru32();
        return ret;
    }
    Point rpoint()
    {
        Point ret;
        ret.x = ru16();
        ret.y = ru16();
        return ret;
    }

    Type::Ref call(std::string t)
    {
        if (m_parser_table.find(t) != m_parser_table.end())
        {
            auto& method = m_parser_table[t];
            return method();
        }
        return nullptr;
    }
};

class ABR : private SerializedStreamReader
{
    bool section_desc();
    bool section_samp();
    bool section_patt();
public:
    std::vector<std::shared_ptr<Descriptor>> m_presets;
    std::map<std::string /*uid*/, std::shared_ptr<Image>> m_patterns;
    std::map<std::string /*uid*/, std::shared_ptr<Image>> m_samples;

    bool open(const std::string& path);
    std::vector<std::shared_ptr<Brush>> compute_brushes(const std::string& path);
};