#include "rdp_worker/common/json.hpp"

#include <cctype>
#include <cstdint>
#include <sstream>
#include <stdexcept>

namespace rdp_worker::common {

JsonValue::JsonValue() : value(nullptr) {}
JsonValue::JsonValue(std::nullptr_t) : value(nullptr) {}
JsonValue::JsonValue(bool input) : value(input) {}
JsonValue::JsonValue(double input) : value(input) {}
JsonValue::JsonValue(int input) : value(static_cast<double>(input)) {}
JsonValue::JsonValue(const char* input) : value(std::string(input)) {}
JsonValue::JsonValue(std::string input) : value(std::move(input)) {}
JsonValue::JsonValue(JsonArray input) : value(std::move(input)) {}
JsonValue::JsonValue(JsonObject input) : value(std::move(input)) {}

bool JsonValue::IsObject() const { return std::holds_alternative<JsonObject>(value); }
bool JsonValue::IsArray() const { return std::holds_alternative<JsonArray>(value); }
bool JsonValue::IsString() const { return std::holds_alternative<std::string>(value); }
bool JsonValue::IsBool() const { return std::holds_alternative<bool>(value); }
bool JsonValue::IsNumber() const { return std::holds_alternative<double>(value); }
const JsonObject& JsonValue::AsObject() const { return std::get<JsonObject>(value); }
const JsonArray& JsonValue::AsArray() const { return std::get<JsonArray>(value); }
const std::string& JsonValue::AsString() const { return std::get<std::string>(value); }
bool JsonValue::AsBool() const { return std::get<bool>(value); }
double JsonValue::AsNumber() const { return std::get<double>(value); }

namespace {

void AppendUtf8(std::string& output, std::uint32_t codepoint) {
    if (codepoint <= 0x7F) {
        output.push_back(static_cast<char>(codepoint));
    } else if (codepoint <= 0x7FF) {
        output.push_back(static_cast<char>(0xC0 | (codepoint >> 6)));
        output.push_back(static_cast<char>(0x80 | (codepoint & 0x3F)));
    } else if (codepoint <= 0xFFFF) {
        output.push_back(static_cast<char>(0xE0 | (codepoint >> 12)));
        output.push_back(static_cast<char>(0x80 | ((codepoint >> 6) & 0x3F)));
        output.push_back(static_cast<char>(0x80 | (codepoint & 0x3F)));
    } else if (codepoint <= 0x10FFFF) {
        output.push_back(static_cast<char>(0xF0 | (codepoint >> 18)));
        output.push_back(static_cast<char>(0x80 | ((codepoint >> 12) & 0x3F)));
        output.push_back(static_cast<char>(0x80 | ((codepoint >> 6) & 0x3F)));
        output.push_back(static_cast<char>(0x80 | (codepoint & 0x3F)));
    } else {
        throw std::runtime_error("invalid JSON unicode escape");
    }
}

class Parser {
public:
    explicit Parser(const std::string& input) : input_(input), index_(0) {}

    JsonValue Parse() {
        SkipWhitespace();
        JsonValue value = ParseValue();
        SkipWhitespace();
        if (index_ != input_.size()) {
            throw std::runtime_error("unexpected trailing JSON data");
        }
        return value;
    }

private:
    JsonValue ParseValue() {
        SkipWhitespace();
        if (Match("null")) {
            return JsonValue(nullptr);
        }
        if (Match("true")) {
            return JsonValue(true);
        }
        if (Match("false")) {
            return JsonValue(false);
        }
        if (Peek() == '"') {
            return JsonValue(ParseString());
        }
        if (Peek() == '{') {
            return JsonValue(ParseObject());
        }
        if (Peek() == '[') {
            return JsonValue(ParseArray());
        }
        if (Peek() == '-' || std::isdigit(static_cast<unsigned char>(Peek()))) {
            return JsonValue(ParseNumber());
        }
        throw std::runtime_error("unexpected JSON token");
    }

