#include "ImXui/XurFormat.h"
#include <fstream>
#include <cstring>
#include <algorithm>
#include <iostream>

namespace ImXui {
namespace XurFormat {

XurPropertyValue::XurPropertyValue() : type(XurPropertyType_Integer), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f) {}
XurPropertyValue::XurPropertyValue(bool v) : type(XurPropertyType_Bool), boolVal(v), intVal(0), uintVal(0), floatVal(0.0f) {}
XurPropertyValue::XurPropertyValue(int32_t v) : type(XurPropertyType_Integer), boolVal(false), intVal(v), uintVal(0), floatVal(0.0f) {}
XurPropertyValue::XurPropertyValue(uint32_t v) : type(XurPropertyType_Unsigned), boolVal(false), intVal(0), uintVal(v), floatVal(0.0f) {}
XurPropertyValue::XurPropertyValue(float v) : type(XurPropertyType_Float), boolVal(false), intVal(0), uintVal(0), floatVal(v) {}
XurPropertyValue::XurPropertyValue(const std::wstring& v) : type(XurPropertyType_String), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), stringVal(v) {}
XurPropertyValue::XurPropertyValue(const XurVector& v) : type(XurPropertyType_Vector), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), vectorVal(v) {}
XurPropertyValue::XurPropertyValue(const XurQuaternion& v) : type(XurPropertyType_Quaternion), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), quatVal(v) {}
XurPropertyValue::XurPropertyValue(const XurColor& v) : type(XurPropertyType_Colour), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), colorVal(v) {}
XurPropertyValue::XurPropertyValue(const XurFigure& v) : type(XurPropertyType_Custom), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), figureVal(v) {}
XurPropertyValue::XurPropertyValue(const std::vector<XurProperty>& v) : type(XurPropertyType_Object), boolVal(false), intVal(0), uintVal(0), floatVal(0.0f), compoundVal(v) {}
XurPropertyValue::XurPropertyValue(const XurPropertyValue& other) : type(other.type), boolVal(other.boolVal), intVal(other.intVal), uintVal(other.uintVal), floatVal(other.floatVal), stringVal(other.stringVal), vectorVal(other.vectorVal), quatVal(other.quatVal), colorVal(other.colorVal), figureVal(other.figureVal), compoundVal(other.compoundVal) {}
XurPropertyValue& XurPropertyValue::operator=(const XurPropertyValue& other) {
    if (this != &other) {
        type = other.type; boolVal = other.boolVal; intVal = other.intVal; uintVal = other.uintVal; floatVal = other.floatVal;
        stringVal = other.stringVal; vectorVal = other.vectorVal; quatVal = other.quatVal; colorVal = other.colorVal;
        figureVal = other.figureVal; compoundVal = other.compoundVal;
    }
    return *this;
}
XurPropertyValue::~XurPropertyValue() {}


static const uint8_t* s_readEnd = nullptr;
static bool s_readTruncated = false;
static bool s_maskDesync = false;

static void resetReadState(const uint8_t* end) {
    s_readEnd = end;
    s_readTruncated = false;
    s_maskDesync = false;
}
static bool canRead(const uint8_t* p, size_t n) {
    if (p == NULL || s_readEnd == NULL || p + n > s_readEnd) {
        s_readTruncated = true;
        return false;
    }
    return true;
}

static uint16_t readU16(const uint8_t*& p) {
    if (!canRead(p, 2)) return 0;
    uint16_t v; memcpy(&v, p, 2); p += 2; return swap16(v);
}
static int16_t readS16(const uint8_t*& p) {
    return (int16_t)readU16(p);
}
static uint32_t readU32(const uint8_t*& p) {
    if (!canRead(p, 4)) return 0;
    uint32_t v; memcpy(&v, p, 4); p += 4; return swap32(v);
}
static int32_t readS32(const uint8_t*& p) {
    return (int32_t)readU32(p);
}
static float readF32(const uint8_t*& p) {
    if (!canRead(p, 4)) return 0.0f;
    uint32_t v; memcpy(&v, p, 4); p += 4; return swapFloat(v);
}
static uint8_t readU8(const uint8_t*& p) {
    if (!canRead(p, 1)) return 0;
    return *p++;
}

