#pragma once

#include <istream>
#include <vector>

struct Expr {
public:
  enum class Type { Atom, String, List };

  Type type = Type::Atom;
  std::string value;
  std::vector<Expr> children;

  [[nodiscard]] Expr get(const std::string &key) const {
    if (type != Type::List) {
      throw std::runtime_error("get() called on non-list");
    }
    for (const auto &child : children) {
      if (child.type == Type::List && child.children.size() >= 2 &&
          child.children[0].type == Type::Atom &&
          child.children[0].value == key) {
        Expr result = {.type = Type::List};
        result.children.assign(child.children.begin() + 1,
                               child.children.end());
        return result;
      }
    }
    throw std::runtime_error("key not found: " + key);
  }

  static Expr parse(std::istream &in) { return parse_expr(in); }

private:
  static std::string read_atom(std::istream &in) {
    std::string atom;
    while (true) {
      int next = in.peek();
      if (next == -1 || is_space(static_cast<char>(next)) || next == '(' ||
          next == ')') {
        break;
      }
      atom += static_cast<char>(in.get());
    }
    if (atom.empty()) {
      throw std::runtime_error("empty atom");
    }
    return atom;
  }

  static Expr parse_list(std::istream &in) {
    Expr result = {.type = Expr::Type::List};

    in.get(); // consume bracket

    while (true) {
      while (is_space(static_cast<char>(in.peek()))) {
        in.get();
      }
      if (!in) {
        throw std::runtime_error("unexpected EOF");
      }

      if (in.peek() == ')') {
        in.get();
        break;
      }

      result.children.push_back(parse_expr(in));
    }

    return result;
  }

  static std::string read_string(std::istream &in) {
    std::string result;

    in.get(); // consume the quote

    while (true) {
      if (!in) {
        throw std::runtime_error("unexpected EOF in string");
      }

      auto c = static_cast<char>(in.get());
      if (c == '"') {
        break;
      }
      result += c;
    }
    return result;
  }

  static Expr parse_expr(std::istream &in) {
    while (is_space(static_cast<char>(in.peek()))) {
      in.get();
    }
    if (!in) {
      throw std::runtime_error("unexpected EOF");
    }

    if (in.peek() == '(') {
      return parse_list(in);
    }
    if (in.peek() == ')') {
      throw std::runtime_error("unexpected ')'");
    }

    if (in.peek() == '"') {
      return {.type = Type::String, .value = read_string(in)};
    }

    return {.type = Type::Atom, .value = read_atom(in)};
  }

  static bool is_space(char c) {
    return c == ' ' || c == '\t' || c == '\n' || c == '\v' || c == '\r';
  }
};

inline std::ostream &operator<<(std::ostream &out, const Expr &e) {
  switch (e.type) {
  case Expr::Type::Atom:
    out << e.value;
    break;
  case Expr::Type::String:
    out << '"' << e.value << '"';
    break;
  case Expr::Type::List:
    out << '(';
    for (size_t i = 0; i < e.children.size(); ++i) {
      if (i > 0) {
        out << ' ';
      }
      out << e.children[i];
    }
    out << ')';
    break;
  }
  return out;
}