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

/* trenutni modul */
unique_ptr<Module> theModule;
/* kontekst */
LLVMContext theContext;
/* mapa koja preslikava lokalne promenljive u njihove alokatore */
map<string, AllocaInst*> namedValues;
/* builder za pravljenje instrukcija */
IRBuilder<> builder(theContext);
/* optimizacija funkcija */
unique_ptr<legacy::FunctionPassManager> theFPM;
/* mapa koja preslikava imena funkcija u njihove potpise */
map<string, PrototypeAST> functionProtos;

/* codegen() metodi za sve klase */

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

Value* VariableExprAST::codegen() const {
  AllocaInst *v = namedValues[_name];
  if (!v) {
    cerr << "Nepostojeca promenljiva " << _name << endl;
    return nullptr;
  }
  return builder.CreateLoad(v);
}

Value* BinaryExprAST::codegen() const {
  if (_op == '=') {
    Value *r = _rhs->codegen();
    /* ako je dynamic_cast podrazano
    VariableExprAST* lhs = dynamic_cast<VariableExprAST*>(_lhs);
    if (!lhs) {
      cerr << "Desna strana dodele mora da bude promenljiva" << endl;
      return nullptr;
    }
    */
    VariableExprAST* lhs = (VariableExprAST*)_lhs;
    string name = lhs->getName();
    AllocaInst* alloca = namedValues[name];
    return builder.CreateStore(r, alloca);
  }
  Value *l = _lhs->codegen();
  Value *r = _rhs->codegen();
  if (!l || !r)
    return nullptr;
  switch(_op) {
  case '+':
    return builder.CreateFAdd(l, r, "tmpadd");
  case '*':
    return builder.CreateFMul(l, r, "tmpmul");
  case '-':
    return builder.CreateFSub(l, r, "tmpsub");
  case '<':
    l = builder.CreateFCmpULT(l, r, "tmpu1");
    return builder.CreateUIToFP(l, Type::getDoubleTy(theContext), "tmpcmp");
  case ':':
    return r;
  default:
    return nullptr;
  }
}

Value* CallExprAST::codegen() const {
  Function *f = getFunction(_name);
  if (!f) {
    cerr << "Funkcija " << _name << " nije definisana" << endl;
    return nullptr;
  }
  if (f->arg_size() != _exps.size()) {
    cerr << "Poziv fje " << _name << "nema " << f->arg_size() << "argumenatar" << endl;
    return nullptr;
  }
  vector<Value*> argV;
  for (auto e : _exps) {
    argV.push_back(e->codegen());
    if (!argV.back())
      return nullptr;
  }
  return builder.CreateCall(f, argV, "calltmp");
}

Value* IfExprAST::codegen() const {
  Value* condV = _cond->codegen();
  if (!condV)
    return nullptr;
  
  Value* tmp = builder.CreateFCmpONE(condV, ConstantFP::get(theContext, APFloat(0.0)), "ifcond");

  Function* theFunction = builder.GetInsertBlock()->getParent();

  BasicBlock *thenBB = BasicBlock::Create(theContext, "then", theFunction);
  BasicBlock *elseBB = BasicBlock::Create(theContext, "else");
  BasicBlock *mergeBB = BasicBlock::Create(theContext, "ifcont");

  builder.CreateCondBr(tmp, thenBB, elseBB);

  builder.SetInsertPoint(thenBB);
  Value* thenV = _then->codegen();
  if (!thenV)
    return nullptr;
  builder.CreateBr(mergeBB);
  thenBB = builder.GetInsertBlock();

  theFunction->getBasicBlockList().push_back(elseBB);
  builder.SetInsertPoint(elseBB);
  Value* elseV = _else->codegen();
  if (!elseV)
    return nullptr;
  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);
  
  BasicBlock *loopBB = BasicBlock::Create(theContext, "loop", theFunction);

  Value* startV = _start->codegen();
  if (!startV)
    return nullptr;
  builder.CreateStore(startV, alloca);
  builder.CreateBr(loopBB);

  builder.SetInsertPoint(loopBB);
  
  AllocaInst* oldValue = namedValues[_varName];
  namedValues[_varName] = alloca;

  Value* stopV = _stop->codegen();
  if (!stopV)
    return nullptr;
  stopV = builder.CreateFCmpONE(stopV, ConstantFP::get(theContext, APFloat(0.0)), "loopcond");

  BasicBlock *loopBB1 = BasicBlock::Create(theContext, "loop1", theFunction);
  BasicBlock *afterLoopBB = BasicBlock::Create(theContext, "afterloop", theFunction);
  builder.CreateCondBr(stopV, loopBB1, afterLoopBB);

  builder.SetInsertPoint(loopBB1);
  
  if (!_body->codegen())
    return nullptr;

  Value* stepV = nullptr;
  if (_step) {
    stepV = _step->codegen();
    if (!stepV)
      return nullptr;
  }
  else
    stepV = ConstantFP::get(theContext, APFloat(1.0));

  Value* curVal = builder.CreateLoad(alloca);
  Value* nextVar = builder.CreateFAdd(curVal, stepV, "nextvar");
  builder.CreateStore(nextVar, alloca);
  builder.CreateBr(loopBB);
  
  builder.SetInsertPoint(afterLoopBB);

  if (oldValue)
    namedValues[_varName] = oldValue;
  else
    namedValues.erase(_varName);
  
  return ConstantFP::get(theContext, APFloat(0.0));
}

