import z3.scala._

object ScalaZ3Interface {

  def z3Encode(z3: Z3Context, solver: Z3Solver, z3Variables: Map[Identifier, Z3AST], e: Expr): Z3AST = e match {
    case Var(id) => z3Variables(id)
    case BooleanLiteral(true) => z3.mkTrue()
    case BooleanLiteral(false) => z3.mkFalse()
    case Equal(e1, e2) =>
      val e1AST = z3Encode(z3, solver, z3Variables, e1)
      val e2AST = z3Encode(z3, solver, z3Variables, e2)
      z3.mkEq(e1AST, e2AST)
    case Implies(e1, e2) =>
      val e1AST = z3Encode(z3, solver, z3Variables, e1)
      val e2AST = z3Encode(z3, solver, z3Variables, e2)
      z3.mkImplies(e1AST, e2AST)
    case And(e1, e2) =>
      val e1AST = z3Encode(z3, solver, z3Variables, e1)
      val e2AST = z3Encode(z3, solver, z3Variables, e2)
      z3.mkAnd(e1AST, e2AST)
    case Or(e1, e2) =>
      val e1AST = z3Encode(z3, solver, z3Variables, e1)
      val e2AST = z3Encode(z3, solver, z3Variables, e2)
      z3.mkOr(e1AST, e2AST)
    case Xor(e1, e2) =>
      val e1AST = z3Encode(z3, solver, z3Variables, e1)
      val e2AST = z3Encode(z3, solver, z3Variables, e2)
      z3.mkXor(e1AST, e2AST)
    case Not(e0) =>
      val e0AST = z3Encode(z3, solver, z3Variables, e0)
      z3.mkNot(e0AST)
  }

  def checkSat(e: Expr): Option[Map[Identifier, Boolean]] = {
    val factory: Solver = Solver.getFactory
    val z3: Z3Context = factory.getContext
    val solver: Z3Solver = z3.mkSolver
    val boolSort = z3.mkBoolSort

    val z3Variables = e.fv.map(id => id -> z3.mkConst(z3.mkStringSymbol(id.toString), boolSort)).toMap
    val z3Formula = z3Encode(z3, solver, z3Variables, e)

    solver.assertCnstr(z3Formula)

    val solverResponse = solver.check

    val res: Option[Map[Identifier, Boolean]] = solverResponse match {
      case None => None
      case Some(true) =>
        val model = solver.getModel
        Some(z3Variables.map { case (v, z3v) => v -> model.evalAs[Boolean](z3v).getOrElse(false) })
      case Some(false) => None
    }

    res
  }
}
