#pragma once

#include "basic_printers.hpp"
#include "statement_nodes.hpp"
#include "tree_sitter_wrapper.hpp"

#include <cstddef>
#include <optional>
#include <unordered_map>
#include <vector>

namespace names {

// IN PROGRESS
class NameTree {
public:
  NameTree() {
    nodes_.emplace_back(); // root
  }

  bool insert(const std::string &path, nodes::Statement &&statement);

  // bool insert(const std::string &path, const nodes::Statement &statement);

  nodes::CombineResult insert_combine(const std::string &path,
                                      nodes::Statement &&statement);

  std::optional<nodes::Statement *> find(const std::string &path);

  std::optional<const nodes::Statement *> find(const std::string &path) const;

  void print(printers::Printer &printer) const;

  // TODO
  void add_statement_children_to_tree();

private:
  struct Node {
  public:
    std::optional<nodes::Statement *> get_statement() {
      if (statement_.has_value()) {
        return &statement_.value();
      }
      return std::nullopt;
    }

    std::optional<const nodes::Statement *> get_statement() const {
      if (statement_.has_value()) {
        return &statement_.value();
      }
      return std::nullopt;
    }

    bool set_statement(nodes::Statement &&statement) {
      if (statement_.has_value()) {
        return false;
      }
      statement_ = std::move(statement);
      return true;
    }

    bool set_statement(const nodes::Statement &statement) {
      if (statement_.has_value()) {
        return false;
      }
      statement_ = statement;
      return true;
    }

    std::optional<size_t> find(size_t name_id) const {
      auto name_iter = children_.find(name_id);

      if (name_iter == children_.end()) {
        return std::nullopt;
      }

      return name_iter->second;
    }

    bool insert(size_t name_id, size_t node_id) {
      if (children_.count(name_id) != 0) {
        return false;
      }
      children_[name_id] = node_id;
      return true;
    }

    void print(const std::vector<Node> &nodes,
               printers::Printer &printer) const {
      printer.print("Node - has statement: ");
      printer.print(statement_.has_value() ? "true" : "false");
      printer.indent();
      for (auto &iter : children_) {
        printer.new_indent_line();
        printer.print(std::to_string(iter.first) + " -> ");
        nodes[iter.second].print(nodes, printer);
      }
      printer.deindent();
      printer.new_indent_line();
    }

  private:
    std::unordered_map<size_t, size_t> children_;
    std::optional<nodes::Statement> statement_;
  };

  size_t get_root() const { return 0; }

  size_t add_node(size_t current_node, const std::vector<size_t> &name_ids,
                  size_t current_name = 0);

  std::optional<size_t> find_node(size_t current_node,
                                  const std::vector<size_t> &name_ids,
                                  size_t current_name = 0) const;

  std::vector<size_t> slice_path_to_name_ids(const std::string &path);

  std::optional<std::vector<size_t>>
  slice_path_to_existing_name_ids(const std::string &path) const;

  size_t add_name_id(const std::string &name);

  std::optional<size_t> find_name_id(const std::string &name) const;

private:
  std::unordered_map<std::string, size_t> name_to_id_;
  std::vector<Node> nodes_;
};

} // namespace names