#include <emit.hh>
#include <vector>

namespace xlang {

class EmitFunctionVisitor : public xlangBaseVisitor {};

EmitVisitor::EmitVisitor(std::string_view outputfile) : output{outputfile} {}

std::any EmitVisitor::visitFile(xlangParser::FileContext *ctx) {
  output << "data $printformat = { b \"%d\\n\", b 0 }" << std::endl;
  visitChildren(ctx);
  return {};
}

std::any EmitVisitor::visitFunction(xlangParser::FunctionContext *ctx) {
  output << std::endl
         << "export function w $" << ctx->Identifier()->getSymbol()->getText()
         << "(";
  if (auto arg_list = ctx->argumentList()) {
    for (const auto arg : arg_list->Identifier()) {
      output << "w %_" << arg->getSymbol()->getText() << ", ";
    }
  }
  output << ")" << std::endl << "{" << std::endl << "@start" << std::endl;
  blockcount = 0;
  visitChildren(ctx);
  output << "  ret 0" << std::endl;
  output << "}" << std::endl;
  return {};
}

std::any EmitVisitor::visitStatement(xlangParser::StatementContext *ctx) {
  if (auto identifier = ctx->Identifier()) {
    tmpcount = 0;
    visitChildren(ctx);
    output << "  %_" << identifier->getSymbol()->getText() << " = w copy %t"
           << tmpcount - 1 << std::endl;
    return {};
  }
  if (ctx->If()) {
    int block = blockcount++;
    output << "@if" << block << "_expr" << std::endl;
    tmpcount = 0;
    visitExpr(ctx->expr());
    output << "  jnz %t" << tmpcount - 1 << ", @if" << block << "_then, @if"
           << block << "_else" << std::endl;
    output << "@if" << block << "_then" << std::endl;
    visitBlock(ctx->block(0));
    output << "  jmp @if" << block << "_end" << std::endl;
    output << "@if" << block << "_else" << std::endl;
    if (ctx->Else()) {
      visitBlock(ctx->block(1));
    }
    output << "@if" << block << "_end" << std::endl;
    return {};
  }
  if (ctx->While()) {
    int block = blockcount++;
    output << "@while" << block << "_expr" << std::endl;
    tmpcount = 0;
    visitExpr(ctx->expr());
    output << "  jnz %t" << tmpcount - 1 << ", @while" << block
           << "_body, @while" << block << "_end" << std::endl;
    output << "@while" << block << "_body" << std::endl;
    visitBlock(ctx->block(0));
    output << "  jmp @while" << block << "_expr" << std::endl;
    output << "@while" << block << "_end" << std::endl;
    return {};
  }
  if (ctx->Return()) {
    tmpcount = 0;
    visitChildren(ctx);
    output << "  ret %t" << tmpcount - 1 << std::endl;
    output << "@dead" << blockcount++ << std::endl;
    return {};
  }
  if (ctx->Print()) {
    tmpcount = 0;
    visitChildren(ctx);
    output << "  call $printf(l $printformat, ..., w %t" << tmpcount - 1 << ")"
           << std::endl;
    return {};
  }
  // unreachable
  return {};
}

#define OPERATOR(Operator, visitLeft, left, visitRight, right, ssa_op)  \
  if (ctx->Operator()) {                                                \
    visitLeft(ctx->left);                                               \
    int l = tmpcount - 1;                                               \
    visitRight(ctx->right);                                             \
    output << "  %t" << tmpcount << " = w " ssa_op " %t" << l << ", %t" \
           << tmpcount - 1 << std::endl;                                \
    tmpcount++;                                                         \
    return {};                                                          \
  }

std::any EmitVisitor::visitExpr(xlangParser::ExprContext *ctx) {
  OPERATOR(Less, visitSum, sum(0), visitSum, sum(1), "csltw");
  OPERATOR(LessEqual, visitSum, sum(0), visitSum, sum(1), "cslew");
  OPERATOR(Greater, visitSum, sum(0), visitSum, sum(1), "csgtw");
  OPERATOR(GreaterEqual, visitSum, sum(0), visitSum, sum(1), "csgew");
  OPERATOR(Equal, visitSum, sum(0), visitSum, sum(1), "ceqw");
  OPERATOR(NotEqual, visitSum, sum(0), visitSum, sum(1), "cnew");
  visitSum(ctx->sum(0));
  return {};
}

std::any EmitVisitor::visitSum(xlangParser::SumContext *ctx) {
  OPERATOR(Plus, visitSum, sum(), visitTerm, term(), "add");
  OPERATOR(Minus, visitSum, sum(), visitTerm, term(), "sub");
  visitTerm(ctx->term());
  return {};
}

std::any EmitVisitor::visitTerm(xlangParser::TermContext *ctx) {
  OPERATOR(Mul, visitTerm, term(), visitFactor, factor(), "mul");
  OPERATOR(Div, visitTerm, term(), visitFactor, factor(), "div");
  visitFactor(ctx->factor());
  return {};
}

std::any EmitVisitor::visitFactor(xlangParser::FactorContext *ctx) {
  if (auto integer = ctx->Integer()) {
    output << "  %t" << tmpcount++ << " = w copy "
           << integer->getSymbol()->getText() << std::endl;
    return {};
  }
  if (auto identifier = ctx->Identifier()) {
    if (ctx->LeftParen()) {
      std::vector<int> args;
      if (auto expr_list = ctx->exprList()) {
        for (auto expr : expr_list->expr()) {
          visitExpr(expr);
          args.push_back(tmpcount - 1);
        }
      }
      output << "  %t" << tmpcount++ << " = w call $"
             << identifier->getSymbol()->getText() << "(";
      for (auto arg : args) {
        output << "w %t" << arg << ", ";
      }
      output << ")" << std::endl;
    } else {
      output << "  %t" << tmpcount++ << " = w copy %_"
             << identifier->getSymbol()->getText() << std::endl;
    }
    return {};
  }
  if (auto sum = ctx->sum()) {
    visitSum(sum);
    return {};
  }
  // unreachable
  return {};
}

}  // namespace xlang