#ifndef __AST_HPP__
#define __AST_HPP__

#include <iostream>
#include <vector>

class AST {
public:
  virtual ~AST() {}
  virtual void printOn(std::ostream &out) const = 0;
};

inline std::ostream &operator<<(std::ostream &out, const AST &t) {
  t.printOn(out);
  return out;
}

class Expr : public AST {
public:
  virtual void compile() const = 0;
};

class Stmt : public AST {
public:
  virtual void compile() const = 0;
};

class Id : public Expr {
public:
  Id(char v) : var(v) {}
  virtual void printOn(std::ostream &out) const override {
    out << "Id(" << var << ")";
  }
  virtual void compile() const override {
    std::cout << "  pushl " << 4 * (var - 'a') << "(%edi)\n";
  }

private:
  char var;
};

class Const : public Expr {
public:
  Const(int n) : num(n) {}
  virtual void printOn(std::ostream &out) const override {
    out << "Const(" << num << ")";
  }
  virtual void compile() const override {
    std::cout << "  pushl $" << num << "\n";
  }

private:
  int num;
};

class BinOp : public Expr {
public:
  BinOp(Expr *l, char o, Expr *r) : left(l), op(o), right(r) {}
  ~BinOp() {
    delete left;
    delete right;
  }
  virtual void printOn(std::ostream &out) const override {
    out << op << "(" << *left << ", " << *right << ")";
  }
  virtual void compile() const override {
    left->compile();
    right->compile();
    std::cout << "  popl %ebx\n"  // right
              << "  popl %eax\n"; // left
    switch (op) {
    case '+':
      std::cout << "  addl %ebx, %eax\n"
                << "  pushl %eax\n";
      break;
    case '-':
      std::cout << "  subl %ebx, %eax\n"
                << "  pushl %eax\n";
      break;
    case '*':
      std::cout << "  mull %ebx\n"
                << "  pushl %eax\n";
      break;
    case '/':
      std::cout << "  cdq\n"
                << "  divl %ebx\n"
                << "  pushl %eax\n";
      break;
    case '%':
      std::cout << "  cdq\n"
                << "  divl %ebx\n"
                << "  pushl %edx\n";
      break;
    }
  }

private:
  Expr *left;
  char op;
  Expr *right;
};

class Let : public Stmt {
public:
  Let(char v, Expr *e) : var(v), expr(e) {}
  ~Let() { delete expr; }
  virtual void printOn(std::ostream &out) const override {
    out << "Let(" << var << " = " << *expr << ")";
  }
  virtual void compile() const override {
    expr->compile();
    std::cout << "  popl %eax\n"
              << "  movl %eax, " << 4 * (var - 'a') << "(%edi)\n";
  }

private:
  char var;
  Expr *expr;
};

class Print : public Stmt {
public:
  Print(Expr *e) : expr(e) {}
  ~Print() { delete expr; }
  virtual void printOn(std::ostream &out) const override {
    out << "Print(" << *expr << ")";
  }
  virtual void compile() const override {
    std::cout << "  pushl %edi\n";   // Save %edi to protect it.
    expr->compile();                 // Push number to print.
    std::cout << "  subl $8, %esp\n" // Space for access link and result addr.
              << "  call _writeInteger\n"
              << "  addl $12, %esp\n" // Clear _writeInteger frame.
              << "  movl $NL, %eax\n"
              << "  pushl %eax\n"    // Push string to print.
              << "  subl $8, %esp\n" // Space for access link and result addr.
              << "  call _writeString\n"
              << "  addl $12, %esp\n" // Clear _writeInteger frame.
              << "  popl %edi\n";     // Restore %edi.
  }

private:
  Expr *expr;
};

class If : public Stmt {
public:
  If(Expr *c, Stmt *s1, Stmt *s2 = nullptr) : cond(c), stmt1(s1), stmt2(s2) {}
  ~If() {
    delete cond;
    delete stmt1;
    delete stmt2;
  }
  virtual void printOn(std::ostream &out) const override {
    out << "If(" << *cond << ", " << *stmt1;
    if (stmt2 != nullptr)
      out << ", " << *stmt2;
    out << ")";
  }
  virtual void compile() const override {
    static int counter = 0;
    cond->compile();
    int l_false = counter++;
    std::cout << "  popl %eax\n"
              << "  andl %eax, %eax\n"
              << "  jz Lif" << l_false << "\n";
    stmt1->compile();
    int l_end = counter++;
    std::cout << "  jmp Lif" << l_end << "\n"
              << "Lif" << l_false << ":\n";
    if (stmt2 != nullptr)
      stmt2->compile();
    std::cout << "Lif" << l_end << ":\n";
  }

private:
  Expr *cond;
  Stmt *stmt1;
  Stmt *stmt2;
};

class For : public Stmt {
public:
  For(Expr *e, Stmt *s) : expr(e), stmt(s) {}
  ~For() {
    delete expr;
    delete stmt;
  }
  virtual void printOn(std::ostream &out) const override {
    out << "For(" << *expr << ", " << *stmt << ")";
  }
  virtual void compile() const override {
    static int counter = 0;
    expr->compile();
    int l_loop = counter++;
    std::cout << "Lfor" << l_loop << ":\n"
              << "  popl %eax\n"
              << "  cmpl $0, %eax\n";
    int l_end = counter++;
    std::cout << "  jle Lfor" << l_end << "\n"
              << "  decl %eax\n"
              << "  pushl %eax\n";
    stmt->compile();
    std::cout << "  jmp Lfor" << l_loop << "\n"
              << "Lfor" << l_end << ":\n";
  }

private:
  Expr *expr;
  Stmt *stmt;
};

class Block : public Stmt {
public:
  Block() : stmt_list() {}
  ~Block() {
    for (Stmt *s : stmt_list)
      delete s;
  }
  void append(Stmt *s) { stmt_list.push_back(s); }
  virtual void printOn(std::ostream &out) const override {
    out << "Block(";
    bool first = true;
    for (Stmt *s : stmt_list) {
      if (!first)
        out << ", ";
      first = false;
      out << *s;
    }
    out << ")";
  }
  virtual void compile() const override {
    for (Stmt *s : stmt_list)
      s->compile();
  }

private:
  std::vector<Stmt *> stmt_list;
};

inline void prologue() {
  std::cout << ".text\n"
            << ".global _start\n"
            << "\n"
            << "_start:\n"
            << "  movl $var, %edi\n"
            << "\n";
}

inline void epilogue() {
  std::cout << "\n"
            << "  movl $1, %eax\n"
            << "  movl $0, %ebx\n"
            << "  int $0x80\n"
            << "\n"
            << ".data\n"
            << "var:\n"
            << ".rept 26\n"
            << ".long 0\n"
            << ".endr\n"
            << "NL:\n"
            << ".asciz \"\\n\"\n"; // .asciz "\n"
}

#endif