static std::wstring readUTF16String(const uint8_t*& p, int16_t length) {
    std::wstring result(length, L'\0');
    for (int16_t i = 0; i < length; i++) {
        if (!canRead(p, 2)) return result;
        uint16_t ch = readU16(p);
        result[i] = (wchar_t)ch;
    }
    return result;
}

static std::string ws2s(const std::wstring& ws) {
    return std::string(ws.begin(), ws.end());
}
static std::wstring s2ws(const std::string& s) {
    return std::wstring(s.begin(), s.end());
}

static bool parseStrn(const uint8_t* data, uint32_t length, std::vector<std::wstring>& strings) {
    strings.clear();
    strings.push_back(L"");
    const uint8_t* p = data;
    const uint8_t* end = data + length;
    while (p + 2 <= end) {
        int16_t strLen = readS16(p);
        if (strLen < 0 || p + strLen * 2 > end) break;
        strings.push_back(readUTF16String(p, strLen));
    }
    return strings.size() > 0;
}

static bool parseVect(const uint8_t* data, uint32_t length, std::vector<XurVector>& vectors) {
    const uint8_t* p = data;
    const uint8_t* end = data + length;
    while (p + 12 <= end) {
        XurVector v;
        v.x = readF32(p); v.y = readF32(p); v.z = readF32(p);
        vectors.push_back(v);
    }
    return true;
}

static bool parseQuat(const uint8_t* data, uint32_t length, std::vector<XurQuaternion>& quaternions) {
    const uint8_t* p = data;
    const uint8_t* end = data + length;
    while (p + 16 <= end) {
        XurQuaternion q;
        q.x = readF32(p); q.y = readF32(p); q.z = readF32(p); q.w = readF32(p);
        quaternions.push_back(q);
    }
    return true;
}

static bool parseCust(const uint8_t* data, uint32_t length, std::vector<XurFigure>& figures) {
    const uint8_t* p = data;
    const uint8_t* end = data + length;
    while (p < end) {
        int32_t dataLength = readS32(p);
        if (dataLength < 0 || (size_t)dataLength > (size_t)(end - p)) break;
        const uint8_t* figureEnd = p + dataLength;
        XurFigure fig;
        fig.boundingBox.x = readF32(p); fig.boundingBox.y = readF32(p);
        int32_t pointCount = readS32(p);
        for (int32_t i = 0; i < pointCount && p + 24 <= figureEnd; i++) {
            XurBezierPoint bp;
            bp.point.x = readF32(p); bp.point.y = readF32(p);
            bp.controlOne.x = readF32(p); bp.controlOne.y = readF32(p);
            bp.controlTwo.x = readF32(p); bp.controlTwo.y = readF32(p);
            fig.points.push_back(bp);
        }
        figures.push_back(fig);
        p = figureEnd;
    }
    return true;
}

static XurObject readObject(const uint8_t*& p, const uint8_t* end,
                             const std::vector<std::wstring>& strings,
                             const std::vector<XurVector>& vectors,
                             const std::vector<XurQuaternion>& quaternions,
                             const std::vector<XurFigure>& figures);

