#ifndef __AST_HPP__
#define __AST_HPP__

#include <iostream>
#include <vector>

#include "symbol.hpp"

void yyerror(const char *msg);

inline std::ostream &operator<<(std::ostream &out, Type t) {
  switch (t) {
  case TYPE_int:
    out << "int";
    break;
  case TYPE_bool:
    out << "bool";
    break;
  }
  return out;
}

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

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;
  void type_check(Type t) {
    sem();
    if (type != t)
      yyerror("Type mismatch");
  }

protected:
  Type type;
};

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

extern std::vector<int> rt_stack;

class Id : public Expr {
public:
  Id(char v) : var(v), nestingDiff(-1), offset(-1) {}
  virtual void printOn(std::ostream &out) const override {
    out << "Id(" << var << "@" << offset << "/" << nestingDiff << ")";
  }
  virtual void compile() const override {
    if (nestingDiff == 0)
      // Local variable.
      std::cout << "  pushl " << 4 * offset << "(%ebp)\n";
    else {
      // Non-local variable; follow nestingDiff access links.
      std::cout << "  movl 0(%ebp), %esi\n";
      for (int i = 1; i < nestingDiff; ++i)
        std::cout << "  movl 0(%esi), %esi\n";
      std::cout << "  pushl " << 4 * offset << "(%esi)\n";
    }
  }
  virtual void sem() override {
    SymbolEntry *e = st.lookup(var);
    type = e->type;
    nestingDiff = st.getCurrentNesting() - e->nesting;
    offset = e->offset;
  }

private:
  char var;
  int nestingDiff, offset;
};

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";
  }
  virtual void sem() override { type = TYPE_int; }

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;
    }
  }
  virtual void sem() override {
    left->type_check(TYPE_int);
    right->type_check(TYPE_int);
    switch (op) {
    case '+':
    case '-':
    case '*':
    case '/':
    case '%':
      type = TYPE_int;
      break;
    case '=':
    case '<':
    case '>':
      type = TYPE_bool;
      break;
    }
  }

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

class Let : public Stmt {
public:
  Let(char v, Expr *e) : var(v), nestingDiff(-1), offset(-1), expr(e) {}
  ~Let() { delete expr; }
  virtual void printOn(std::ostream &out) const override {
    out << "Let(" << var << "@" << offset << "/" << nestingDiff << " = "
        << *expr << ")";
  }
  virtual void compile() const override {
    expr->compile();
    std::cout << "  popl %eax\n";
    if (nestingDiff == 0)
      // Local variable.
      std::cout << "  movl %eax, " << 4 * offset << "(%ebp)\n";
    else {
      // Non-local variable; follow nestingDiff access links.
      std::cout << "  movl 0(%ebp), %esi\n";
      for (int i = 1; i < nestingDiff; ++i)
        std::cout << "  movl 0(%esi), %esi\n";
      std::cout << "  movl %eax, " << 4 * offset << "(%esi)\n";
    }
  }
  virtual void sem() override {
    SymbolEntry *lhs = st.lookup(var);
    expr->type_check(lhs->type);
    nestingDiff = st.getCurrentNesting() - lhs->nesting;
    offset = lhs->offset;
  }

private:
  char var;
  int nestingDiff, offset;
  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 {
    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.
  }
  virtual void sem() override { expr->type_check(TYPE_int); }

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";
  }
  virtual void sem() override {
    cond->type_check(TYPE_bool);
    stmt1->sem();
    if (stmt2 != nullptr)
      stmt2->sem();
  }

private:
  Expr *cond;
  Stmt *stmt1, *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";
  }
  virtual void sem() override {
    expr->type_check(TYPE_int);
    stmt->sem();
  }

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

class Decl : public AST {
public:
  Decl(char c, Type t) : var(c), type(t) {}
  virtual void printOn(std::ostream &out) const override {
    out << "Decl(" << var << " : " << type << ")";
  }
  virtual void sem() override { st.insert(var, type); }

private:
  char var;
  Type type;
};

class Block : public Stmt {
public:
  Block() : decl_list(), stmt_list(), size(0) {}
  ~Block() {
    for (Decl *d : decl_list)
      delete d;
    for (Stmt *s : stmt_list)
      delete s;
  }
  void append_decl(Decl *d) { decl_list.push_back(d); }
  void append_stmt(Stmt *s) { stmt_list.push_back(s); }
  void merge(Block *b) {
    stmt_list = b->stmt_list;
    b->stmt_list.clear();
    delete b;
  }
  virtual void printOn(std::ostream &out) const override {
    out << "Block(";
    bool first = true;
    for (Decl *d : decl_list) {
      if (!first)
        out << ", ";
      first = false;
      out << *d;
    }
    for (Stmt *s : stmt_list) {
      if (!first)
        out << ", ";
      first = false;
      out << *s;
    }
    out << ")";
  }
  virtual void compile() const override {
    for (int i = 0; i < size; ++i)
      std::cout << "  pushl $0\n";      // Push zeroes for all local variables.
    std::cout << "  pushl %ebp\n"       // Push previous frame pointer.
              << "  movl %esp, %ebp\n"; // Set new frame pointer.
    for (Stmt *s : stmt_list)
      s->compile();
    std::cout << "  popl %ebp\n"                       // Restore frame pointer.
              << "  addl $" << 4 * size << ", %esp\n"; // Clear local variables.
  }
  virtual void sem() override {
    st.openScope();
    for (Decl *d : decl_list)
      d->sem();
    for (Stmt *s : stmt_list)
      s->sem();
    size = st.getSizeOfCurrentScope();
    st.closeScope();
  }

private:
  std::vector<Decl *> decl_list;
  std::vector<Stmt *> stmt_list;
  int size;
};

inline void prologue() {
  std::cout << ".text\n"
            << ".global _start\n"
            << "\n"
            << "_start:\n"
            << "\n";
}

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

#endif