Value* VarExprAST::codegen() const {
  Function* theFunction = builder.GetInsertBlock()->getParent();
  vector<AllocaInst*> oldAllocas;
  for (unsigned i = 0; i < _list.size(); i++) {
    string name = _list[i].first;
    AllocaInst* alloca = createEntryBlockAlloca(theFunction, name);
    oldAllocas.push_back(namedValues[name]);
    namedValues[name] = alloca;
    Value *startV;
    if (_list[i].second)
      startV = _list[i].second->codegen();
    else
      startV = ConstantFP::get(theContext, APFloat(0.0));
    if (!startV)
      return nullptr;
    builder.CreateStore(startV, alloca);
  }

  Value* bodyV = _body->codegen();
  if (!bodyV)
    return nullptr;

  for (unsigned i = 0; i < _list.size(); i++) {
    if (oldAllocas[i])
      namedValues[_list[i].first] = oldAllocas[i];
    else
      namedValues.erase(_list[i].first);
  }
  
  return bodyV;
}

Function* PrototypeAST::codegen() const {
  vector<Type*> v(_args.size(), Type::getDoubleTy(theContext));
  FunctionType* fp = FunctionType::get(Type::getDoubleTy(theContext), v, false);
  Function* f = Function::Create(fp, Function::ExternalLinkage, _name, theModule.get());
  int j = 0;
  for (auto &a : f->args())
    a.setName(_args[j++]);
  return f;
}

Function* FunctionAST::codegen() const {
  functionProtos.insert(pair<string, PrototypeAST>(_proto.getName(),_proto));
  Function* theFunction = getFunction(_proto.getName());
  if (!theFunction)
    theFunction = (Function*)_proto.codegen();
  if (!theFunction)
    return nullptr;
  BasicBlock *bb = BasicBlock::Create(theContext, "entry", theFunction);
  builder.SetInsertPoint(bb);
  namedValues.clear();
  for (auto &a: theFunction->args()) {
    //namedValues[a.getName()] = &a;
    AllocaInst* alloca = createEntryBlockAlloca(theFunction, a.getName());
    namedValues[a.getName()] = alloca;
    builder.CreateStore(&a, alloca);
  }
  if (Value *ret = _definition->codegen()) {
    builder.CreateRet(ret);
    verifyFunction(*theFunction);
    theFPM->run(*theFunction);
    return theFunction;
  }
  else {
    theFunction->eraseFromParent();
    return nullptr;
  }
}

void initializeModuleAndPassManager() {
  theModule = make_unique<Module>("my module", theContext);
  theFPM = make_unique<legacy::FunctionPassManager>(theModule.get());
  //theFPM->add(createInstructionCombiningPass());
  //theFPM->add(createReassociatePass());
  //theFPM->add(createGVNPass());
  //theFPM->add(createCFGSimplificationPass());
  theFPM->doInitialization();
}

Function* getFunction(string s) {
  Function *f = theModule->getFunction(s);
  if (f)
    return f;
  auto i = functionProtos.find(s);
  if (i != functionProtos.end())
    return i->second.codegen();
  return nullptr;
}

AllocaInst* createEntryBlockAlloca(Function* theFunction, const string& name) {
  IRBuilder<> b(&theFunction->getEntryBlock(), theFunction->getEntryBlock().begin());
  return b.CreateAlloca(Type::getDoubleTy(theContext), 0, name.c_str());
}