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

namespace xlang {

std::string EmitVisitor::tmp() {
  last = "%t" + std::to_string(tmpcount++);
  return last;
}

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 param_list = ctx->parameterList()) {
    for (const auto param : param_list->Identifier()) {
      output << "w %_" << param->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;
    visitValue(ctx->value());
    output << "  %_" << identifier->getSymbol()->getText() << " = w copy "
           << last << std::endl;
    return {};
  }
  if (ctx->If()) {
    int block = blockcount++;

    output << "@if" << block << "_expr" << std::endl;
    tmpcount = 0;
    visitCondition(ctx->condition());
    output << "  jnz " << last << ", @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++;

    loopstack.push_back(block);
    output << "@while" << block << "_expr" << std::endl;
    tmpcount = 0;
    visitCondition(ctx->condition());
    output << "  jnz " << last << ", @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;
    loopstack.pop_back();
    return {};
  }
  if (ctx->Break()) {
    output << "  jmp @while" << loopstack.back() << "_end" << std::endl;
    output << "@dead" << blockcount++ << std::endl;
  }
  if (ctx->Continue()) {
    output << "  jmp @while" << loopstack.back() << "_expr" << std::endl;
    output << "@dead" << blockcount++ << std::endl;
  }
  if (ctx->Return()) {
    tmpcount = 0;
    visitValue(ctx->value());
    output << "  ret " << last << std::endl;
    output << "@dead" << blockcount++ << std::endl;
    return {};
  }
  if (ctx->Print()) {
    tmpcount = 0;
    visitValue(ctx->value());
    output << "  call $printf(l $printformat, ..., w " << last << ")"
           << std::endl;
    return {};
  }
  // unreachable
  return {};
}

#define OPERATOR(Operator, visitLeft, left, visitRight, right, ssa_op) \
  if (ctx->Operator()) {                                               \
    visitLeft(ctx->left);                                              \
    std::string l = last;                                              \
    visitRight(ctx->right);                                            \
    std::string r = last;                                              \
    output << "  " << tmp() << " = w " ssa_op " " << l << ", " << r    \
           << std::endl;                                               \
    return {};                                                         \
  }

std::any EmitVisitor::visitCondition(xlangParser::ConditionContext *ctx) {
  OPERATOR(And, visitCondition, condition(), visitBoolean, boolean(), "and");
  OPERATOR(Or, visitCondition, condition(), visitBoolean, boolean(), "or");
  OPERATOR(Xor, visitCondition, condition(), visitBoolean, boolean(), "xor");
  visitBoolean(ctx->boolean());
  return {};
}

std::any EmitVisitor::visitBoolean(xlangParser::BooleanContext *ctx) {
  if (ctx->Not()) {
    visitBoolean(ctx->boolean());
    std::string b = last;
    output << "  " << tmp() << " = w xor 1, " << b << std::endl;
    return {};
  }
  if (ctx->True()) {
    last = "1";
    return {};
  }
  if (ctx->False()) {
    last = "0";
    return {};
  }
  OPERATOR(Less, visitExpr, expr(0), visitExpr, expr(1), "csltw");
  OPERATOR(LessEqual, visitExpr, expr(0), visitExpr, expr(1), "cslew");
  OPERATOR(Greater, visitExpr, expr(0), visitExpr, expr(1), "csgtw");
  OPERATOR(GreaterEqual, visitExpr, expr(0), visitExpr, expr(1), "csgew");
  OPERATOR(Equal, visitExpr, expr(0), visitExpr, expr(1), "ceqw");
  OPERATOR(NotEqual, visitExpr, expr(0), visitExpr, expr(1), "cnew");
  visitChildren(ctx);
  return {};
}

std::any EmitVisitor::visitExpr(xlangParser::ExprContext *ctx) {
  OPERATOR(Plus, visitExpr, expr(), visitTerm, term(), "add");
  OPERATOR(Minus, visitExpr, expr(), visitTerm, term(), "sub");
  OPERATOR(BitAnd, visitExpr, expr(), visitTerm, term(), "and");
  OPERATOR(BitOr, visitExpr, expr(), visitTerm, term(), "or");
  OPERATOR(BitXor, visitExpr, expr(), visitTerm, term(), "xor");
  OPERATOR(ShiftLeft, visitExpr, expr(), visitTerm, term(), "shl");
  OPERATOR(ShiftRight, visitExpr, expr(), visitTerm, term(), "shr");
  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");
  OPERATOR(Rem, visitTerm, term(), visitFactor, factor(), "rem");
  visitFactor(ctx->factor());
  return {};
}

std::any EmitVisitor::visitFactor(xlangParser::FactorContext *ctx) {
  if (ctx->Minus()) {
    visitFactor(ctx->factor());
    std::string f = last;
    output << "  " << tmp() << " = w neg " << f << std::endl;
    return {};
  }
  if (ctx->BitNot()) {
    visitFactor(ctx->factor());
    std::string f = last;
    output << "  " << tmp() << " = w xor -1, " << f << std::endl;
    return {};
  }
  if (auto integer = ctx->Integer()) {
    last = integer->getSymbol()->getText();
    return {};
  }
  visitChildren(ctx);
  return {};
}

std::any EmitVisitor::visitVariable(xlangParser::VariableContext *ctx) {
  auto name = ctx->Identifier()->getSymbol()->getText();

  if (ctx->LeftParen()) {
    std::vector<std::string> args;

    if (auto arg_list = ctx->argumentList()) {
      for (auto value : arg_list->value()) {
        visitValue(value);
        args.emplace_back(last);
      }
    }
    output << "  " << tmp() << " = w call $" << name << "(";
    for (auto arg : args) {
      output << "w " << arg << ", ";
    }
    output << ")" << std::endl;
  } else {
    last = "%_" + name;
  }
  return {};
}

}  // namespace xlang