#include "Ast.h"

using namespace ast;

// attributes added from chapter 3 of the tutorial
std::unique_ptr<Module> AstTree::TheModule = make_unique<Module>("my cool jit", getGlobalContext());;
IRBuilder<> AstTree::Builder(getGlobalContext());
std::map<std::string, Value *> AstTree::NamedValues;

// error method from chapter 3
// HACK: duplicated function already present in example
std::unique_ptr<ExprAST> Error_Ast(const char *Str) {
  fprintf(stderr, "Error: %s\n", Str);
  return nullptr;
}

Value *ErrorV(const char *Str) {
  Error_Ast(Str);
  return nullptr;
}

Value *NumberExprAST::codegen() {
  return ConstantFP::get(getGlobalContext(), APFloat(Val));
}

Value *VariableExprAST::codegen() {
  // Look this variable up in the function.
  Value *V = AstTree::NamedValues[Name];
  if (!V)
  ErrorV("Unknown variable name");
  return V;
}

Value *BinaryExprAST::codegen() {
  Value *L = LHS->codegen();
  Value *R = RHS->codegen();
  if (!L || !R)
  return nullptr;

  switch (Op) {
    case '+':
    return AstTree::Builder.CreateFAdd(L, R, "addtmp");
    case '-':
    return AstTree::Builder.CreateFSub(L, R, "subtmp");
    case '*':
    return AstTree::Builder.CreateFMul(L, R, "multmp");
    case '<':
    L = AstTree::Builder.CreateFCmpULT(L, R, "cmptmp");
    // Convert bool 0/1 to double 0.0 or 1.0
    return AstTree::Builder.CreateUIToFP(L, Type::getDoubleTy(getGlobalContext()),
    "booltmp");
    default:
    return ErrorV("invalid binary operator");
  }
}

Value *CallExprAST::codegen() {
  // Look up the name in the global module table.
  Function *CalleeF = AstTree::TheModule->getFunction(Callee);
  if (!CalleeF)
  return ErrorV("Unknown function referenced");

  // If argument mismatch error.
  if (CalleeF->arg_size() != Args.size())
  return ErrorV("Incorrect # arguments passed");

  std::vector<Value *> ArgsV;
  for (unsigned i = 0, e = Args.size(); i != e; ++i) {
    ArgsV.push_back(Args[i]->codegen());
    if (!ArgsV.back())
    return nullptr;
  }
  return AstTree::Builder.CreateCall(CalleeF, ArgsV, "calltmp");
}

Function *PrototypeAST::codegen() {
  // Make the function type:  double(double,double) etc.
  std::vector<Type *> Doubles(Args.size(),
                              Type::getDoubleTy(getGlobalContext()));
  FunctionType *FT =
      FunctionType::get(Type::getDoubleTy(getGlobalContext()), Doubles, false);

  Function *F =
      Function::Create(FT, Function::ExternalLinkage, Name, AstTree::TheModule.get());

  // Set names for all arguments.
  unsigned Idx = 0;
  for (auto &Arg : F->args())
    Arg.setName(Args[Idx++]);

  return F;
}

Function *FunctionAST::codegen() {
  // First, check for an existing function from a previous 'extern' declaration.
  Function *TheFunction = AstTree::TheModule->getFunction(Proto->getName());

  if (!TheFunction)
  TheFunction = Proto->codegen();

  if (!TheFunction)
  return nullptr;

  if (!TheFunction->empty())
  return (Function*)ErrorV("Function cannot be redefined.");

  // Create a new basic block to start insertion into.
  BasicBlock *BB = BasicBlock::Create(getGlobalContext(), "entry", TheFunction);
  AstTree::Builder.SetInsertPoint(BB);

  // Record the function arguments in the NamedValues map.
  AstTree::NamedValues.clear();
  for (auto &Arg : TheFunction->args())
  AstTree::NamedValues[Arg.getName()] = &Arg;

  if (Value *RetVal = Body->codegen()) {
    // Finish off the function.
    AstTree::Builder.CreateRet(RetVal);

    // Validate the generated code, checking for consistency.
    verifyFunction(*TheFunction);

    return TheFunction;
  }

  // Error reading body, remove function.
  TheFunction->eraseFromParent();
  return nullptr;
}