static XurPropertyValue readPropValue(const uint8_t*& p, XurPropertyType type,
                                       const std::vector<std::wstring>& strings,
                                       const std::vector<XurVector>& vectors,
                                       const std::vector<XurQuaternion>& quaternions,
                                       const std::vector<XurFigure>& figures) {
    switch (type) {
    case XurPropertyType_Bool:
        return XurPropertyValue((bool)(readU8(p) != 0));
    case XurPropertyType_Integer:
        return XurPropertyValue(readS32(p));
    case XurPropertyType_Unsigned:
        return XurPropertyValue(readU32(p));
    case XurPropertyType_Float:
        return XurPropertyValue(readF32(p));
    case XurPropertyType_Colour: {
        XurColor c;
        c.a = readU8(p); c.r = readU8(p); c.g = readU8(p); c.b = readU8(p);
        return XurPropertyValue(c);
    }
    case XurPropertyType_String: {
        int16_t idx = readS16(p);
        if (idx >= 0 && idx < (int16_t)strings.size())
            return XurPropertyValue(strings[idx]);
        return XurPropertyValue(std::wstring(L""));
    }
    case XurPropertyType_Vector: {
        int32_t idx = readS32(p);
        if (idx >= 0 && idx < (int32_t)vectors.size())
            return XurPropertyValue(vectors[idx]);
        XurVector zero = {0,0,0};
        return XurPropertyValue(zero);
    }
    case XurPropertyType_Quaternion: {
        int32_t idx = readS32(p);
        if (idx >= 0 && idx < (int32_t)quaternions.size())
            return XurPropertyValue(quaternions[idx]);
        XurQuaternion zero = {0,0,0,0};
        return XurPropertyValue(zero);
    }
    case XurPropertyType_Custom: {
        (void)readS32(p);
        if (!figures.empty())
            return XurPropertyValue(figures[0]);
        return XurPropertyValue(XurFigure());
    }
    case XurPropertyType_Object: {
        int16_t valuesCount = readS16(p); (void)valuesCount;
        uint8_t depth = readU8(p); (void)depth;
        uint8_t packedByte = readU8(p);
        if (packedByte == 0) return XurPropertyValue(std::vector<XurProperty>());
        int maskCount = packedByte & 0x07;
        std::vector<uint8_t> masks;
        for (int i = 0; i < maskCount; i++) masks.push_back(readU8(p));
        std::reverse(masks.begin(), masks.end());
        for (size_t mi = 0; mi < masks.size(); mi++) {
            for (int bit = 0; bit < 8; bit++) {
                if ((masks[mi] >> bit) & 1) {


                    if (!canRead(p, 4)) return XurPropertyValue(std::vector<XurProperty>());
                    p += 4;
                }
            }
        }
        return XurPropertyValue(std::vector<XurProperty>());
    }
    }
    return XurPropertyValue((int32_t)0);
}

static std::vector<XurProperty> readProperties(const uint8_t*& p, const uint8_t* end,
                                                 const std::string& className,
                                                 const std::vector<std::wstring>& strings,
                                                 const std::vector<XurVector>& vectors,
                                                 const std::vector<XurQuaternion>& quaternions,
                                                 const std::vector<XurFigure>& figures) {
    std::vector<XurProperty> props;
    if (p + 2 > end) return props;
    int16_t totalPropertyCount = readS16(p);

    auto hierarchy = XuiClassRegistry::getHierarchy(className);
    if (hierarchy.empty()) {
        return props;
    }

    std::reverse(hierarchy.begin(), hierarchy.end());
    for (size_t hi = 0; hi < hierarchy.size(); ++hi) {
        const XuiClassInfo* cls = hierarchy[hi];
        if (p >= end) break;

        uint8_t packedByte = readU8(p);
        if (packedByte == 0x00) continue;
        if (cls->props.empty()) continue;
        if ((int16_t)props.size() >= totalPropertyCount) continue;

        int maskCount = packedByte & 0x07;
        std::vector<uint8_t> masks;
        for (int i = 0; i < maskCount && p < end; i++)
            masks.push_back(readU8(p));
        std::reverse(masks.begin(), masks.end());

        int bitIdx = 0;
        for (size_t mi = 0; mi < masks.size(); ++mi) {
            uint8_t mask = masks[mi];
            for (int bit = 0; bit < 8; bit++, bitIdx++) {
                if (!((mask >> bit) & 1)) continue;
                if (bitIdx >= (int)cls->props.size()) {


                    s_maskDesync = true;
                    bitIdx = (int)cls->props.size();
                    break;
                }
                const auto& propDef = cls->props[bitIdx];
                XurProperty prop;
                prop.name = propDef.name;
                prop.type = propDef.type;
                if (propDef.indexed) {
                    uint8_t indexCount = readU8(p);
                    prop.value = readPropValue(p, propDef.type, strings, vectors, quaternions, figures);
                    for (uint8_t idx = 1; idx < indexCount; idx++) {
                        readPropValue(p, propDef.type, strings, vectors, quaternions, figures);
                    }
                } else {
                    prop.value = readPropValue(p, propDef.type, strings, vectors, quaternions, figures);
                }
                props.push_back(prop);
                if (s_readTruncated || s_maskDesync) break;
            }
            if (s_readTruncated || s_maskDesync) break;
        }
        if (s_readTruncated || s_maskDesync) break;
    }

    while ((int16_t)props.size() < totalPropertyCount && p < end && !s_readTruncated && !s_maskDesync) {
        XurProperty unknown;
        unknown.name = "unknown";
        unknown.type = XurPropertyType_Integer;
        unknown.value = XurPropertyValue((int32_t)0);
        props.push_back(unknown);
        if (!canRead(p, 4)) break;
        p += 4;
    }

    return props;
}

