module T = Sidekick_cic_lib.Term
module Const = Sidekick_cic_lib.Const
module Level = Sidekick_cic_lib.Level
module Inductive = Sidekick_cic_lib.Inductive

let ( let@ ) = ( @@ )

(** Map indexes to objects *)
module Idx = struct
  type t = {
    names: string Vec.t;
    levels: Level.t Vec.t;
    terms: Lean_term.t Vec.t;
  }

  (* create. The slot 0 is already filled with name "" *)
  let create () : t =
    let names = Vec.create () in
    Vec.push names "";
    let levels = Vec.create () in
    let terms = Vec.create () in
    { names; levels; terms }

  (** ensure that index [n] is valid *)
  let ensure_size (v : _ Vec.t) ~dummy (n : int) : unit =
    if n >= Vec.size v then Vec.ensure_size v ~elt:dummy n

  let set_name self i n : unit =
    ensure_size self.names ~dummy:"" (i + 1);
    Vec.set self.names i n

  let get_name self i : string =
    match Vec.get self.names i with
    | s -> s
    | exception _ -> Error.errorf "invalid name at %d" i

  let set_term self i t : unit =
    ensure_size self.terms ~dummy:Lean_term.dummy (i + 1);
    Vec.set self.terms i t

  let get_term self i : Lean_term.t =
    match Vec.get self.terms i with
    | s -> s
    | exception _ -> Error.errorf "invalid term at %d" i

  let set_level self i l : unit =
    ensure_size self.levels ~dummy:Level.zero (i + 1);
    Vec.set self.levels i l

  let get_level self i : Level.t =
    match Vec.get self.levels i with
    | s -> s
    | exception _ -> Error.errorf "invalid level at %d" i

  let dump out (self : t) : unit =
    let pp_with_idx ppx out (i, x) = Fmt.fprintf out "[%d]=%a" i ppx x in
    Fmt.fprintf out
      "(@[idx {@ names: [@[%a@]],@ terms: [@[%a@]],@ levels: [@[%a@]]}@])"
      (Fmt.iter ~sep:(Fmt.return ";@ ") (pp_with_idx Fmt.Dump.string))
      (Vec.to_iter self.names |> Iter.zip_i)
      (Fmt.iter ~sep:(Fmt.return ";@ ") (pp_with_idx Lean_term.pp))
      (Vec.to_iter self.terms |> Iter.zip_i)
      (Fmt.iter ~sep:(Fmt.return ";@ ") Level.pp)
      (Vec.to_iter self.levels)
end

(** Join two names with "." *)
let name_join a b =
  if a = "" then
    b
  else
    a ^ "." ^ b

let binder_of_string = function
  | "#BD" -> Lean_term.BD
  | "#BI" -> Lean_term.BI
  | "#BS" -> Lean_term.BS
  | "#BC" -> Lean_term.BC
  | s -> failwith (Printf.sprintf "invalid binder: %S" s)

