#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");
  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");
}

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