(* We handle _very_ simple arithmetic expressions.
   Really nothing crazy yet, hopefully enough to handle
   most [while x < 10 do x = x + 1 done] loops! *)
type operator = Add | Sub | Mul | Div | Mod | Lt | Gt | Eq
[@@deriving to_yojson]

let exec_op = function
  | Add -> Int.add
  | Sub -> Int.sub
  | Mul -> Int.mul
  | Div -> Int.div
  | Mod -> ( mod )
  | Lt -> fun a b -> if a < b then 1 else 0
  | Gt -> fun a b -> if a > b then 1 else 0
  | Eq -> fun a b -> if Int.equal a b then 1 else 0

type expr =
  | Int of int
  | Var of string
  | Binop of operator * expr * expr
  | Neg of expr
  | Assign of operator * string * expr
  | Ternary of (expr * expr * expr)
[@@deriving to_yojson]

module Make (S : Types.State) = struct
  let eval initial_state expr =
    let lookup state s =
      match S.lookup state ~param:s with
      | Some [ Ast.WordLiteral n ] when Option.is_some (int_of_string_opt n) ->
          int_of_string n
      | _ -> 0
    in
    let update state s i =
      match S.update state ~param:s [ Ast.WordLiteral (string_of_int i) ] with
      | Ok s -> s
      | Error m -> failwith m
    in
    let rec calc state = function
      | Int i -> (state, i)
      | Var v -> (state, lookup state v)
      | Binop (op, e1, e2) ->
          let state, v1 = calc state e1 in
          let state, v2 = calc state e2 in
          (state, exec_op op v1 v2)
      | Neg n ->
          let state, v1 = calc state n in
          (state, Int.neg v1)
      | Assign (op, var, e) ->
          let current_v = lookup state var in
          let state, v1 = calc state e in
          let nv = exec_op op current_v v1 in
          (update state var nv, nv)
      | Ternary (e1, e2, e3) ->
          let state, v1 = calc state e1 in
          if Int.equal v1 Int.zero then calc state e3 else calc state e2
    in
    calc initial_state expr
end