#include "ast.hpp"

#include <iostream>


LLVMContext TheContext;
map<string, AllocaInst*> NamedValues;
IRBuilder<> Builder(TheContext);
Module* TheModule;
FunctionPassManager *TheFPM;

ExprAST::~ExprAST() {

}

InnerExprAST::InnerExprAST(ExprAST* a) {
  _nodes.resize(0);
  _nodes.push_back(a);
}

InnerExprAST::InnerExprAST(ExprAST* a, ExprAST* b) {
  _nodes.resize(0);
  _nodes.push_back(a);
  _nodes.push_back(b);
}

InnerExprAST::InnerExprAST(ExprAST* a, ExprAST* b, ExprAST* c) {
  _nodes.resize(0);
  _nodes.push_back(a);
  _nodes.push_back(b);
  _nodes.push_back(c);
}

InnerExprAST::InnerExprAST(ExprAST* a, ExprAST* b, ExprAST* c, ExprAST* d) {
  _nodes.resize(0);
  _nodes.push_back(a);
  _nodes.push_back(b);
  _nodes.push_back(c);
  _nodes.push_back(d);
}

InnerExprAST::InnerExprAST(vector<ExprAST*> a) {
  _nodes = a;
}

InnerExprAST::~InnerExprAST() {
  for (unsigned i = 0; i < _nodes.size(); i++)
    delete _nodes[i];
}

FunctionAST::~FunctionAST() {
  delete Body;
}

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

Value* VariableExprAST::codegen() const {
  AllocaInst* tmp = NamedValues[Name];
  if (tmp == NULL) {
    cerr << "Promenljiva " + Name + " nije definisana" << endl;
    return NULL;
  }
  return Builder.CreateLoad(Type::getDoubleTy(TheContext), tmp, Name.c_str());
}

Value* AddExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  return Builder.CreateFAdd(l, d, "addtmp");
}

Value* SubExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  return Builder.CreateFSub(l, d, "subtmp");
}

Value* MulExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  return Builder.CreateFMul(l, d, "multmp");
}

Value* DivExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  return Builder.CreateFDiv(l, d, "divtmp");
}

Value* LTExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  l = Builder.CreateFCmpULT(l, d, "lttmp");
  return Builder.CreateUIToFP(l, Type::getDoubleTy(TheContext) , "booltmp");
}

Value* GTExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  l = Builder.CreateFCmpUGT(l, d, "gttmp");
  return Builder.CreateUIToFP(l, Type::getDoubleTy(TheContext) , "booltmp");
}

Value* SeqExprAST::codegen() const {
  Value *l = _nodes[0]->codegen();
  Value *d = _nodes[1]->codegen();
  if (l == NULL || d == NULL)
    return NULL;
  return d;
}

Value* AssignExprAST::codegen() const {
  Value* tmp = _nodes[0]->codegen();
  if (tmp == NULL)
    return NULL;
  AllocaInst* Alloca = NamedValues[Name];
  if (Alloca == NULL) {
    cerr << "Promenljiva " + Name + " nije definisana" << endl;
    return NULL;
  }
  Builder.CreateStore(tmp, Alloca);
  return tmp;
}

Value* CallExprAST::codegen() const {
  Function *CalleeF = TheModule->getFunction(Callee);
  if (!CalleeF) {
    cerr << "Fja " << Callee << " nije definisana" << endl;
    return NULL;
  }
  
  // If argument mismatch error.
  if (CalleeF->arg_size() != _nodes.size()) {
    cout << "Fja " << Callee << " prima " << CalleeF->arg_size() << " argumenata" << endl;
    return NULL;
  }
  
  vector<Value *> ArgsV;
  for (unsigned i = 0, e = _nodes.size(); i != e; ++i) {
    ArgsV.push_back(_nodes[i]->codegen());
    if (!ArgsV.back())
      return NULL;
  }
  
  return Builder.CreateCall(CalleeF, ArgsV, "calltmp");
}

