open Merlin_kernel
open Merlin_utils
open Ocaml_parsing
open Ocaml_typing

let _ =
  Logs.set_reporter (Logs_fmt.reporter ());
  Logs.set_level (Some Debug)

module F = Fmt

module Log = (val Logs.src_log
                    (Logs.Src.create "dot_of_tast"
                       ~doc:"dot_of_tast.ml logger") : Logs.LOG)

module L = struct
  let stack = ref []
  let push l = stack := l :: !stack
  let peek () = match !stack with [] -> None | x :: _ -> Some x
  let pop () = match !stack with [] -> () | _ :: xs -> stack := xs
end

let vertex_attributes = ref (fun _ -> [])
let edge_attributes = ref (fun _ -> [])
let graph_name = ref None

module G = Graph.Imperative.Digraph.ConcreteBidirectional (struct
  type t = string

  let compare = String.compare
  let equal = String.equal
  let hash = Hashtbl.hash
end)

module Dot = Graph.Graphviz.Dot (struct
  include G

  let edge_attributes k = !edge_attributes k
  let default_edge_attributes _ = []
  let vertex_name k = k
  let vertex_attributes k = !vertex_attributes k
  let default_vertex_attributes _ = []
  let get_subgraph _ = None

  let graph_attributes _ =
    match !graph_name with None -> [] | Some n -> [ `Label n ]
end)

let str_of_exp (exp : Typedtree.expression) =
  Fmt.str "%a" Pprintast.expression (Untypeast.untype_expression exp)

let str_of_ident (ident : Ident.t) = Ident.name ident

let pp_value_binding_list ppf (vbl : Typedtree.value_binding list) =
  Fmt.pf ppf "%a"
    (Fmt.parens
       (Fmt.list (fun ppf (ident, _loc, _ty) -> Ident.print ppf ident)))
    (Typedtree.let_bound_idents_full vbl)

let str_of_primary_value (exp : Typedtree.expression) =
  match exp.exp_desc with
  | Texp_ident _ | Texp_constant _ -> str_of_exp exp
  | Texp_apply (exp, _) -> str_of_exp exp
  | Texp_function { param; _ } ->
      Fmt.str "fun %s" (str_of_ident param)
  | Texp_let (_, vbl, _) -> Fmt.str "let %a" pp_value_binding_list vbl
  | Texp_ifthenelse (_, _, Some _) -> Fmt.str "if/then/else"
  | Texp_ifthenelse (_, _, None) -> Fmt.str "if/then"
  | _ -> "????"

let graph_of_impl (s : Typedtree.structure) : G.t =
  let g = G.create () in
  let iterator =
    {
      Tast_iterator.default_iterator with
      expr =
        (fun iter exp ->
          match exp.exp_desc with
          | Texp_ident _ | Texp_constant _ | Texp_apply _
          | Texp_function _ | Texp_let _ | Texp_ifthenelse _ -> (
              let src = F.str "\"%s\"" (str_of_primary_value exp) in
              Log.debug (fun m -> m "src: %s" src);
              L.push (F.str "%s" src);
              (match exp.exp_desc with
              | Texp_apply (_exp, list) ->
                  (* Tast_iterator.default_iterator.expr iter exp *)
                  (* skip _exp *)
                  List.iter
                    (fun (_, o) -> Option.iter (iter.expr iter) o)
                    list
              | _ -> Tast_iterator.default_iterator.expr iter exp);
              L.pop ();
              match L.peek () with
              | Some dst ->
                  Log.debug (fun m -> m "G.add_edge %s %s" src dst);
                  G.add_edge g src dst
              | None -> ())
          | _ ->
              Tast_iterator.default_iterator.expr iter exp;
              Log.info (fun m ->
                  m "dot_of_impl unknown:@ %a" Pprintast.expression
                    (Untypeast.untype_expression exp)));
    }
  in
  iterator.structure iterator s;
  Log.info (fun m ->
      m "g_of_impl fprint_graph:@ %a" Dot.fprint_graph g);
  g

let merlin_parse str : G.t =
  let config, _command_args =
    Mconfig.parse_arguments ~wd:(Sys.getcwd ())
      ~warning:(fun _ -> ())
      (List.map snd []) [] Mconfig.initial []
  in
  File_id.with_cache @@ fun () ->
  let source = Msource.make str in
  let pipeline = Mpipeline.make config source in
  Mpipeline.with_pipeline pipeline @@ fun () ->
  Log.info (fun m -> m "merlin_parse...");
  let typer = Mpipeline.typer_result pipeline in
  let typedtree = Mtyper.get_typedtree typer in
  let _mbrowse = Mbrowse.of_typedtree typedtree in
  match typedtree with
  | `Implementation t ->
      Log.info (fun m ->
          m "Untyped:@ %a" Pprintast.structure
            (Untypeast.untype_structure t));
      Log.info (fun m ->
          m "Printtyped:@ %a" Ocaml_typing.Printtyped.implementation t);
      graph_of_impl t
  | `Interface _ -> G.create ()

let dot_of_tast ?(fname = "x.dot") str : unit =
  let oc = open_out fname in
  F.pf
    (Format.formatter_of_out_channel oc)
    "%a" Dot.fprint_graph (merlin_parse str);
  close_out oc