module SMT.Unsafe.Z3 where

open import Agda.Primitive using (Set)
open import Agda.Builtin.Bool using (Bool; true; false)
open import Agda.Builtin.List using (List; []; _∷_)
open import Agda.Builtin.Nat using (Nat; zero; suc)
import Agda.Builtin.Reflection as R
open import Agda.Builtin.Sigma using (Σ; _,_)
open import Agda.Builtin.String
  using (String; primShowNat; primStringAppend; primStringEquality)
open import Agda.Builtin.Unit using (⊤)

open import SMT.Core using (Problem; Statement; counterexampleCheckSMT; counterexampleQuerySMT)

infixr  _++_
infixl  _>>=_

postulate
  execTC : String → List String → String →
           R.TC (Σ Nat (λ _ → Σ String (λ _ → String)))

{-# BUILTIN AGDATCMEXEC execTC #-}

_++_ : String → String → String
_++_ = primStringAppend

_>>=_ : ∀ {A B : Set} → R.TC A → (A → R.TC B) → R.TC B
_>>=_ = R.bindTC

return : ∀ {A : Set} → A → R.TC A
return = R.returnTC

vinfo : R.ArgInfo
vinfo = R.arg-info R.visible (R.modality R.relevant R.quantity-ω)

varg : ∀ {A : Set} → A → R.Arg A
varg = R.arg vinfo

z3Executable : String
z3Executable = "z3"

z3Args : List String
z3Args = "-in" ∷ "-smt2" ∷ []

failString : String → R.TC String
failString message = R.typeError (R.strErr message ∷ [])

showExit : Nat → String
showExit zero = "0"
showExit (suc n) = primShowNat (suc n)

runExecutable : String → List String → String → R.TC String
runExecutable exe args stdin =
  execTC exe args stdin >>= λ where
    (zero , (stdout , stderr)) → return stdout
    (suc n , (stdout , stderr)) →
      failString
        ( "external command failed with exit code " ++ showExit (suc n) ++
          "\nstdout:\n" ++ stdout ++ "\nstderr:\n" ++ stderr
        )

runZ3TC : String → R.TC String
runZ3TC script = runExecutable z3Executable z3Args script

z3VersionTC : R.TC String
z3VersionTC = runExecutable z3Executable ("--version" ∷ []) ""

fillStringTC : String → R.Term → R.TC ⊤
fillStringTC s hole = R.unify hole (R.lit (R.string s))

statementTypeTerm : R.Term → R.Term
statementTypeTerm p = R.def (quote Statement) (varg p ∷ [])

fillStatementTypeTC : ∀ {Γ} → Problem Γ → R.Term → R.TC ⊤
fillStatementTypeTC p hole =
  R.quoteTC p >>= λ pTerm →
  R.unify hole (statementTypeTerm pTerm)

checkedStatementResultTC : ∀ {Γ} → Problem Γ → R.Term → String → R.TC ⊤
checkedStatementResultTC p hole out with primStringEquality out "unsat\n"
... | true = fillStatementTypeTC p hole
... | false with primStringEquality out "sat\n"
...   | true =
  runZ3TC (counterexampleQuerySMT p) >>= λ model →
  R.typeError
    ( R.strErr
      ( "Z3 found a counterexample; refusing to produce an assumption type.\n" ++
        "Counterexample query output:\n" ++ model
      )
      ∷ []
    )
...   | false =
  R.typeError
    ( R.strErr
      ( "Z3 did not return a supported check-sat result; refusing to produce an assumption type.\n" ++
        "Raw output:\n" ++ out
      )
      ∷ []
    )

checkedStatementTC : ∀ {Γ} → Problem Γ → R.Term → R.TC ⊤
checkedStatementTC p hole =
  runZ3TC (counterexampleCheckSMT p) >>= checkedStatementResultTC p hole

macro
  z3Version : R.Term → R.TC ⊤
  z3Version hole =
    z3VersionTC >>= λ out →
    fillStringTC out hole

  z3Run : String → R.Term → R.TC ⊤
  z3Run script hole =
    runZ3TC script >>= λ out →
    fillStringTC out hole

  z3CounterexampleQuery : ∀ {Γ} → Problem Γ → R.Term → R.TC ⊤
  z3CounterexampleQuery p hole =
    runZ3TC (counterexampleQuerySMT p) >>= λ out →
    fillStringTC out hole

  z3CounterexampleCheck : ∀ {Γ} → Problem Γ → R.Term → R.TC ⊤
  z3CounterexampleCheck p hole =
    runZ3TC (counterexampleCheckSMT p) >>= λ out →
    fillStringTC out hole

  z3CheckedStatement : ∀ {Γ} → Problem Γ → R.Term → R.TC ⊤
  z3CheckedStatement = checkedStatementTC