    JsonObject ParseObject() {
        Expect('{');
        JsonObject object;
        SkipWhitespace();
        if (Peek() == '}') {
            Advance();
            return object;
        }
        while (true) {
            const std::string key = ParseString();
            SkipWhitespace();
            Expect(':');
            object.emplace(key, ParseValue());
            SkipWhitespace();
            if (Peek() == '}') {
                Advance();
                break;
            }
            Expect(',');
        }
        return object;
    }

    JsonArray ParseArray() {
        Expect('[');
        JsonArray array;
        SkipWhitespace();
        if (Peek() == ']') {
            Advance();
            return array;
        }
        while (true) {
            array.emplace_back(ParseValue());
            SkipWhitespace();
            if (Peek() == ']') {
                Advance();
                break;
            }
            Expect(',');
        }
        return array;
    }

    std::string ParseString() {
        Expect('"');
        std::string output;
        while (index_ < input_.size()) {
            const char current = input_[index_++];
            if (current == '"') {
                return output;
            }
            if (current == '\\') {
                if (index_ >= input_.size()) {
                    throw std::runtime_error("invalid JSON escape");
                }
                const char escaped = input_[index_++];
                switch (escaped) {
                    case '"':
                    case '\\':
                    case '/':
                        output.push_back(escaped);
                        break;
                    case 'b':
                        output.push_back('\b');
                        break;
                    case 'f':
                        output.push_back('\f');
                        break;
                    case 'n':
                        output.push_back('\n');
                        break;
                    case 'r':
                        output.push_back('\r');
                        break;
                    case 't':
                        output.push_back('\t');
                        break;
                    case 'u':
                        AppendUtf8(output, ParseUnicodeEscape());
                        break;
                    default:
                        throw std::runtime_error("unsupported JSON escape");
                }
                continue;
            }
            output.push_back(current);
        }
        throw std::runtime_error("unterminated JSON string");
    }

    std::uint32_t ParseUnicodeEscape() {
        const std::uint32_t first = ParseHexQuad();
        if (first >= 0xD800 && first <= 0xDBFF) {
            if (index_ + 1 >= input_.size() || input_[index_] != '\\' || input_[index_ + 1] != 'u') {
                throw std::runtime_error("invalid JSON unicode surrogate");
            }
            index_ += 2;
            const std::uint32_t second = ParseHexQuad();
            if (second < 0xDC00 || second > 0xDFFF) {
                throw std::runtime_error("invalid JSON unicode surrogate");
            }
            return 0x10000 + (((first - 0xD800) << 10) | (second - 0xDC00));
        }
        if (first >= 0xDC00 && first <= 0xDFFF) {
            throw std::runtime_error("invalid JSON unicode surrogate");
        }
        return first;
    }

    std::uint32_t ParseHexQuad() {
        if (index_ + 4 > input_.size()) {
            throw std::runtime_error("invalid JSON unicode escape");
        }
        std::uint32_t value = 0;
        for (int i = 0; i < 4; ++i) {
            const char ch = input_[index_++];
            value <<= 4;
            if (ch >= '0' && ch <= '9') {
                value |= static_cast<std::uint32_t>(ch - '0');
            } else if (ch >= 'a' && ch <= 'f') {
                value |= static_cast<std::uint32_t>(ch - 'a' + 10);
            } else if (ch >= 'A' && ch <= 'F') {
                value |= static_cast<std::uint32_t>(ch - 'A' + 10);
            } else {
                throw std::runtime_error("invalid JSON unicode escape");
            }
        }
        return value;
    }

    double ParseNumber() {
        const std::size_t start = index_;
        if (Peek() == '-') {
            Advance();
        }
        while (std::isdigit(static_cast<unsigned char>(Peek()))) {
            Advance();
        }
        if (Peek() == '.') {
            Advance();
            while (std::isdigit(static_cast<unsigned char>(Peek()))) {
                Advance();
            }
        }
        return std::stod(input_.substr(start, index_ - start));
    }

