#include "JsonParser.h"
#include <Core/Memory/MemoryArena.h>
#include <Core/Common/CoreUtils.h>
#include <ctype.h>
#include <stdlib.h>
#include <string.h>

namespace Romeo
{
    struct JsonParserState
    {
        Juliet::Arena* Arena;
        const char* Ptr;
        const char* End;
    };

    static void SkipWhitespace(JsonParserState& state)
    {
        while (state.Ptr < state.End)
        {
            char c = *state.Ptr;
            if (c == ' ' || c == '\t' || c == '\n' || c == '\r')
            {
                state.Ptr++;
            }
            else
            {
                break;
            }
        }
    }

    static bool ParseString(JsonParserState& state, Juliet::String& outStr)
    {
        if (state.Ptr >= state.End || *state.Ptr != '"')
        {
            return false;
        }
        state.Ptr++;
        const char* start = state.Ptr;
        
        size_t len = 0;
        bool hasEscape = false;
        while (state.Ptr < state.End)
        {
            char c = *state.Ptr;
            if (c == '"')
            {
                break;
            }
            else if (c == '\\')
            {
                hasEscape = true;
                state.Ptr++;
                if (state.Ptr >= state.End)
                {
                    return false;
                }
            }
            state.Ptr++;
            len++;
        }
        
        if (state.Ptr >= state.End)
        {
            return false;
        }
        
        const char* endQuote = state.Ptr;
        state.Ptr++;
        
        char* buf = Juliet::ArenaPushArray<char>(state.Arena, len + 1 JULIET_DEBUG_PARAM("JsonParsedStr"));
        
        if (hasEscape)
        {
            size_t writeIdx = 0;
            const char* readPtr = start;
            while (readPtr < endQuote)
            {
                char c = *readPtr;
                if (c == '\\')
                {
                    readPtr++;
                    if (readPtr >= endQuote)
                    {
                        break;
                    }
                    char esc = *readPtr;
                    if (esc == 'n') { buf[writeIdx++] = '\n'; }
                    else if (esc == 'r') { buf[writeIdx++] = '\r'; }
                    else if (esc == 't') { buf[writeIdx++] = '\t'; }
                    else if (esc == 'b') { buf[writeIdx++] = '\b'; }
                    else if (esc == 'f') { buf[writeIdx++] = '\f'; }
                    else if (esc == '"') { buf[writeIdx++] = '"'; }
                    else if (esc == '\\') { buf[writeIdx++] = '\\'; }
                    else if (esc == '/') { buf[writeIdx++] = '/'; }
                    else if (esc == 'u')
                    {
                        if (readPtr + 4 < endQuote)
                        {
                            buf[writeIdx++] = '?';
                            readPtr += 4;
                        }
                    }
                    else
                    {
                        buf[writeIdx++] = esc;
                    }
                }
                else
                {
                    buf[writeIdx++] = c;
                }
                readPtr++;
            }
            buf[writeIdx] = '\0';
            outStr.Data = buf;
            outStr.Size = writeIdx;
        }
        else
        {
            memcpy(buf, start, len);
            buf[len] = '\0';
            outStr.Data = buf;
            outStr.Size = len;
        }
        
        return true;
    }

    static bool ParseValue(JsonParserState& state, JsonValue& outVal);

