// 简单 AST 调试打印，便于前端验证。

#include "ast/AstNodes.h"

#include <iostream>
#include <ostream>
#include <string>

namespace ast {

static void PrintIndent(int depth) {
  for (int i = 0; i < depth; ++i) std::cout << "  ";
}

static void PrintLine(int depth, const std::string& text) {
  PrintIndent(depth);
  std::cout << text << "\n";
}

static void DumpExpr(const Expr* expr, int depth) {
  if (!expr) {
    PrintLine(depth, "NullExpr");
    return;
  }

  if (auto num = dynamic_cast<const NumberExpr*>(expr)) {
    PrintLine(depth, "NumberExpr value=" + std::to_string(num->value));
    return;
  }
  if (auto var = dynamic_cast<const VarExpr*>(expr)) {
    PrintLine(depth, "VarExpr name=" + var->name);
    return;
  }
  if (auto bin = dynamic_cast<const BinaryExpr*>(expr)) {
    const char* op = "?";
    switch (bin->op) {
      case BinaryOp::Add:
        op = "Add";
        break;
      case BinaryOp::Sub:
        op = "Sub";
        break;
      case BinaryOp::Mul:
        op = "Mul";
        break;
      case BinaryOp::Div:
        op = "Div";
        break;
    }
    PrintLine(depth, std::string("BinaryExpr op=") + op);
    DumpExpr(bin->lhs.get(), depth + 1);
    DumpExpr(bin->rhs.get(), depth + 1);
    return;
  }

  PrintLine(depth, "UnknownExpr");
}

void PrintAST(const CompUnit& cu) {
  PrintLine(0, "CompUnit");
  if (!cu.func) {
    PrintLine(1, "<empty>");
    std::cout << "\n";
    return;
  }

  PrintLine(1, "FuncDef name=" + cu.func->name);
  const auto& body = cu.func->body;
  if (!body) {
    PrintLine(2, "Block <null>");
    return;
  }

  PrintLine(2, "Block");
  for (const auto& decl : body->varDecls) {
    PrintLine(3, "VarDecl name=" + decl->name);
    if (decl->init) {
      PrintLine(4, "Init");
      DumpExpr(decl->init.get(), 5);
    }
  }
  for (const auto& stmt : body->stmts) {
    if (auto ret = dynamic_cast<ReturnStmt*>(stmt.get())) {
      PrintLine(3, "ReturnStmt");
      DumpExpr(ret->value.get(), 4);
    } else {
      PrintLine(3, "UnknownStmt");
    }
  }
  std::cout << "\n";
}

namespace {

std::string EscapeDotLabel(const std::string& s) {
  std::string out;
  out.reserve(s.size());
  for (char c : s) {
    if (c == '\\' || c == '"') out.push_back('\\');
    out.push_back(c);
  }
  return out;
}

class DotWriter {
 public:
  explicit DotWriter(std::ostream& os) : os_(os) {}

  int AddNode(const std::string& label) {
    int id = next_id_++;
    os_ << "  n" << id << " [label=\"" << EscapeDotLabel(label) << "\"];\n";
    return id;
  }

  void AddEdge(int from, int to) { os_ << "  n" << from << " -> n" << to << ";\n"; }

 private:
  std::ostream& os_;
  int next_id_ = 0;
};

int EmitExpr(const Expr* expr, DotWriter& dot) {
  if (!expr) return dot.AddNode("NullExpr");

  if (auto num = dynamic_cast<const NumberExpr*>(expr)) {
    return dot.AddNode("Number(" + std::to_string(num->value) + ")");
  }
  if (auto var = dynamic_cast<const VarExpr*>(expr)) {
    return dot.AddNode("Var(" + var->name + ")");
  }
  if (auto bin = dynamic_cast<const BinaryExpr*>(expr)) {
    const char* op = "?";
    switch (bin->op) {
      case BinaryOp::Add:
        op = "Add";
        break;
      case BinaryOp::Sub:
        op = "Sub";
        break;
      case BinaryOp::Mul:
        op = "Mul";
        break;
      case BinaryOp::Div:
        op = "Div";
        break;
    }
    int id = dot.AddNode(std::string("Binary(") + op + ")");
    int lhs = EmitExpr(bin->lhs.get(), dot);
    int rhs = EmitExpr(bin->rhs.get(), dot);
    dot.AddEdge(id, lhs);
    dot.AddEdge(id, rhs);
    return id;
  }
  return dot.AddNode("UnknownExpr");
}

}  // namespace

void PrintASTDot(const CompUnit& cu, std::ostream& os) {
  os << "digraph AST {\n";
  os << "  rankdir=TB;\n";
  os << "  node [shape=box, fontname=\"Courier\"];\n";

  DotWriter dot(os);
  int cu_id = dot.AddNode("CompUnit");
  if (cu.func) {
    int func_id = dot.AddNode("FuncDef(" + cu.func->name + ")");
    dot.AddEdge(cu_id, func_id);
    if (cu.func->body) {
      int block_id = dot.AddNode("Block");
      dot.AddEdge(func_id, block_id);

      for (const auto& decl : cu.func->body->varDecls) {
        int decl_id = dot.AddNode("VarDecl(" + decl->name + ")");
        dot.AddEdge(block_id, decl_id);
        if (decl->init) {
          int init_id = EmitExpr(decl->init.get(), dot);
          dot.AddEdge(decl_id, init_id);
        }
      }

      for (const auto& stmt : cu.func->body->stmts) {
        if (auto ret = dynamic_cast<ReturnStmt*>(stmt.get())) {
          int ret_id = dot.AddNode("Return");
          dot.AddEdge(block_id, ret_id);
          int value_id = EmitExpr(ret->value.get(), dot);
          dot.AddEdge(ret_id, value_id);
        } else {
          int stmt_id = dot.AddNode("Stmt");
          dot.AddEdge(block_id, stmt_id);
        }
      }
    }
  }
  os << "}\n";
}

}  // namespace ast