#include "frontend/SyntaxTreePrinter.h"

#include <string>
#include <vector>

#include "SysYParser.h"

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;
}

bool KeepImportantToken(const std::string& token_name) {
  return token_name == "Ident" || token_name == "Number" ||
         token_name == "Assign" || token_name == "AddOp";
}

std::string PrettyPrimary(SysYParser::PrimaryContext* primary) {
  if (!primary) {
    return "";
  }
  if (primary->Number()) {
    return primary->Number()->getText();
  }
  if (primary->Ident()) {
    return primary->Ident()->getText();
  }
  if (primary->exp()) {
    return "(" + primary->exp()->getText() + ")";
  }
  return primary->getText();
}

std::string PrettyAddExp(SysYParser::AddExpContext* add_exp) {
  if (!add_exp) {
    return "";
  }
  const auto terms = add_exp->primary();
  if (terms.empty()) {
    return "";
  }
  std::string out = PrettyPrimary(terms[0]);
  for (size_t i = 1; i < terms.size(); ++i) {
    out += " + " + PrettyPrimary(terms[i]);
  }
  return out;
}

std::string PrettyExp(SysYParser::ExpContext* exp) {
  if (!exp || !exp->addExp()) {
    return "";
  }
  return PrettyAddExp(exp->addExp());
}

std::string PrettyRuleText(antlr4::ParserRuleContext* rule) {
  if (!rule) {
    return "";
  }
  if (auto* var_decl = dynamic_cast<SysYParser::VarDeclContext*>(rule)) {
    std::string out = "int " + var_decl->Ident()->getText();
    if (var_decl->exp()) {
      out += " = " + PrettyExp(var_decl->exp());
    }
    out += ";";
    return out;
  }
  if (auto* ret = dynamic_cast<SysYParser::ReturnStmtContext*>(rule)) {
    return "return " + PrettyExp(ret->exp()) + ";";
  }
  if (dynamic_cast<SysYParser::FuncDefContext*>(rule) != nullptr) {
    return "int main()";
  }
  if (auto* stmt = dynamic_cast<SysYParser::StmtContext*>(rule)) {
    if (stmt->varDecl()) {
      return PrettyRuleText(stmt->varDecl());
    }
    if (stmt->returnStmt()) {
      return PrettyRuleText(stmt->returnStmt());
    }
  }
  if (auto* exp = dynamic_cast<SysYParser::ExpContext*>(rule)) {
    return PrettyExp(exp);
  }
  if (auto* add_exp = dynamic_cast<SysYParser::AddExpContext*>(rule)) {
    return PrettyAddExp(add_exp);
  }
  if (auto* primary = dynamic_cast<SysYParser::PrimaryContext*>(rule)) {
    return PrettyPrimary(primary);
  }
  return "";
}

bool HasVisibleNode(antlr4::tree::ParseTree* node, antlr4::Parser* parser) {
  auto* terminal = dynamic_cast<antlr4::tree::TerminalNode*>(node);
  if (terminal) {
    const std::string token_name = GetTokenName(terminal->getSymbol(), parser);
    return KeepImportantToken(token_name);
  }
  for (auto* child : node->children) {
    if (HasVisibleNode(child, parser)) {
      return true;
    }
  }
  return false;
}

std::string RuleName(antlr4::tree::ParseTree* node, antlr4::Parser* parser) {
  auto* rule = dynamic_cast<antlr4::ParserRuleContext*>(node);
  if (!parser || !rule) {
    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";
}

std::string NodeLabel(antlr4::tree::ParseTree* node, antlr4::Parser* parser) {
  auto* terminal = dynamic_cast<antlr4::tree::TerminalNode*>(node);
  if (terminal) {
    return GetTokenName(terminal->getSymbol(), parser) + ": " + node->getText();
  }

  const std::string rule_name = RuleName(node, parser);
  auto* rule = dynamic_cast<antlr4::ParserRuleContext*>(node);
  const std::string pretty = PrettyRuleText(rule);
  if (!pretty.empty()) {
    return rule_name + " (" + pretty + ")";
  }
  return rule_name;
}

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

  if (is_root) {
    os << NodeLabel(node, parser) << "\n";
  } else {
    os << prefix << (is_last ? "└── " : "├── ") << NodeLabel(node, parser)
       << "\n";
  }

  std::vector<antlr4::tree::ParseTree*> children;
  for (auto* child : node->children) {
    if (HasVisibleNode(child, parser)) {
      children.push_back(child);
    }
  }

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

}  // namespace

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