module Syn = Spice_syntax.Ast

module Id = struct
  type t = {
    name : string;
    uid : int;
  }

  let to_string { name; uid } = Fmt.str "%s#%d" name uid
  let compare a b = Int.compare a.uid b.uid
  let hash a = a.uid

  let equal a b =
    if a.uid = b.uid then (
      assert (String.equal a.name b.name);
      true)
    else false
end

type name = Syn.name
type literal = Syn.literal
type binop = Syn.binop

(* TODO: split between arithmetic operators vs comparison operators *)

type arg = Arg of Id.t [@@unboxed]
type ele = Ele of arg * name

type exp =
  | Let of Id.t * rhs * exp
  | Set of ele * arg * exp
  | If of arg * exp * exp
  | Jump of Id.t * arg list
  | Ret of arg

and rhs =
  | Literal of literal
  | Get of ele
  | Call of ele * arg list
  | Binop of binop * arg * arg
  | Cont of cont
  | Obj of block

and block = {
  slots : name list;
  mthds : mthd list;
}

and mthd = {
  name : string;
  defn : cont;
}

and cont = {
  params : Id.t list;
  body : exp;
}

type entrypoint = cont

(* pretty printer *)

let pp_list pp_ele ppf list =
  Fmt.pf ppf "[";
  List.iteri
    (fun i ele ->
      if i > 0 then Fmt.pf ppf ",";
      pp_ele ppf ele)
    list;
  Fmt.pf ppf "]"

let pp_name ppf name = Fmt.pf ppf "%S" name

let rec pp_exp ppf exp =
  Fmt.pf ppf "[";
  let rec loop i exp =
    if i > 0 then Fmt.pf ppf ",";
    match exp with
    | Let (name, rhs, rest) ->
        Fmt.pf ppf "{\"let\":%a%a}" pp_id name pp_rhs rhs;
        loop (i + 1) rest
    | Set (e, v, rest) ->
        Fmt.pf ppf "{\"set\":%a,\"value\":%a}" pp_ele e pp_arg v;
        loop (i + 1) rest
    | If (v1, e2, e3) ->
        Fmt.pf ppf "{\"if\":%a,\"then\":%a,\"else\":%a}]" pp_arg v1 pp_exp e2 pp_exp e3
    | Jump (tgt, args) ->
        Fmt.pf ppf "{\"jump\":%a,\"args\":%a}]" pp_id tgt (pp_list pp_arg) args
    | Ret v -> Fmt.pf ppf "{\"ret\":%a}]" pp_arg v
  in
  loop 0 exp

and pp_id ppf id = Fmt.pf ppf "%S" (Id.to_string id)

and pp_arg ppf = function
  | Arg id -> pp_id ppf id

and pp_ele ppf = function
  | Ele (Arg ob, el) -> Fmt.pf ppf "\"%s.%s\"" (Id.to_string ob) el

and pp_rhs ppf = function
  | Literal (Int n) -> Fmt.pf ppf ",\"int\":%s" (Int64.to_string n)
  | Literal l -> Fmt.pf ppf ",\"lit\":%S" (Syn.string_of_literal l)
  | Get e -> Fmt.pf ppf ",\"get\":%a" pp_ele e
  | Call (fn, args) ->
      Fmt.pf ppf ",\"call\":%a,\"args\":%a" pp_ele fn (pp_list pp_arg) args
  | Binop (op, v1, v2) ->
      Fmt.pf
        ppf
        ",\"binop\":%S,\"lhs\":%a,\"rhs\":%a"
        (Syn.string_of_binop op)
        pp_arg
        v1
        pp_arg
        v2
  | Cont { params; body } ->
      Fmt.pf ppf ",\"cont\":%a,\"body\":%a" (pp_list pp_id) params pp_exp body
  | Obj { slots; mthds } ->
      Fmt.pf ppf ",\"obj\":{\"slots\":%a,\"mthds\":{" (pp_list pp_name) slots;
      List.iteri
        (fun i { name; defn = { params; body } } ->
          if i > 0 then Fmt.pf ppf ",";
          Fmt.pf
            ppf
            "%S:{\"params\":%a,\"body\":%a}"
            name
            (pp_list pp_id)
            params
            pp_exp
            body)
        mthds;
      Fmt.pf ppf "}}"

let pp_entrypoint ppf { params; body } =
  Fmt.pf ppf "{\"inputs\":%a,\"program\":%a}" (pp_list pp_id) params pp_exp body