#include "LoopUtils.h"
#include "llvm/Analysis/DomTreeUpdater.h"
#include "llvm/IR/BasicBlock.h"
#include "llvm/IR/Constants.h"
#include "llvm/IR/Dominators.h"
#include "llvm/IR/GlobalValue.h"
#include "llvm/IR/IRBuilder.h"
#include "llvm/IR/Module.h"
#define DEBUG_TYPE "runtime-safety"

// STARTING - 0
// LOOPING - 1
// CACHING - 2
// INVALID - 3

class RuntimeInsertion {

  Module *M;
  Loop *L;
  GlobalVariable *GV;
  ModuleAnalysisManager &AM;
  FunctionAnalysisManager *FM;

public:
  RuntimeInsertion(Module *M, Loop *L, GlobalVariable *GV,
                   ModuleAnalysisManager &AM)
      : M(M), L(L), GV(GV), AM(AM) {}

  void setFunctionAnalysismanager() {
    FM = &AM.getResult<FunctionAnalysisManagerModuleProxy>(*M).getManager();
  }

  void addSafetyChecksToLoopRegion() {
    /* get Loop information and Dominator Tree information */
    Function *F = L->getHeader()->getParent();
    setFunctionAnalysismanager();
    LoopInfo *LI = &FM->getResult<LoopAnalysis>(*F);
    DominatorTree *DT = &FM->getResult<DominatorTreeAnalysis>(*F);
    addSafetyChecksToLoopRegionStart(LI, DT);
    addSafetyChecksToLoopRegionEnd(LI, DT);
    DT->recalculate(*F);
  }

  void addSafetyChecksToLoopRegionStart(LoopInfo *LI, DominatorTree *DT) {

    LLVMContext &Cntx = M->getContext();
    Type *I32Ty = Type::getInt32Ty(Cntx);
    Constant *i32_val1 = ConstantInt::get(I32Ty, 1, true);
    Constant *i32_val2 = ConstantInt::get(I32Ty, 2, true);
    Constant *i32_val3 = ConstantInt::get(I32Ty, 3, true);

    /* Get the preheader block of the loop.  */
    BasicBlock *preHeader = L->getLoopPreheader();

    /* If no preheader then create a dedicated preheader */
    if (!preHeader)
      preHeader = InsertPreheaderForLoop(L, DT, LI, nullptr, false);
    assert(preHeader != nullptr);

    /* Create the condition check for statevariable
      (StateVar  == STARTING || StateVar == LOOPING) */
    IRBuilder<> B(preHeader->getTerminator());
    LoadInst *Load = B.CreateLoad(I32Ty, GV);
    Value *cmpInst = B.CreateICmp(ICmpInst::ICMP_ULT, Load, i32_val2);

    /* Insert the check in preheader and split the block to if-then-else
       old preheader:
         statevar check
         cmp
         Branch equal else:
       then:
         statevar = LOOPING
       else:
         statevar = INVALID
         clearcache
     merge block:  */

    Instruction *ThenBrInst = nullptr, *ElseBrInst = nullptr;
    SplitBlockAndInsertIfThenElse(cmpInst, preHeader->getTerminator(),
                                  &ThenBrInst, &ElseBrInst);

    /* Add statevariable changes to If block */
    B.SetInsertPoint(ThenBrInst);
    B.CreateStore(i32_val1, GV);

    /* Add statevariable changes to Else block */
    B.SetInsertPoint(ElseBrInst);
    B.CreateStore(i32_val3, GV);

    /* Add call to clearing cache in Else block */
  }

