#include "SymbolicInterpreter.h"

#include <ctime>
#include <fstream>

std::ostream &operator<<(std::ostream &OS, const Address &A) {
  if (A.Type == A.Memory) {
    OS << A.Addr;
  } else {
    OS << "R" << A.Addr;
  }
  return OS;
}

int SymbolicInterpreter::NewInput(int *Ptr, int ID) {
  int Ret = 0;
  if (Inputs.find(ID) != Inputs.end()) {
    Ret = Inputs[ID];
  } else {
    Ret = std::rand();
    Inputs[ID] = Ret;
  }
  Address X(Ptr);
  std::string InputName = "X" + std::to_string(NumOfInputs);
  z3::expr SE = Ctx.int_const(InputName.c_str());
  /// Type(Mem) : MemoryTy = std::map<Address, z3::expr>
  Mem.insert(std::make_pair(X, SE));
  NumOfInputs++;
  return Ret;
}

SymbolicInterpreter SI;

void print(std::ostream &OS) {
  OS << "=== Inputs ===" << std::endl;
  int Idx = 0;
  for (auto &E : SI.getInputs()) {
    OS << "X" << E.first << " : " << E.second << std::endl;
  }
  OS << std::endl;
  OS << "=== Symbolic Memory ===" << std::endl;
  for (auto &E : SI.getMemory()) {
    OS << E.first << " : " << E.second << std::endl;
  }
  OS << std::endl;

  OS << "=== Path Condition ===" << std::endl;
  for (auto &E : SI.getPathCondition()) {
    std::string BID = "B" + std::to_string(E.first);
    OS << BID << " : " << E.second << std::endl;
  }
}

extern "C" void __DSE_Exit__() {
  z3::solver Solver(SI.getContext());
  std::ofstream Branch(BranchFile);
  for (auto &E : SI.getPathCondition()) {
    std::string BID = "B" + std::to_string(E.first);
    Solver.add(E.second);
    Branch << E.second << "\n";
  }
  std::ofstream Smt2(FormulaFile);
  Smt2 << Solver.to_smt2();
  std::ofstream Log(LogFile);
  print(Log);
}

extern "C" void __DSE_Init__() {
  std::srand(std::time(nullptr));
  std::string Line;
  std::ifstream Input(InputFile);
  if (Input.is_open()) {
    while (getline(Input, Line)) {
      int ID = std::stoi(Line.substr(1, Line.find(",")));
      int Val = std::stoi(Line.substr(Line.find(",") + 1));
      SI.getInputs()[ID] = Val;
    }
  }
  std::atexit(__DSE_Exit__);
}

extern "C" void __DSE_Input__(int *X, int ID) {
    /// 地址 *X 的值变为 (int)SI.NewInput(X, ID)
    /// Inputs<0, V0> => %x = V0
    *X = (int)SI.NewInput(X, ID);
}

/// 这里我按照自己的逻辑改了框架
extern "C" void __DSE_Branch__(int RID, int BID) {
  if(SI.getStack().top().to_string() == "0"){
      SI.getStack().pop();
      return;
  }
  MemoryTy &Mem = SI.getMemory();
  Address Addr(RID);
  /// (Addr(RID), SE)
  z3::expr SE = Mem.at(Addr);
  SI.getPathCondition().push_back(std::make_pair(BID, SE));
}

extern "C" void __DSE_Const__(int X) {
  z3::expr SE = SI.getContext().int_val(X);
  SI.getStack().push(SE);
}

extern "C" void __DSE_Register__(int X) {
  std::string Name = "R" + std::to_string(X);
  z3::expr SE = SI.getContext().int_const(Name.c_str());
  SI.getStack().push(SE);
}