static XurNamedFrame readNamedFrame(const uint8_t*& p, const uint8_t* end,
                                      const std::vector<std::wstring>& strings) {
    XurNamedFrame nf;
    if (p + 9 > end) { p = end; return nf; }
    int16_t nameIdx = readS16(p);
    nf.name = (nameIdx >= 0 && nameIdx < (int16_t)strings.size()) ? strings[nameIdx] : L"";
    nf.keyframe = readS32(p);
    nf.command = (XurNamedFrameCommand)readU8(p);
    int16_t paramIdx = readS16(p);
    nf.targetParameter = (paramIdx >= 0 && paramIdx < (int16_t)strings.size()) ? strings[paramIdx] : L"";
    return nf;
}

static XurTimeline readTimeline(const uint8_t*& p, const uint8_t* end,
                                  const std::vector<std::wstring>& strings,
                                  const std::vector<XurVector>& vectors,
                                  const std::vector<XurQuaternion>& quaternions,
                                  const std::vector<XurFigure>& figures,
                                  const std::string& parentClassName) {
    XurTimeline tl;
    if (p + 6 > end) { p = end; return tl; }
    int16_t elemNameIdx = readS16(p);
    tl.elementName = (elemNameIdx >= 0 && elemNameIdx < (int16_t)strings.size()) ? strings[elemNameIdx] : L"";

    int32_t propDefCount = readS32(p);
    std::string targetClass = parentClassName;
    std::vector<XurPropertyType> animatedTypes;
    for (int32_t i = 0; i < propDefCount && p < end; i++) {
        uint8_t packedByte = readU8(p);
        uint8_t classIndex = readU8(p);
        uint8_t classDepth = packedByte & 0x7F;
        for (uint8_t d = 0; d < classDepth && p < end; d++) {
            readU8(p);
        }
        if (packedByte & 0x80) {
            if (p + 4 > end) break;
            readS32(p);
        }
        animatedTypes.push_back(XurPropertyType_Integer);
    }

    if (p + 4 > end) return tl;
    int32_t kfCount = readS32(p);
    for (int32_t k = 0; k < kfCount && p < end; k++) {
        XurKeyframe kf;
        if (p + 8 > end) break;
        kf.time = readS32(p);
        kf.interpolation = (XurInterpolationType)readU8(p);
        kf.easeIn = readU8(p);
        kf.easeOut = readU8(p);
        kf.easeScale = readU8(p);

        for (size_t ati = 0; ati < animatedTypes.size(); ++ati) {
            if (p >= end) break;
            XurProperty kfProp;
            kfProp.name = "anim";
            kfProp.type = animatedTypes[ati];
            kfProp.value = readPropValue(p, animatedTypes[ati], strings, vectors, quaternions, figures);
            kf.properties.push_back(kfProp);
        }
        tl.keyframes.push_back(kf);
    }
    return tl;
}

static int s_readDepth = 0;
static int s_totalObjects = 0;