  void addSafetyChecksToLoopRegionEnd(LoopInfo *LI, DominatorTree *DT) {
    LLVMContext &Cntx = M->getContext();
    Type *I32Ty = Type::getInt32Ty(Cntx);
    Constant *i32_valm2 = ConstantInt::get(I32Ty, -2, true);
    Constant *i32_val2 = ConstantInt::get(I32Ty, 2, true);
    Constant *i32_val3 = ConstantInt::get(I32Ty, 3, true);


    /* Check dedicated exit blocks if not create them */
    if (!L->hasDedicatedExits())
      formDedicatedExitBlocks(L, DT, LI, nullptr, true);

    SmallVector<BasicBlock *, 4> Exits;
    L->getUniqueExitBlocks(Exits);
    for (auto *ExitBB : Exits) {
      LLVM_DEBUG(dbgs() << "ExitBBs\n");
      IRBuilder<> B(ExitBB->getTerminator());
      LoadInst *Load = B.CreateLoad(I32Ty, GV);
      Value *And = B.CreateAnd(Load, i32_valm2);
      Value *cmpInst = B.CreateICmp(ICmpInst::ICMP_EQ, Load, i32_val2);

      /* In each exit blocks create condition for state variables
       exit block:
         statevar check
         cmp
         Branch equal else:
       then:
         statevar = LOOPING
       else:
         statevar = INVALID
         clearcache
       exit block merge:  */
      Instruction *ThenBrInst = SplitBlockAndInsertIfThen(
          cmpInst, ExitBB->getTerminator(), false);
      /* Add statevariable changes to If block */
      B.SetInsertPoint(ThenBrInst);
      B.CreateStore(i32_val3, GV);
      /// TODO: Connect this call with adv-mem-opt pass
      // Clear the cache by calling aocc_clear_cache
      Type *ResultTy = Type::getVoidTy(Cntx);
      FunctionType *FnType = FunctionType::get(ResultTy, {}, false);
      FunctionCallee LibcallFn =
          M->getOrInsertFunction("aocc_clear_cache", FnType);
      CallInst *CacheClearCall = B.CreateCall(LibcallFn);
      /* Add call to clearing cache in Else block */
    }
  }

  Instruction *InsertBlockBetweenPredecessors(BasicBlock *BB,
                                              BasicBlock *PredBB,
                                              DomTreeUpdater &DTU) {

    BasicBlock *NewBB = SplitEdge(PredBB, BB);
    NewBB->setName(PredBB->getName() + ".split");
    Instruction *NewTerm = NewBB->getTerminator();

    // FIXME: SplitEdge does not yet take a DTU, so we include the split edge
    //        in the update set here.
    DTU.applyUpdates({{DominatorTree::Delete, PredBB, BB},
                      {DominatorTree::Insert, PredBB, NewBB},
                      {DominatorTree::Insert, NewBB, BB}});
    return NewTerm;
  }

  void addSafetyChecksToUnsafeBlocks(std::set<BasicBlock *> UnsafeBlocks) {
    LLVMContext &Cntx = M->getContext();
    Type *I32Ty = Type::getInt32Ty(Cntx);
    Constant *i32_val2 = ConstantInt::get(I32Ty, 2, true);
    Constant *i32_val3 = ConstantInt::get(I32Ty, 3, true);
    setFunctionAnalysismanager();

    /* For unsafe blocks region where we are not able to analyze due to indirect
       calls or failed devirtualization we will split the edge between the block
       and its predecessors. we will insert if-then block which checks state is
       caching and if so make it invalid */
    for (BasicBlock *BB : UnsafeBlocks) {
      Function *F = BB->getParent();
      DominatorTree *DT = &FM->getResult<DominatorTreeAnalysis>(*F);
      LoopInfo *LI = &FM->getResult<LoopAnalysis>(*F);
      std::optional<DomTreeUpdater> DTU;
      if (DT)
        DTU.emplace(DT, DomTreeUpdater::UpdateStrategy::Lazy);

      Instruction *newBBTerm = BB->getFirstNonPHI();
      IRBuilder<> B(BB->getFirstNonPHI());
      B.SetInsertPoint(BB->getFirstNonPHI());
      LoadInst *Load = B.CreateLoad(I32Ty, GV);
      Value *cmpInst = B.CreateICmp(ICmpInst::ICMP_EQ, Load, i32_val2);
      Instruction *ThenBrInst = SplitBlockAndInsertIfThen(
          cmpInst, newBBTerm, false, nullptr, &(*DTU), LI);
      B.SetInsertPoint(ThenBrInst);
      B.CreateStore(i32_val3, GV);
    }
  }

  void
  addSafetyChecksToUnsafeExceptionBlocks(std::set<BasicBlock *> UnsafeBlocks) {
    LLVMContext &Cntx = M->getContext();
    Type *I32Ty = Type::getInt32Ty(Cntx);
    Constant *i32_val3 = ConstantInt::get(I32Ty, 3, true);
    setFunctionAnalysismanager();

    /* For unsafe exception blocks  we will split the edge between the block and
       its predecessors. we will insert state as invalid in the added new block
     */
    for (BasicBlock *BB : UnsafeBlocks) {
      Function *F = BB->getParent();
      DominatorTree *DT = &FM->getResult<DominatorTreeAnalysis>(*F);
      std::optional<DomTreeUpdater> DTU;
      if (DT)
        DTU.emplace(DT, DomTreeUpdater::UpdateStrategy::Lazy);

      if (BB->isLandingPad()) {
        StoreInst *St = new StoreInst(i32_val3, GV, BB->getTerminator());
        continue;
      }

      for (auto it = pred_begin(BB), et = pred_end(BB); it != et; ++it) {
        BasicBlock *predecessor = *it;
        Instruction *newBBTerm =
            InsertBlockBetweenPredecessors(BB, predecessor, *DTU);
        IRBuilder<> B(newBBTerm);
        B.SetInsertPoint(newBBTerm);
        B.CreateStore(i32_val3, GV);
      }
    }
  }