Value* IfExprAST::codegen() const {
  Value* CondV = _nodes[0]->codegen();
  if (CondV == NULL)
    return NULL;
  CondV = Builder.CreateFCmpONE(
      CondV, ConstantFP::get(TheContext, APFloat(0.0)), "ifcond");
  Function *TheFunction = Builder.GetInsertBlock()->getParent();
  
  // Create blocks for the then and else cases.  Insert the 'then' block at the
  // end of the function.
  BasicBlock *ThenBB =
    BasicBlock::Create(TheContext, "then", TheFunction);
  BasicBlock *ElseBB = BasicBlock::Create(TheContext, "else");
  BasicBlock *MergeBB = BasicBlock::Create(TheContext, "ifcont");
  
  Builder.CreateCondBr(CondV, ThenBB, ElseBB);

  Builder.SetInsertPoint(ThenBB);
  Value* ThenV = _nodes[1]->codegen();
  if (ThenV == NULL)
    return NULL;
  Builder.CreateBr(MergeBB);
  ThenBB = Builder.GetInsertBlock();

  TheFunction->getBasicBlockList().push_back(ElseBB);
  Builder.SetInsertPoint(ElseBB);
  Value* ElseV = _nodes[2]->codegen();
  if (ElseV == NULL)
    return NULL;
  Builder.CreateBr(MergeBB);
  ElseBB = Builder.GetInsertBlock();

  TheFunction->getBasicBlockList().push_back(MergeBB);
  Builder.SetInsertPoint(MergeBB);
  PHINode* PN = Builder.CreatePHI(Type::getDoubleTy(TheContext), 2, "iftmp");
  PN->addIncoming(ThenV, ThenBB);
  PN->addIncoming(ElseV, ElseBB);
  return PN;
}

Value* ForExprAST::codegen() const {
  Function *TheFunction = Builder.GetInsertBlock()->getParent();
  AllocaInst* Alloca = CreateEntryBlockAlloca(TheFunction, VarName);
  AllocaInst* OldAlloca = NamedValues[VarName];
  NamedValues[VarName] = Alloca;
  
  Value* StartV = _nodes[0]->codegen();
  if (StartV == NULL)
    return NULL;
  Builder.CreateStore(StartV, Alloca);

  BasicBlock* LoopBB = BasicBlock::Create(TheContext, "loop", TheFunction);
  Builder.CreateBr(LoopBB);

  Builder.SetInsertPoint(LoopBB);

  Value* BodyV = _nodes[3]->codegen();
  if (BodyV == NULL)
    return NULL;

  Value* StepV = NULL;
  if (_nodes[2] != NULL)
    StepV = _nodes[2]->codegen();
  else
    StepV = ConstantFP::get(TheContext, APFloat(1.0));
  if (StepV == NULL)
    return NULL;

  Value *curVal = Builder.CreateLoad(Type::getDoubleTy(TheContext), Alloca, VarName.c_str());
  Value *NextVar = Builder.CreateFAdd(curVal, StepV, "nextvar");
  Builder.CreateStore(NextVar, Alloca);
  
  Value *CondV = _nodes[1]->codegen();
  if (CondV == NULL)
    return NULL;
  CondV = Builder.CreateFCmpONE(CondV, ConstantFP::get(TheContext, APFloat(0.0)), "loopcond");

  
  BasicBlock* AfterLoopBB = BasicBlock::Create(TheContext, "afterloop", TheFunction);
  Builder.CreateCondBr(CondV, LoopBB, AfterLoopBB);

  Builder.SetInsertPoint(AfterLoopBB);

  if (OldAlloca != NULL)
    NamedValues[VarName] = OldAlloca;
  else
    NamedValues.erase(VarName);
  
  return ConstantFP::get(TheContext, APFloat(0.0));
}

VarExprAST::~VarExprAST() {
  for (unsigned i = 0; i < V.size(); i++)
    delete V[i].second;
}

