#pragma once

#include <cstddef>
#include <map>
#include <memory>
#include <optional>
#include <string>
#include <utility>
#include <variant>
#include <vector>

#include "weaseljson.h"

struct JsonNumber : std::string {};
using JsonValue = std::variant<std::nullptr_t, bool, std::string, JsonNumber,
                               std::unique_ptr<struct JsonArray>,
                               std::unique_ptr<struct JsonObject>>;
struct JsonArray : std::vector<JsonValue> {};
struct JsonObject : std::map<std::string, JsonValue> {};

struct ReadValueState {
  JsonValue result;
  std::vector<JsonValue> valueStack;
  std::vector<std::string> keyStack;
  std::vector</*bool*/ int> isKeyStack;
  void on_end_value() {
    auto object = std::move(valueStack.back());
    valueStack.pop_back();
    if (valueStack.empty()) {
      result = std::move(object);
      return;
    }
    auto i = valueStack.back().index();
    switch (i) {
    case 0: // null
    case 1: // bool
    case 2: // string
    case 3: // number
      __builtin_unreachable();
    case 4: // array
      std::get<std::unique_ptr<JsonArray>>(valueStack.back())
          ->push_back(std::move(object));
      return;
    case 5: // object
      if (std::exchange(isKeyStack.back(), !isKeyStack.back())) {
        keyStack.push_back(std::move(std::get<std::string>(object)));
      } else {
        std::get<std::unique_ptr<JsonObject>>(valueStack.back())
            ->emplace(std::move(keyStack.back()), std::move(object));
        keyStack.pop_back();
      }
      return;
    }
  }
};

inline WeaselJsonCallbacks readValueCallbacks() {
  WeaselJsonCallbacks result;
  result.on_begin_object = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(std::make_unique<JsonObject>());
    state->isKeyStack.push_back(true);
  };
  result.on_end_object = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->isKeyStack.pop_back();
    state->on_end_value();
  };
  result.on_begin_string = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(std::string());
  };
  result.on_string_data = +[](void *p, const char *buf, int len) {
    auto *state = (ReadValueState *)p;
    std::get<std::string>(state->valueStack.back()).append(buf, len);
  };
  result.on_end_string = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->on_end_value();
  };
  result.on_begin_array = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(std::make_unique<JsonArray>());
  };
  result.on_end_array = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->on_end_value();
  };
  result.on_begin_number = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(JsonNumber());
  };
  result.on_number_data = +[](void *p, const char *buf, int len) {
    auto *state = (ReadValueState *)p;
    std::get<JsonNumber>(state->valueStack.back()).append(buf, len);
  };
  result.on_end_number = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->on_end_value();
  };
  result.on_true_literal = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(true);
    state->on_end_value();
  };
  result.on_false_literal = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(false);
    state->on_end_value();
  };
  result.on_null_literal = +[](void *p) {
    auto *state = (ReadValueState *)p;
    state->valueStack.emplace_back(nullptr);
    state->on_end_value();
  };
  return result;
}

inline std::string escapeAsJsonString(std::string_view s) {
  std::string result;
  for (uint8_t c : s) {
    switch (c) {
    case '\"':
      result.append(R"(\")");
      break;
    case '\\':
      result.append(R"(\\)");
      break;
    case '\b':
      result.append(R"(\b)");
      break;
    case '\f':
      result.append(R"(\f)");
      break;
    case '\n':
      result.append(R"(\n)");
      break;
    case '\r':
      result.append(R"(\r)");
      break;
    case '\t':
      result.append(R"(\t)");
      break;
    default:
      if (c < 0x20) {
        const char *hex = "0123456789abcdef";
        result.append(R"(\u00)");
        result.push_back(hex[c >> 4]);
        result.push_back(hex[c & 15]);
      } else {
        // TODO check if valid utf-8
        result.push_back(c);
      }
    }
  }
  return result;
}

inline std::string toString(JsonValue const &jsonValue) {
  switch (jsonValue.index()) {
  case 0: // null
    return "null";
  case 1: // bool
    return std::get<bool>(jsonValue) ? "true" : "false";
  case 2: // string
    return "\"" + escapeAsJsonString(std::get<std::string>(jsonValue)) + "\"";
  case 3: // number
    return std::get<JsonNumber>(jsonValue);
  case 4: // array
  {
    std::string result = "[";
    std::string delimiter = "";
    for (auto const &v : *std::get<std::unique_ptr<JsonArray>>(jsonValue)) {
      result += delimiter + toString(v);
      delimiter = ", ";
    }
    return result + "]";
  }
  case 5: // object
  {
    std::string result = "{";
    std::string delimiter = "";
    for (auto const &[k, v] :
         *std::get<std::unique_ptr<JsonObject>>(jsonValue)) {
      result += delimiter + "\"" + escapeAsJsonString(k) + "\": " + toString(v);
      delimiter = ", ";
    }
    return result + "}";
  }
  }
  __builtin_unreachable();
}

inline std::optional<JsonValue> toValue(std::string copy, int stride) {
  ReadValueState state;
  auto c = readValueCallbacks();
  std::unique_ptr<WeaselJsonParser, decltype(&WeaselJsonParser_destroy)> parser{
      WeaselJsonParser_create(1024, &c, &state), WeaselJsonParser_destroy};
  if (stride == 0) {
    if (WeaselJsonParser_parse(parser.get(), copy.data(), copy.size()) !=
        WeaselJson_AGAIN) {
      return std::nullopt;
    }
  } else {
    for (int i = 0; i < copy.size(); i += stride) {
      if (WeaselJsonParser_parse(parser.get(), copy.data() + i,
                                 std::min<int>(stride, copy.size() - i)) !=
          WeaselJson_AGAIN) {
        return std::nullopt;
      }
    }
  }
  if (WeaselJsonParser_parse(parser.get(), nullptr, 0) != WeaselJson_OK) {
    return std::nullopt;
  }
  return std::move(state.result);
}