import Lean
import Cegar

open Lean
open Cegar

structure Input where
  states : Nat
  keys : Nat
  initial : List Nat
  bad : List Nat
  edges : List (Nat * Nat)
  seed : Array Nat
  deriving FromJson

structure Output where
  badWalk : Option (List Nat)
  rounds : Nat
  deriving ToJson

def verify (input : Input) : Except String Output := do
  if positive : 0 < input.keys then
    if input.seed.size != input.states || !(input.seed.all (. < input.keys)) then
      throw "seed must give one valid key for each concrete state"
    if !(input.initial.all (. < input.states)) || !(input.bad.all (. < input.states)) ||
        !(input.edges.all (fun (x, y) => x < input.states && y < input.states)) then
      throw "boundary and edge indices must belong to the concrete carrier"
    let graph : Graph (Fin input.states) :=
      <fun x => input.initial.contains x.val,
       fun x => input.bad.contains x.val,
       fun x y => input.edges.contains (x.val, y.val)>
    let seed (x : Fin input.states) : Fin input.keys :=
      <input.seed[x.val]?.getD 0 % input.keys, Nat.mod_lt _ positive>
    let result := run graph (finEnumeration input.states) (finEnumeration input.keys) seed
    pure <(verdictCode result.verdict).map (..map Fin.val), result.rounds>
  else
    throw "keys must be positive; one unused key suffices for the empty carrier"

def main (args : List String) : IO UInt32 := do
  let [file] := args | IO.eprintln "usage: cegar_cli <batch.json>"; return 2
  let raw <- IO.FS.readFile file
  let result := do
    let json <- Json.parse raw
    let inputs <- fromJson? (alpha := List Input) json
    inputs.mapM verify
  match result with
  | .ok rows => IO.println (toJson rows).compress; return 0
  | .error error => IO.eprintln error; return 1