Value* VarExprAST::codegen() const {
  vector<AllocaInst*> oldAllocas;
  Function* TheFunction = Builder.GetInsertBlock()->getParent();
  for (unsigned i = 0; i < V.size(); i++) {
    AllocaInst* Alloca = CreateEntryBlockAlloca(TheFunction, V[i].first);
    Value* Tmp = NULL;
    if (V[i].second != NULL)
      Tmp = V[i].second->codegen();
    else
      Tmp = ConstantFP::get(TheContext, APFloat(0.0));
    if (Tmp == NULL)
      return NULL;
    
    oldAllocas.push_back(NamedValues[V[i].first]);
    NamedValues[V[i].first] = Alloca;
    Builder.CreateStore(Tmp, Alloca);
  }

  Value* Res = _nodes[0]->codegen();
  if (Res == NULL)
    return NULL;

  for (unsigned i = 0; i < oldAllocas.size(); i++)
    if (oldAllocas[i] != NULL)
      NamedValues[V[i].first] = oldAllocas[i];
    else
      NamedValues.erase(V[i].first);
  
  return Res;
}


Function* PrototypeAST::codegen() const {
  vector<Type*> Doubles(Args.size(),
                             Type::getDoubleTy(TheContext));
  FunctionType *FT =
    FunctionType::get(Type::getDoubleTy(TheContext), Doubles, false);
  
  Function *F =
    Function::Create(FT, Function::ExternalLinkage, Name, TheModule);
  
  unsigned Idx = 0;
  for (auto &Arg : F->args())
    Arg.setName(Args[Idx++]);
  
  return F;
}

Function* FunctionAST::codegen() const {
  Function *TheFunction = TheModule->getFunction(Proto.getName());

  if (!TheFunction)
    TheFunction = Proto.codegen();

  if (!TheFunction)
    return NULL;

  if (!TheFunction->empty()) {
    cerr << "Nije dozvoljeno predefinisanje fje " << Proto.getName() << endl;
    return NULL;
  }

  // Create a new basic block to start insertion into.
  BasicBlock *BB = BasicBlock::Create(TheContext, "entry", TheFunction);
  Builder.SetInsertPoint(BB);
  
  // Record the function arguments in the NamedValues map.
  NamedValues.clear();
  for (auto &Arg : TheFunction->args()) {
    AllocaInst* Alloca = CreateEntryBlockAlloca(TheFunction, Arg.getName());
    NamedValues[Arg.getName()] = Alloca;
    Builder.CreateStore(&Arg, Alloca);
  }
  
  if (Value *RetVal = Body->codegen()) {
    // Finish off the function.
    Builder.CreateRet(RetVal);
    
    // Validate the generated code, checking for consistency.
    verifyFunction(*TheFunction);

    TheFPM->run(*TheFunction);
    
    return TheFunction;
  }
  TheFunction->eraseFromParent();
  return NULL;
}

void InitializeModuleAndPassManager(void) {
  TheModule = new Module("my_module", TheContext);

  // Create a new pass manager attached to it.
  TheFPM = new FunctionPassManager(TheModule);

  // Do simple "peephole" optimizations and bit-twiddling optzns.
  TheFPM->add(createInstructionCombiningPass());
  // Reassociate expressions.
  TheFPM->add(createReassociatePass());
  // Eliminate Common SubExpressions.
  TheFPM->add(createGVNPass());
  // Simplify the control flow graph (deleting unreachable blocks, etc).
  TheFPM->add(createCFGSimplificationPass());
  TheFPM->add(createPromoteMemoryToRegisterPass());

  TheFPM->doInitialization();
}

AllocaInst *CreateEntryBlockAlloca(Function *TheFunction, const string &VarName) {
  IRBuilder<> TmpB(&TheFunction->getEntryBlock(), TheFunction->getEntryBlock().begin());
  return TmpB.CreateAlloca(Type::getDoubleTy(TheContext), 0, VarName.c_str());
}