let process_files ~max_err ~quiet (files : string list) : unit =
  let start = Mtime_clock.now () in

  Log.debugf 1 (fun k ->
      k "(@[process-files %a@])" Fmt.Dump.(list string) files);

  let n_lines = ref 0 in

  (* get a level. 0 means level 0. *)
  let get_level idx i =
    if i = 0 then
      Level.zero
    else
      Idx.get_level idx i
  in

  let proc_file (file : string) : unit =
    let@ ic = CCIO.with_in file in
    let idx = Idx.create () in

    Parse.parse ~max_errors:max_err (`In ic)
      (module struct
        let line _ = incr n_lines

        let ns ~at i s : unit =
          Idx.set_name idx at (name_join (Idx.get_name idx i) s)

        let ni ~at i n : unit =
          Idx.set_name idx at (name_join (Idx.get_name idx i) (string_of_int n))

        let us ~at i : unit =
          Idx.set_level idx at (Level.succ @@ get_level idx i)

        let um ~at i j : unit =
          Idx.set_level idx at (Level.max (get_level idx i) (get_level idx j))

        let uim ~at i j : unit =
          Idx.set_level idx at (Level.imax (get_level idx i) (get_level idx j))

        let up ~at i : unit =
          Idx.set_level idx at (Level.var @@ Idx.get_name idx i)

        let ev ~at i : unit = Idx.set_term idx at (Lean_term.BVar i)

        let ea ~at i j : unit =
          Idx.set_term idx at
            (Lean_term.App (Idx.get_term idx i, Idx.get_term idx j))

        let ec ~at i_name i_args : unit =
          let name = Idx.get_name idx i_name in
          let args = List.map (Idx.get_level idx) i_args in
          Idx.set_term idx at (Lean_term.Const (name, args))

        let es ~at i : unit =
          Idx.set_term idx at (Lean_term.Sort (Idx.get_level idx i))

        let el ~at b n i j : unit =
          Idx.set_term idx at
            (Lean_term.Lam
               ( binder_of_string b,
                 Idx.get_name idx n,
                 Idx.get_term idx i,
                 Idx.get_term idx j ))

        let ep ~at b n i j : unit =
          Idx.set_term idx at
            (Lean_term.Pi
               ( binder_of_string b,
                 Idx.get_name idx n,
                 Idx.get_term idx i,
                 Idx.get_term idx j ))

        let ind ~n_params ~nidx ~tyidx ~intros ~univ_params : unit =
          let name = Idx.get_name idx nidx in
          let ty = Lean_elab.elab_top (Idx.get_term idx tyidx) in
          let cstors =
            List.map
              (fun (i, j) ->
                {
                  Inductive.c_name = Idx.get_name idx i;
                  c_ty = Lean_elab.elab_top (Idx.get_term idx j);
                })
              intros
          in
          let univ_params =
            List.map (fun i -> Idx.get_name idx i) univ_params
          in

          let spec = { Inductive.name; ty; n_params; univ_params; cstors } in

          if not quiet then (
            Fmt.eprintf "@[<2>@{<Green>Ind@}@ %a@]@." Inductive.pp_spec spec;
            (* TODO: register to the conversion relation *)
            Fmt.eprintf "is valid: %B@." (Inductive.is_valid spec)
          );
          ()

        let def ~nidx ~tyidx ~validx ~univ_params : unit =
          let name = Idx.get_name idx nidx in
          let _ty = Lean_elab.elab_top (Idx.get_term idx tyidx) in
          let _val = Lean_elab.elab_top (Idx.get_term idx validx) in
          let _univ_params = List.map (Idx.get_name idx) univ_params in
          if not quiet then
            Fmt.eprintf "@[<2>@{<Blue>Def@}@ %s@]@." name

        let after_line () =
          if quiet then (
            if !n_lines mod 1000 = 0 then Fmt.eprintf "\rlines: %d%!" !n_lines
          ) else
            Fmt.eprintf "%a@." Idx.dump idx
      end)
  in

  List.iter proc_file files;

  let elapsed =
    (Mtime.span (Mtime_clock.now ()) start |> Mtime.Span.to_float_ns) /. 1e9
  in
  Log.debugf 1 (fun k ->
      k "number of lines processed: %d in %.4fs (%.2f/s)" !n_lines elapsed
        (float !n_lines /. elapsed));
  ()

let () =
  let files = ref [] in
  let max_err = ref max_int in
  let quiet = ref false in
  let opts =
    [
      "--debug", Arg.Int Log.set_debug, " set debug level";
      "-d", Arg.Int Log.set_debug, " like --debug";
      ( "--max-err",
        Arg.Set_int max_err,
        " maximum number of errors before bailing out" );
      ( "-c",
        Arg.Unit (fun () -> CCFormat.set_color_default true),
        " enable colors" );
      ( "-nc",
        Arg.Unit (fun () -> CCFormat.set_color_default false),
        " disable colors" );
      ( "-q",
        Arg.Set quiet,
        " quiet: suppress verbose output, show progress every 1000 lines" );
      "--quiet", Arg.Set quiet, " like -q";
    ]
    |> Arg.align
  in

  Arg.parse opts (CCList.Ref.push files) "leancheck file+";
  if !files = [] then failwith "provide at least one file";

  process_files ~max_err:!max_err ~quiet:!quiet (List.rev !files)