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

unique_ptr<Module> TheModule;
LLVMContext TheContext;
map<string, Value*> NamedValues;
IRBuilder<> Builder(TheContext);
unique_ptr<legacy::FunctionPassManager> TheFPM;
unique_ptr<KaleidoscopeJIT> TheJIT;
map<string, PrototypeAST> FunctionProtos;

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

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

Value* BinaryExprAST::codegen() const {
  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");
  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 {
  Value* startV = _start->codegen();
  if (!startV)
    return nullptr;

  Function* TheFunction = Builder.GetInsertBlock()->getParent();
  
  BasicBlock *loopBB = BasicBlock::Create(TheContext, "loop", TheFunction);
  BasicBlock *beforeLoopBB = Builder.GetInsertBlock();

  Builder.CreateBr(loopBB);

  Builder.SetInsertPoint(loopBB);
  PHINode *Variable = Builder.CreatePHI(Type::getDoubleTy(TheContext), 2, _varName.c_str());
  Variable->addIncoming(startV, beforeLoopBB);

  Value* oldValue = NamedValues[_varName];
  NamedValues[_varName] = startV;

  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* nextVar = Builder.CreateFAdd(Variable, stepV, "nextvar");

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

  BasicBlock *afterLoopBB = BasicBlock::Create(TheContext, "afterloop", TheFunction);
  loopBB = Builder.GetInsertBlock();
  Builder.CreateCondBr(stopV, loopBB, afterLoopBB);
  Variable->addIncoming(nextVar, loopBB);
  
  Builder.SetInsertPoint(afterLoopBB);

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

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