static XurObject readObject(const uint8_t*& p, const uint8_t* end,
                             const std::vector<std::wstring>& strings,
                             const std::vector<XurVector>& vectors,
                             const std::vector<XurQuaternion>& quaternions,
                             const std::vector<XurFigure>& figures) {
    XurObject obj;
    if (p + 3 > end) { p = end; return obj; }

    if (++s_readDepth > 500) { p = end; return obj; }
    if (++s_totalObjects > 50000) { p = end; return obj; }

    int16_t classNameIndex = readS16(p);
    uint8_t flags = readU8(p);
    bool idxOk = (classNameIndex >= 0 && classNameIndex < (int16_t)strings.size());
    obj.className = ws2s(strings[idxOk ? classNameIndex : 0]);

    if (obj.className.empty()) { s_readDepth--; return obj; }

    if ((flags & 0x01) && p < end) {
        obj.properties = readProperties(p, end, obj.className, strings, vectors, quaternions, figures);
    }

    if ((flags & 0x02) && p + 4 <= end) {
        int32_t childCount = readS32(p);
        if (childCount > 0 && childCount < 10000) {
            for (int32_t i = 0; i < childCount && p < end; i++) {
                XurObject child = readObject(p, end, strings, vectors, quaternions, figures);
                if (child.className.empty()) break;
                obj.children.push_back(std::move(child));
            }
        }
    }

    if ((flags & 0x04) && p < end) {
        int32_t nfCount = readS32(p);
        if (nfCount > 0 && nfCount < 50000) {
            for (int32_t i = 0; i < nfCount && p < end; i++) {
                obj.namedFrames.push_back(readNamedFrame(p, end, strings));
            }
        }

        if (!obj.children.empty() && p + 4 <= end) {
            int32_t tlCount = readS32(p);
            if (tlCount > 0 && tlCount < 50000) {
                for (int32_t i = 0; i < tlCount && p < end; i++) {
                    obj.timelines.push_back(readTimeline(p, end, strings, vectors, quaternions, figures, obj.className));
                }
            }
        }
    }

    s_readDepth--;
    return obj;
}

static void resetObjectTracking() {
    s_readDepth = 0;
    s_totalObjects = 0;
}

std::unique_ptr<XurDocument> XurParser::parse(const void* data, size_t size) {
    if (!data || size < sizeof(XurHeader)) return nullptr;

    resetObjectTracking();
    XuiClassRegistry::initialize();

    std::unique_ptr<XurDocument> doc(new XurDocument());
    const uint8_t* bytes = static_cast<const uint8_t*>(data);
    doc->rawData.assign(bytes, bytes + size);

    memcpy(&doc->header, bytes, sizeof(XurHeader));
    if (!doc->header.isValid()) return nullptr;

    const uint8_t* p = bytes + sizeof(XurHeader);
    const uint8_t* fileEnd = bytes + size;

    if (doc->header.getFlags() & 0x01) {
        if (p + sizeof(XurCountHeader) > fileEnd) return nullptr;
        p += sizeof(XurCountHeader);
    }

    uint16_t sectionCount = doc->header.getSectionsCount();
    if (sectionCount > 20) return nullptr;
    std::vector<XurSectionEntry> entries;
    for (uint16_t i = 0; i < sectionCount; i++) {
        if (p + 12 > fileEnd) return nullptr;
        XurSectionEntry entry;
        memcpy(&entry, p, 12);
        entry.fromBE();
        entries.push_back(entry);
        p += 12;
    }

    for (size_t ei = 0; ei < entries.size(); ++ei) {
        const XurSectionEntry& entry = entries[ei];
        if (entry.offset + entry.length > size) continue;
        const uint8_t* sectionData = bytes + entry.offset;


        resetReadState(sectionData + entry.length);

        switch (entry.magic) {
        case MAGIC_STRN:
            parseStrn(sectionData, entry.length, doc->strings);
            break;
        case MAGIC_VECT:
            parseVect(sectionData, entry.length, doc->vectors);
            break;
        case MAGIC_QUAT:
            parseQuat(sectionData, entry.length, doc->quaternions);
            break;
        case MAGIC_CUST:
            parseCust(sectionData, entry.length, doc->figures);
            break;
        case MAGIC_DATA: {


            resetReadState(sectionData + entry.length);
            doc->rootObject = readObject(sectionData, sectionData + entry.length,
                                           doc->strings, doc->vectors,
                                           doc->quaternions, doc->figures);
            break;
        }
        }
    }

    return doc;
}

std::unique_ptr<XurDocument> XurParser::parseFromFile(const std::string& path) {
    std::ifstream ifs(path, std::ios::binary | std::ios::ate);
    if (!ifs) return nullptr;
    std::streamsize fileSize = ifs.tellg();
    ifs.seekg(0, std::ios::beg);
    std::vector<char> buf(fileSize);
    if (!ifs.read(buf.data(), fileSize)) return nullptr;
    return parse(buf.data(), buf.size());
}

bool XurParser::validate(const XurDocument* doc) {
    return doc && doc->header.isValid();
}

}
}