
#pragma once

#include <memory>
#include <string>
#include <vector>

namespace ast {

enum class BinaryOp { Add, Sub, Mul, Div };

struct Expr {
  virtual ~Expr() = default;
};

struct NumberExpr : Expr {
  int value{};
  explicit NumberExpr(int v) : value(v) {}
};

struct VarExpr : Expr {
  std::string name;
  explicit VarExpr(std::string n) : name(std::move(n)) {}
};

struct BinaryExpr : Expr {
  BinaryOp op;
  std::shared_ptr<Expr> lhs;
  std::shared_ptr<Expr> rhs;
  BinaryExpr(BinaryOp op, std::shared_ptr<Expr> lhs, std::shared_ptr<Expr> rhs)
      : op(op), lhs(std::move(lhs)), rhs(std::move(rhs)) {}
};

struct Stmt {
  virtual ~Stmt() = default;
};

struct ReturnStmt : Stmt {
  std::shared_ptr<Expr> value;
  explicit ReturnStmt(std::shared_ptr<Expr> v) : value(std::move(v)) {}
};

struct VarDecl {
  std::string name;
  std::shared_ptr<Expr> init;  // nullptr if no initializer
  VarDecl(std::string n, std::shared_ptr<Expr> i)
      : name(std::move(n)), init(std::move(i)) {}
};

struct Block {
  std::vector<std::shared_ptr<VarDecl>> varDecls;
  std::vector<std::shared_ptr<Stmt>> stmts;
};

struct FuncDef {
  std::string name;
  std::shared_ptr<Block> body;
  FuncDef(std::string n, std::shared_ptr<Block> b)
      : name(std::move(n)), body(std::move(b)) {}
};

struct CompUnit {
  std::shared_ptr<FuncDef> func;
  explicit CompUnit(std::shared_ptr<FuncDef> f) : func(std::move(f)) {}
};

// 调试打印
void PrintAST(const CompUnit& cu);

}  // namespace ast