  Function *
  specializeCandidateFunctionAndAddSafetyChecks(Function *candidateFunction,
                                                ValueToValueMapTy *FCloneVMap) {
    LLVMContext &Cntx = M->getContext();
    Type *I32Ty = Type::getInt32Ty(Cntx);
    Constant *i32_valm1 = ConstantInt::get(I32Ty, -1, true);
    Constant *i32_val2 = ConstantInt::get(I32Ty, 2, true);
    std::vector<CallBase *> CallSites;
    Function *clonedFunction = nullptr;
    if (candidateFunction) {
      // Clone the candidate function.
      clonedFunction = CloneFunction(candidateFunction, *FCloneVMap);
      clonedFunction->setName(clonedFunction->getName() + ".iacaching.clone");
      clonedFunction->setLinkage(candidateFunction->getLinkage());

      // Get all the call site for the candidate function.
      for (User *U : candidateFunction->users()) {
        LLVM_DEBUG(dbgs() << "User" << *U << "\n");
        if (CallBase *CI = dyn_cast<CallBase>(U)) {
          CallSites.push_back(CI);
        }
      }

      // For each call sites of the candidate function.
      for (CallBase *CI : CallSites) {
        /* Create statevariable check and insert if then else checks.
           Create Phi nodes for the return value and update the uses
           of return value with new phi node.
           from:
            ret = foo();
           to:
            if (stateVar=LOOPING || stateVar=CACHING)
              ret1 = foo_special();
            else
              ret2 =  foo();
            ret3 = phi(ret1, ret2);
            ret = foo() <== Orginal call will be deleted.
        */
        IRBuilder<> B(CI);
        B.SetInsertPoint(CI);
        LoadInst *Load = B.CreateLoad(I32Ty, GV);
        Value *Add = B.CreateAdd(Load, i32_valm1);
        Value *cmpInst = B.CreateICmp(ICmpInst::ICMP_ULT, Add, i32_val2);

        Instruction *ThenBrInst = nullptr, *ElseBrInst = nullptr;
        SplitBlockAndInsertIfThenElse(cmpInst, CI, &ThenBrInst, &ElseBrInst);

        CallInst *thenCall = cast<CallInst>(dyn_cast<Instruction>(CI)->clone());
        thenCall->insertBefore(ThenBrInst);
        CallInst *elseCall = cast<CallInst>(dyn_cast<Instruction>(CI)->clone());
        elseCall->insertBefore(ElseBrInst);
        thenCall->setCalledFunction(clonedFunction);

        B.SetInsertPoint(CI);
        PHINode *phi = B.CreatePHI(CI->getType(), 2);
        phi->addIncoming(thenCall, ThenBrInst->getParent());
        phi->addIncoming(elseCall, ElseBrInst->getParent());

        CI->replaceAllUsesWith(phi);
        CI->eraseFromParent();
      }
    }
    return clonedFunction;
  }

void clearCacheAtDestructors(const std::vector<Function *> &Destructors){
  LLVMContext &Cntx = M->getContext();
  Type *I32Ty = Type::getInt32Ty(Cntx);
  Constant *i32_val3 = ConstantInt::get(I32Ty, 3, true);
  for (auto * fn: Destructors){
    IRBuilder<> B(fn->getEntryBlock().getFirstNonPHI());
    B.SetInsertPoint(fn->getEntryBlock().getFirstNonPHI());
    B.CreateStore(i32_val3, GV);
    /// TODO: Connect this call with adv-mem-opt pass
    // Clear the cache by calling aocc_clear_cache
    Type *ResultTy = Type::getVoidTy(Cntx);
    FunctionType *FnType = FunctionType::get(ResultTy, {}, false);
    FunctionCallee LibcallFn =
    M->getOrInsertFunction("aocc_clear_cache", FnType);
    CallInst *CacheClearCall = B.CreateCall(LibcallFn);   
  }
}
};
