#include "frontend/SyntaxTreePrinter.h"

#include <string>

namespace {

std::string GetTokenName(const antlr4::Token* tok, antlr4::Parser* parser) {
  if (!tok || !parser) {
    return "UNKNOWN";
  }
  const int token_type = tok->getType();
  const auto& vocab = parser->getVocabulary();
  std::string token_name(vocab.getSymbolicName(token_type));
  if (token_name.empty()) {
    token_name = std::string(vocab.getLiteralName(token_type));
  }
  if (token_name.empty()) {
    token_name = std::to_string(token_type);
  }
  return token_name;
}

std::string RuleName(antlr4::ParserRuleContext* rule, antlr4::Parser* parser) {
  if (!rule || !parser) {
    return "unknown";
  }
  const int idx = rule->getRuleIndex();
  const auto& names = parser->getRuleNames();
  if (idx >= 0 && idx < static_cast<int>(names.size())) {
    return names[static_cast<size_t>(idx)];
  }
  return "unknown";
}

void PrintSyntaxTreeImpl(antlr4::tree::ParseTree* node, antlr4::Parser* parser,
                         std::ostream& os, const std::string& prefix,
                         bool is_last, bool is_root) {
  if (!node) {
    return;
  }

  std::string label;
  if (auto* terminal = dynamic_cast<antlr4::tree::TerminalNode*>(node)) {
    label = GetTokenName(terminal->getSymbol(), parser) + ": " + terminal->getText();
  } else if (auto* rule = dynamic_cast<antlr4::ParserRuleContext*>(node)) {
    label = RuleName(rule, parser);
  } else {
    label = "unknown";
  }

  if (is_root) {
    os << label << "\n";
  } else {
    os << prefix << (is_last ? "`-- " : "|-- ") << label << "\n";
  }

  const std::string child_prefix =
      is_root ? "" : prefix + (is_last ? "    " : "|   ");
  const size_t child_count = node->children.size();
  for (size_t i = 0; i < child_count; ++i) {
    PrintSyntaxTreeImpl(node->children[i], parser, os, child_prefix,
                        i + 1 == child_count, false);
  }
}

}  // namespace

void PrintSyntaxTree(antlr4::tree::ParseTree* tree, antlr4::Parser* parser,
                     std::ostream& os) {
  PrintSyntaxTreeImpl(tree, parser, os, "", true, true);
}