    static bool ParseObject(JsonParserState& state, JsonValue& outVal)
    {
        if (state.Ptr >= state.End || *state.Ptr != '{')
        {
            return false;
        }
        state.Ptr++;
        outVal.Type = JsonType::Object;
        outVal.ObjectVal.Create(state.Arena JULIET_DEBUG_ONLY(, "JsonObject"));

        while (state.Ptr < state.End)
        {
            SkipWhitespace(state);
            if (state.Ptr < state.End && *state.Ptr == '}')
            {
                state.Ptr++;
                return true;
            }

            Juliet::String key = {};
            if (!ParseString(state, key))
            {
                return false;
            }

            SkipWhitespace(state);
            if (state.Ptr >= state.End || *state.Ptr != ':')
            {
                return false;
            }
            state.Ptr++;

            JsonKeyValue kv = {};
            kv.Key = key;
            kv.Value = Juliet::ArenaPushArray<JsonValue>(state.Arena, 1 JULIET_DEBUG_PARAM("JsonValueRef"));
            new (kv.Value) JsonValue();

            if (!ParseValue(state, *kv.Value))
            {
                return false;
            }

            outVal.ObjectVal.PushBack(kv);

            SkipWhitespace(state);
            if (state.Ptr < state.End && *state.Ptr == ',')
            {
                state.Ptr++;
            }
            else if (state.Ptr < state.End && *state.Ptr == '}')
            {
                state.Ptr++;
                return true;
            }
            else
            {
                return false;
            }
        }
        return false;
    }

    static bool ParseArray(JsonParserState& state, JsonValue& outVal)
    {
        if (state.Ptr >= state.End || *state.Ptr != '[')
        {
            return false;
        }
        state.Ptr++;
        outVal.Type = JsonType::Array;
        outVal.ArrayVal.Create(state.Arena JULIET_DEBUG_ONLY(, "JsonArray"));

        while (state.Ptr < state.End)
        {
            SkipWhitespace(state);
            if (state.Ptr < state.End && *state.Ptr == ']')
            {
                state.Ptr++;
                return true;
            }

            JsonValue val = {};
            if (!ParseValue(state, val))
            {
                return false;
            }

            outVal.ArrayVal.PushBack(val);

            SkipWhitespace(state);
            if (state.Ptr < state.End && *state.Ptr == ',')
            {
                state.Ptr++;
            }
            else if (state.Ptr < state.End && *state.Ptr == ']')
            {
                state.Ptr++;
                return true;
            }
            else
            {
                return false;
            }
        }
        return false;
    }

    static bool ParseNumber(JsonParserState& state, JsonValue& outVal)
    {
        const char* start = state.Ptr;
        while (state.Ptr < state.End)
        {
            char c = *state.Ptr;
            if (isdigit(c) || c == '-' || c == '+' || c == '.' || c == 'e' || c == 'E')
            {
                state.Ptr++;
            }
            else
            {
                break;
            }
        }
        size_t len = static_cast<size_t>(state.Ptr - start);
        if (len == 0)
        {
            return false;
        }

        char temp[64];
        if (len >= sizeof(temp))
        {
            return false;
        }
        memcpy(temp, start, len);
        temp[len] = '\0';

        char* endPtr = nullptr;
        outVal.NumVal = strtod(temp, &endPtr);
        if (endPtr == temp)
        {
            return false;
        }
        outVal.Type = JsonType::Number;
        return true;
    }

    static bool ParseValue(JsonParserState& state, JsonValue& outVal)
    {
        SkipWhitespace(state);
        if (state.Ptr >= state.End)
        {
            return false;
        }

        char c = *state.Ptr;
        if (c == '{')
        {
            return ParseObject(state, outVal);
        }
        else if (c == '[')
        {
            return ParseArray(state, outVal);
        }
        else if (c == '"')
        {
            outVal.Type = JsonType::String;
            return ParseString(state, outVal.StrVal);
        }
        else if (c == 't' || c == 'f')
        {
            size_t remaining = static_cast<size_t>(state.End - state.Ptr);
            if (c == 't' && remaining >= 4 && strncmp(state.Ptr, "true", 4) == 0)
            {
                outVal.Type = JsonType::Bool;
                outVal.BoolVal = true;
                state.Ptr += 4;
                return true;
            }
            else if (c == 'f' && remaining >= 5 && strncmp(state.Ptr, "false", 5) == 0)
            {
                outVal.Type = JsonType::Bool;
                outVal.BoolVal = false;
                state.Ptr += 5;
                return true;
            }
        }
        else if (c == 'n')
        {
            size_t remaining = static_cast<size_t>(state.End - state.Ptr);
            if (remaining >= 4 && strncmp(state.Ptr, "null", 4) == 0)
            {
                outVal.Type = JsonType::Null;
                state.Ptr += 4;
                return true;
            }
        }
        else if (isdigit(c) || c == '-')
        {
            return ParseNumber(state, outVal);
        }

        return false;
    }

