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

LLVMContext TheContext;
Module *TheModule;
IRBuilder<> Builder(TheContext);
map<string, Value*> NamedValues;
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;
}

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

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

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

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, "lesstmp");
  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, "greatertmp");
  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 << "Funkcija " << _f << " nije deklarisana" << endl;
    exit(EXIT_FAILURE);
  }

  vector<Value*> a;
  for (unsigned i = 0; i < _v.size(); i++) {
    a.push_back(_v[i]->codegen());
    if (!a.back())
      return NULL;
  }
  return Builder.CreateCall(F, a, "calltmp");
}

Value* IfThenElseAST::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);
  Builder.CreateBr(LoopBB);
  BasicBlock* BeforeLoopBB = Builder.GetInsertBlock();

  Builder.SetInsertPoint(LoopBB);
  PHINode* Variable = Builder.CreatePHI(Type::getDoubleTy(TheContext), 2, "i");
  Variable->addIncoming(StartVal, BeforeLoopBB);

  Value *OldVal = NamedValues[_s];
  NamedValues[_s] = Variable;

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

  Value *IncVal = _v[2]->codegen();
  if (!IncVal)
    return NULL;
  Value *NextVal = Builder.CreateFAdd(Variable, IncVal, "nextvar");

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

  Value* Tmp = Builder.CreateFCmpONE(CondVal, ConstantFP::get(TheContext, APFloat(0.0)), "loopcond");
  BasicBlock* AfterLoopBB = BasicBlock::Create(TheContext, "afterloop", F);
  Builder.CreateCondBr(Tmp, LoopBB, AfterLoopBB);
  
  LoopBB = Builder.GetInsertBlock();
  Variable->addIncoming(NextVal, LoopBB);

  Builder.SetInsertPoint(AfterLoopBB);

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

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

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

  return F;
}

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

  if (!F) {
    cerr << "Nemoguce je generisati funkciju " << _p.getName() << endl;
    exit(EXIT_FAILURE);
  }

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

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

  NamedValues.clear();
  for (auto &Arg : F->args())
    NamedValues[Arg.getName()] = &Arg;

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

    verifyFunction(*F);

    //TheFPM->run(*F);
    
    return F;
  }
  else {
    F->eraseFromParent();
    return NULL;
  }
}

void InitializeModuleAndPassManager() {
  TheModule = new Module("Moj modul", TheContext);
  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->doInitialization();
}