    void SkipWhitespace() {
        while (index_ < input_.size() && std::isspace(static_cast<unsigned char>(input_[index_]))) {
            ++index_;
        }
    }

    char Peek() const {
        if (index_ >= input_.size()) {
            return '\0';
        }
        return input_[index_];
    }

    void Advance() {
        if (index_ < input_.size()) {
            ++index_;
        }
    }

    void Expect(char expected) {
        SkipWhitespace();
        if (Peek() != expected) {
            throw std::runtime_error("unexpected JSON character");
        }
        Advance();
    }

    bool Match(const char* keyword) {
        const std::size_t length = std::char_traits<char>::length(keyword);
        if (input_.substr(index_, length) == keyword) {
            index_ += length;
            return true;
        }
        return false;
    }

    const std::string& input_;
    std::size_t index_;
};

std::string Escape(const std::string& value) {
    std::ostringstream output;
    for (const char ch : value) {
        switch (ch) {
            case '"':
                output << "\\\"";
                break;
            case '\\':
                output << "\\\\";
                break;
            case '\n':
                output << "\\n";
                break;
            case '\r':
                output << "\\r";
                break;
            case '\t':
                output << "\\t";
                break;
            default:
                output << ch;
                break;
        }
    }
    return output.str();
}

std::string SerializeInternal(const JsonValue& value) {
    if (std::holds_alternative<std::nullptr_t>(value.value)) {
        return "null";
    }
    if (std::holds_alternative<bool>(value.value)) {
        return std::get<bool>(value.value) ? "true" : "false";
    }
    if (std::holds_alternative<double>(value.value)) {
        std::ostringstream output;
        output << std::get<double>(value.value);
        return output.str();
    }
    if (std::holds_alternative<std::string>(value.value)) {
        return "\"" + Escape(std::get<std::string>(value.value)) + "\"";
    }
    if (std::holds_alternative<JsonArray>(value.value)) {
        std::ostringstream output;
        output << "[";
        const auto& array = std::get<JsonArray>(value.value);
        for (std::size_t i = 0; i < array.size(); ++i) {
            if (i > 0) {
                output << ",";
            }
            output << SerializeInternal(array[i]);
        }
        output << "]";
        return output.str();
    }
    std::ostringstream output;
    output << "{";
    const auto& object = std::get<JsonObject>(value.value);
    bool first = true;
    for (const auto& [key, child] : object) {
        if (!first) {
            output << ",";
        }
        first = false;
        output << "\"" << Escape(key) << "\":" << SerializeInternal(child);
    }
    output << "}";
    return output.str();
}

}  // namespace

JsonValue ParseJson(const std::string& input) {
    return Parser(input).Parse();
}

std::string SerializeJson(const JsonValue& value) {
    return SerializeInternal(value);
}

std::optional<std::string> GetString(const JsonObject& object, const std::string& key) {
    auto iterator = object.find(key);
    if (iterator == object.end() || !iterator->second.IsString()) {
        return std::nullopt;
    }
    return iterator->second.AsString();
}

std::optional<bool> GetBool(const JsonObject& object, const std::string& key) {
    auto iterator = object.find(key);
    if (iterator == object.end() || !iterator->second.IsBool()) {
        return std::nullopt;
    }
    return iterator->second.AsBool();
}

std::optional<double> GetNumber(const JsonObject& object, const std::string& key) {
    auto iterator = object.find(key);
    if (iterator == object.end() || !iterator->second.IsNumber()) {
        return std::nullopt;
    }
    return iterator->second.AsNumber();
}

const JsonObject* GetObject(const JsonObject& object, const std::string& key) {
    auto iterator = object.find(key);
    if (iterator == object.end() || !iterator->second.IsObject()) {
        return nullptr;
    }
    return &iterator->second.AsObject();
}

const JsonArray* GetArray(const JsonObject& object, const std::string& key) {
    auto iterator = object.find(key);
    if (iterator == object.end() || !iterator->second.IsArray()) {
        return nullptr;
    }
    return &iterator->second.AsArray();
}

}  // namespace rdp_worker::common