    bool Json_Parse(Juliet::Arena* arena, const Juliet::String& jsonStr, JsonValue& outValue)
    {
        Assert(arena != nullptr);
        if (jsonStr.Size == 0)
        {
            return false;
        }

        JsonParserState state;
        state.Arena = arena;
        state.Ptr = CStr(jsonStr);
        state.End = state.Ptr + jsonStr.Size;

        return ParseValue(state, outValue);
    }

    bool Json_Query(const JsonValue& root, const char* path, Juliet::String& outResult)
    {
        Assert(path != nullptr);
        
        const JsonValue* current = &root;
        const char* p = path;
        
        while (*p)
        {
            const char* start = p;
            while (*p && *p != '/')
            {
                p++;
            }
            size_t len = static_cast<size_t>(p - start);
            if (len == 0)
            {
                if (*p == '/') { p++; }
                continue;
            }

            if (current->Type == JsonType::Object)
            {
                const JsonValue* next = nullptr;
                for (const JsonKeyValue& kv : current->ObjectVal)
                {
                    if (kv.Key.Size == len && strncmp(CStr(kv.Key), start, len) == 0)
                    {
                        next = kv.Value;
                        break;
                    }
                }
                if (!next)
                {
                    return false;
                }
                current = next;
            }
            else if (current->Type == JsonType::Array)
            {
                char tempIndex[32];
                if (len >= sizeof(tempIndex))
                {
                    return false;
                }
                memcpy(tempIndex, start, len);
                tempIndex[len] = '\0';
                
                int index = atoi(tempIndex);
                if (index < 0 || static_cast<size_t>(index) >= current->ArrayVal.Size())
                {
                    return false;
                }
                current = &current->ArrayVal[static_cast<size_t>(index)];
            }
            else
            {
                return false;
            }

            if (*p == '/')
            {
                p++;
            }
        }

        if (current->Type == JsonType::String)
        {
            outResult = current->StrVal;
            return true;
        }

        return false;
    }

    Juliet::String Json_EscapeString(Juliet::Arena* arena, const Juliet::String& str)
    {
        Assert(arena != nullptr);
        
        size_t escLen = 0;
        const unsigned char* p = reinterpret_cast<const unsigned char*>(CStr(str));
        const unsigned char* end = p + str.Size;
        while (p < end)
        {
            unsigned char c = *p;
            if (c == '"' || c == '\\' || c == '\n' || c == '\r' || c == '\t' || c == '\b' || c == '\f')
            {
                escLen += 2;
            }
            else if (c < 32)
            {
                escLen += 6;
            }
            else
            {
                escLen += 1;
            }
            p++;
        }
        
        char* buf = Juliet::ArenaPushArray<char>(arena, escLen + 1 JULIET_DEBUG_PARAM("JsonEscapedStr"));
        char* dst = buf;
        p = reinterpret_cast<const unsigned char*>(CStr(str));
        while (p < end)
        {
            unsigned char c = *p;
            if (c == '"') { *dst++ = '\\'; *dst++ = '"'; }
            else if (c == '\\') { *dst++ = '\\'; *dst++ = '\\'; }
            else if (c == '\n') { *dst++ = '\\'; *dst++ = 'n'; }
            else if (c == '\r') { *dst++ = '\\'; *dst++ = 'r'; }
            else if (c == '\t') { *dst++ = '\\'; *dst++ = 't'; }
            else if (c == '\b') { *dst++ = '\\'; *dst++ = 'b'; }
            else if (c == '\f') { *dst++ = '\\'; *dst++ = 'f'; }
            else if (c < 32)
            {
                snprintf(dst, 7, "\\u00%02x", static_cast<unsigned int>(c));
                dst += 6;
            }
            else { *dst++ = static_cast<char>(c); }
            p++;
        }
        *dst = '\0';
        return { buf, escLen };
    }
}