(** {1 Code generators} *)

module Fmt = CCFormat
let spf = Printf.sprintf
let fpf = Fmt.fprintf

type code =
  | Base of { pp: unit Fmt.printer }
  | Struct of string * code list
  | Sig of string * code list

module Code = struct
  type t = code

  let in_struct m (cs:t list) : t = Struct (m, cs)
  let in_sig m (cs:t list) : t = Sig (m, cs)

  let rec pp_rec out c =
    let ppl = Fmt.(list ~sep:(return "@ ") pp_rec) in
    match c with
      | Base {pp} -> pp out ()
      | Struct (m,cs) ->
        fpf out "@[<hv2>module %s = struct@ %a@;<1 -2>end@]" m ppl cs
      | Sig (m,cs) ->
        fpf out "@[<hv2>module %s : sig@ %a@;<1 -2>end@]" m ppl cs

  let pp out c = fpf out "@[<v>%a@]" pp_rec c
  let to_string c = Fmt.to_string pp c

  let mk_pp pp = Base {pp}
  let mk_str s = Base {pp=Fmt.const Fmt.string s}
end

module Bitfield = struct
  type field = {
    f_name: string;
    f_offset: int;
    f_def: field_def;
  }
  and field_def =
    | F_bit
    | F_int of {width: int}

  type t = {
    name: string;
    mutable fields: field list;
    mutable width: int;
    emit_failure_if_too_wide: bool;
  }

  let make ?(emit_failure_if_too_wide=true) ~name () : t =
    { name; fields=[]; width=0; emit_failure_if_too_wide; }

  let total_width self = self.width

  let field_bit self f_name =
    let f_offset = total_width self in
    let f = {f_name; f_offset; f_def=F_bit} in
    self.fields <- f :: self.fields;
    self.width <- 1 + self.width

  let field_int self ~width f_name : unit =
    let f_offset = total_width self in
    let f = {f_name; f_offset; f_def=F_int {width}} in
    self.fields <- f :: self.fields;
    self.width <- self.width + width

  let empty_name self =
    if self.name = "t" then "empty" else spf "empty_%s" self.name

  let gen_ml self : code =
    Code.mk_pp @@ fun out () ->
    fpf out "@[<v>type %s = int@," self.name;
    fpf out "@[let %s : %s = 0@]@," (empty_name self) self.name;
    List.iter
      (fun f ->
         let inline = "[@inline]" in (* TODO: option to enable/disable that *)
         let off = f.f_offset in
         match f.f_def with
         | F_bit ->
           let x_lsr = if off = 0 then "x" else spf "(x lsr %d)" off in
           fpf out "@[let%s get_%s (x:%s) : bool = (%s land 1) <> 0@]@,"
             inline f.f_name self.name x_lsr;
           let mask_shifted = 1 lsl off in
           fpf out "@[<2>let%s set_%s (v:bool) (x:%s) : %s =@ \
                    if v then x lor %d else x land (lnot %d)@]@,"
             inline f.f_name self.name self.name mask_shifted mask_shifted;
         | F_int {width} ->
           let mask0 = (1 lsl width) - 1 in
           fpf out "@[let%s get_%s (x:%s) : int = ((x lsr %d) land %d)@]@,"
             inline f.f_name self.name off mask0;
           fpf out "@[<2>let%s set_%s (i:int) (x:%s) : %s =@ \
                    assert ((i land %d) == i);@ \
                    ((x land (lnot %d)) lor (i lsl %d))@]@,"
             inline f.f_name self.name self.name
             mask0 (mask0 lsl off) off;
      )
      (List.rev self.fields);
    (* check width *)
    if self.emit_failure_if_too_wide then (
      fpf out "(* check that int size is big enough *)@,\
               @[let () = assert (Sys.int_size >= %d);;@]" (total_width self);
    );
    fpf out "@]"

  let gen_mli self : code =
    Code.mk_pp @@ fun out () ->
    fpf out "@[<v>type %s = private int@," self.name;
    fpf out "@[<v>val %s : %s@," (empty_name self) self.name;
    List.iter
      (fun f ->
         match f.f_def with
         | F_bit ->
           fpf out "@[val get_%s : %s -> bool@]@," f.f_name self.name;
           fpf out "@[val set_%s : bool -> %s -> %s@]@," f.f_name self.name self.name;
         | F_int {width} ->
           fpf out "@[val get_%s : %s -> int@]@,"
             f.f_name self.name;
           fpf out "@,@[(** %d bits integer *)@]@,\
                    @[val set_%s : int -> %s -> %s@]@,"
             width f.f_name self.name self.name;
      )
      (List.rev self.fields);
    fpf out "@]"
end

let emit_chan oc cs =
  let fmt = Fmt.formatter_of_out_channel oc in
  List.iter (fun c -> Fmt.fprintf fmt "@[%a@]@." Code.pp c) cs;
  Fmt.fprintf fmt "@?"

let emit_file file cs =
  CCIO.with_out file (fun oc -> emit_chan oc cs)

let emit_string cs : string =
  Fmt.asprintf "@[<v>%a@]" (Fmt.list ~sep:(Fmt.return "@ ") Code.pp) cs