#include "json.hh"

#include <cassert>
#include <charconv>
#include <cstddef>
#include <cstdint>
#include <memory>
#include <ostream>
#include <string>
#include <string_view>
#include <system_error>
#include <vector>

namespace json {

namespace {

struct StackEntry {
  enum class Type : uint8_t {
    kRoot,
    kObject,
    kArray,
  };

  Type type;
  bool need_comma{false};

  explicit StackEntry(Type type) : type(type) {}
};

class BaseWriter : public Writer {
 public:
  BaseWriter() : stack_({StackEntry(StackEntry::Type::kRoot)}) {}

  void value(std::string_view value) override {
    before_value();
    quote(value);
  }

  void value(int64_t value) override {
    before_value();
    write(value);
  }

  void value(uint64_t value) override {
    before_value();
    write(value);
  }

  void value(float value) override {
    before_value();
    write(value);
  }

  void value(double value) override {
    before_value();
    write(value);
  }

  void value(bool value) override {
    before_value();
    write(value);
  }

  void start_array() override {
    before_value();
    stack_.emplace_back(StackEntry::Type::kArray);
    write('[');
  }

  void end_array() override {
    if (stack_.empty()) {
      assert(false);
      return;
    }
    assert(stack_.back().type == StackEntry::Type::kArray);
    stack_.pop_back();
    write(']');
    if (!stack_.empty())
      stack_.back().need_comma = true;
  }

  void start_object() override {
    before_value();
    stack_.emplace_back(StackEntry::Type::kObject);
    write('{');
  }

  void key(std::string_view name) override {
    if (stack_.empty()) {
      assert(false);
      return;
    }
    assert(stack_.back().type == StackEntry::Type::kObject);
    if (stack_.back().need_comma) {
      write(',');
      stack_.back().need_comma = false;
    }
    quote(name);
    write(':');
  }

  void end_object() override {
    if (stack_.empty()) {
      assert(false);
      return;
    }
    assert(stack_.back().type == StackEntry::Type::kObject);
    stack_.pop_back();
    write('}');
    if (!stack_.empty())
      stack_.back().need_comma = true;
  }

  void clear() override {
    stack_.clear();
    stack_.emplace_back(StackEntry::Type::kRoot);
  }

 protected:
  virtual void write(int64_t value) = 0;
  virtual void write(uint64_t value) = 0;
  virtual void write(float value) = 0;
  virtual void write(double value) = 0;
  virtual void write(bool value) = 0;
  virtual void write(char value) = 0;
  virtual void write(std::string_view value) = 0;

  void write(char const* value) { this->write(std::string_view(value)); }

 private:
  void before_value() {
    if (stack_.empty()) {
      assert(false);
      return;
    }
    if (stack_.back().need_comma) {
      assert(stack_.back().type != StackEntry::Type::kRoot);
      write(',');
    } else {
      stack_.back().need_comma = true;
    }
  }

  void quote(std::string_view str) {
    write('"');
    size_t last = 0;
    while (true) {
      auto pos = need_quote(str, last);
      if (pos == std::string_view::npos) {
        write(str.substr(last));
        break;
      }
      write(str.substr(last, pos - last));
      switch (str[pos]) {
        case '"':
          write("\\\"");
          break;
        case '\\':
          write("\\\\");
          break;
        case '\n':
          write("\\n");
          break;
        case '\r':
          write("\\r");
          break;
        case '\t':
          write("\\t");
          break;
        case '\b':
          write("\\b");
          break;
        case '\f':
          write("\\f");
          break;
        default: {
          char tmp[4];
          write("\\u");
          auto [ptr, ec] = std::to_chars(tmp, tmp + sizeof(tmp), str[pos], 16);
          if (ec == std::errc()) {
            size_t len = ptr - tmp;
            assert(len > 0);
            for (size_t i = 4; i > len; --i)
              write('0');
            write(std::string_view(tmp, len));
          } else {
            assert(false);
          }
          break;
        };
      }
      last = pos + 1;
    }
    write('"');
  }

  static inline size_t need_quote(std::string_view str, size_t offset) {
    for (; offset < str.size(); ++offset) {
      if (str[offset] == '\\' || str[offset] == '"' ||
          (str[offset] >= 0 && str[offset] < ' '))
        return offset;
    }
    return std::string_view::npos;
  }

  std::vector<StackEntry> stack_;
};

class IosWriter : public BaseWriter {
 public:
  explicit IosWriter(std::ostream& out) : out_(out) {}

  void write(std::string_view value) override { out_ << value; }

  void write(int64_t value) override { out_ << value; }

  void write(uint64_t value) override { out_ << value; }

  void write(float value) override { out_ << value; }

  void write(double value) override { out_ << value; }

  void write(bool value) override { out_ << value; }

  void write(char value) override { out_ << value; }

 private:
  std::ostream& out_;
};

class StringWriter : public BaseWriter {
 public:
  explicit StringWriter(std::string& out) : out_(out) {}

  void write(std::string_view value) override { out_.append(value); }

  void write(int64_t value) override {
    auto [ptr, ec] = std::to_chars(tmp_, tmp_ + sizeof(tmp_), value);
    if (ec == std::errc()) {
      out_.append(tmp_, ptr - tmp_);
    } else {
      assert(false);
    }
  }

  void write(uint64_t value) override {
    auto [ptr, ec] = std::to_chars(tmp_, tmp_ + sizeof(tmp_), value);
    if (ec == std::errc()) {
      out_.append(tmp_, ptr - tmp_);
    } else {
      assert(false);
    }
  }

  void write(float value) override {
    auto [ptr, ec] = std::to_chars(tmp_, tmp_ + sizeof(tmp_), value);
    if (ec == std::errc()) {
      out_.append(tmp_, ptr - tmp_);
    } else {
      assert(false);
    }
  }

  void write(double value) override {
    auto [ptr, ec] = std::to_chars(tmp_, tmp_ + sizeof(tmp_), value);
    if (ec == std::errc()) {
      out_.append(tmp_, ptr - tmp_);
    } else {
      assert(false);
    }
  }

  void write(bool value) override { out_.append(value ? "true" : "false"); }

  void write(char value) override { out_.push_back(value); }

  void clear() override {
    BaseWriter::clear();
    out_.clear();
  }

 private:
  std::string& out_;
  char tmp_[100];
};

}  // namespace

void Writer::value(char const* value) { this->value(std::string_view(value)); }

void Writer::value(int value) {
  static_assert(sizeof(int) <= sizeof(int64_t));
  this->value(static_cast<int64_t>(value));
}

std::unique_ptr<Writer> writer(std::string& out) {
  return std::make_unique<StringWriter>(out);
}

std::unique_ptr<Writer> writer(std::ostream& out) {
  return std::make_unique<IosWriter>(out);
}

}  // namespace json