#include "ast.hpp"
#include <iostream>

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

ExprAST::~ExprAST() {

}

InnerExprAST::InnerExprAST(ExprAST* e1) {
  _v.resize(0);
  _v.push_back(e1);
}

InnerExprAST::InnerExprAST(ExprAST* e1, ExprAST* e2) {
  _v.resize(0);
  _v.push_back(e1);
  _v.push_back(e2);
}

InnerExprAST::InnerExprAST(ExprAST* e1, ExprAST* e2, ExprAST* e3) {
  _v.resize(0);
  _v.push_back(e1);
  _v.push_back(e2);
  _v.push_back(e3);
}

InnerExprAST::InnerExprAST(ExprAST* e1, ExprAST* e2, ExprAST* e3, ExprAST* e4) {
  _v.resize(0);
  _v.push_back(e1);
  _v.push_back(e2);
  _v.push_back(e3);
  _v.push_back(e4);
}

InnerExprAST::~InnerExprAST() {
  for (auto e : _v)
    delete e;
}

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

Value* VariableExprAST::codegen() const {
  AllocaInst *Alloca = NamedValues[_v];
  if (!Alloca) {
    cerr << "Nepostojeca promenljiva " << _v << endl;
    exit(EXIT_FAILURE);
  }
  return Builder.CreateLoad(Alloca); 
}

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

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

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

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

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

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

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

Value* CallExprAST::codegen() const {
  Function* F = TheModule->getFunction(_f);
  if (!F) {
    cerr << "Poziv funkcije " << _f << " koja nije definisana";
    exit(EXIT_FAILURE);
  }

  if (F->arg_size() != _v.size()) {
    cerr << "Funkcija " << _f << " ocekuje " << F->arg_size() << " argumenata" ;
    exit(EXIT_FAILURE);
  }

  vector<Value*> a;
  for (unsigned i = 0; i < _v.size(); i++)
    a.push_back(_v[i]->codegen());

  return Builder.CreateCall(F, a, "calltmp");
}

Value* IfThenElse::codegen() const {
  Value *CondV = _v[0]->codegen();
  if (!CondV)
    return NULL;
  Value* Tmp = Builder.CreateFCmpONE(CondV, ConstantFP::get(TheContext, APFloat(0.0)), "ifcond");

  Function *F = Builder.GetInsertBlock()->getParent();
  BasicBlock *ThenBB = BasicBlock::Create(TheContext, "then", F);
  BasicBlock *ElseBB = BasicBlock::Create(TheContext, "else");
  BasicBlock *MergeBB = BasicBlock::Create(TheContext, "ifcont");

  Builder.CreateCondBr(Tmp, ThenBB, ElseBB);

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

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

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

  return PHI;
}

Value* ForIn::codegen() const {
  Value* StartVal = _v[0]->codegen();
  if (!StartVal)
    return NULL;

  Function *F = Builder.GetInsertBlock()->getParent();
  BasicBlock *LoopBB = BasicBlock::Create(TheContext, "loop", F);

  AllocaInst* Alloca = CreateEntryBlockAlloca(F, _s);
  Builder.CreateStore(StartVal, Alloca);
  Builder.CreateBr(LoopBB);

  Builder.SetInsertPoint(LoopBB);
  AllocaInst* OldVal = NamedValues[_s];
  NamedValues[_s] = Alloca;
  
  Value* BodyVal = _v[3]->codegen();
  if (!BodyVal)
    return NULL;

  Value* IncVal = _v[2]->codegen();
  if (!IncVal)
    return NULL;
  Value *CurrVal = Builder.CreateLoad(Alloca);
  Value* NextVar = Builder.CreateFAdd(CurrVal, IncVal, "nextvar");
  Builder.CreateStore(NextVar, Alloca);

  Value* CondV = _v[1]->codegen();
  if (!CondV)
    return NULL;

  Value* Tmp = Builder.CreateFCmpONE(CondV, ConstantFP::get(TheContext, APFloat(0.0)), "loopcond");

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

  Builder.SetInsertPoint(AfterLoopBB);

  LoopBB = Builder.GetInsertBlock();
  
  if (OldVal != NULL)
    NamedValues[_s] = OldVal;
  else
    NamedValues.erase(_s);
  return ConstantFP::get(TheContext, APFloat(0.0));
}

Value* AssignExprAST::codegen() const {
  Value *Val = _v[0]->codegen();
  if (!Val)
    return NULL;
  AllocaInst* Alloca = NamedValues[_s];
  if (!Alloca) {
    cerr << "Promenljiva " << _s << " ne postoji" << endl;
    exit(EXIT_FAILURE);
  }
  Builder.CreateStore(Val, Alloca);
  return Val;
}

Value* VarExprAST::codegen() const {
  Function *F = Builder.GetInsertBlock()->getParent();
  vector<AllocaInst*> OldValues;
  
  for (unsigned i = 0; i < _vars.size(); i++) {
    string VarName = _vars[i].first;
    AllocaInst* Alloca = CreateEntryBlockAlloca(F, VarName);
    Value *InitVal = _vars[i].second->codegen();
    if (!InitVal)
      return NULL;
    Builder.CreateStore(InitVal, Alloca);

    OldValues.push_back(NamedValues[VarName]);
    NamedValues[VarName] = Alloca;
  }

  Value* Val = _v[0]->codegen();
  if (!Val)
    return NULL;

  for (unsigned i = 0; i < _vars.size(); i++) {
    string VarName = _vars[i].first;
    if (OldValues[i])
      NamedValues[VarName] = OldValues[i];
    else
      NamedValues.erase(VarName);
  }

  return Val;
}

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

Function* PrototypeAST::codegen() const {
  vector<Type*> d(_args.size(), Type::getDoubleTy(TheContext));
  FunctionType* FT = FunctionType::get(Type::getDoubleTy(TheContext), d, false);
  Function *F = Function::Create(FT, Function::ExternalLinkage, _f, TheModule);

  unsigned i = 0;
  for (auto &Arg : F->args())
    Arg.setName(_args[i++]);

  return F;
}

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

Function* FunctionAST::codegen() const {
  Function* F = TheModule->getFunction(_proto.getName());
  if (!F)
    F = _proto.codegen();

  if (!F) {
    cerr << "Nemoguce generisanje koda za funkciju " << _proto.getName() << endl;
    exit(EXIT_FAILURE);
  }

  if (!F->empty()) {
    cerr << "Nemoguce je predefinisati fju " << _proto.getName() << endl;
    exit(EXIT_FAILURE);
  }

  BasicBlock *BB = BasicBlock::Create(TheContext, "entry", F);
  Builder.SetInsertPoint(BB);

  NamedValues.clear();
  for (auto &Arg : F->args()) {
    AllocaInst* Alloca = CreateEntryBlockAlloca(F, Arg.getName());
    NamedValues[Arg.getName()] = Alloca;
    Builder.CreateStore(&Arg, Alloca);
  }

  Value* RetVal;
  if ((RetVal = _e->codegen())) {
    Builder.CreateRet(RetVal);

    verifyFunction(*F);

    //TheFPM->run(*F);

    return F;
  }
  else {
    F->eraseFromParent();
    return NULL;
  }
}

void InitializeModuleAndPassManager() {
    TheModule = new llvm::Module("Moj Modul", TheContext);
    TheFPM = new llvm::legacy::FunctionPassManager(TheModule);

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

    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());
}