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

open Graph

module G = struct
  module EDGE = struct
    type t = string * Graphviz.DotAttributes.edge list

    let compare (a, _) (b, _) = String.compare a b
    let default = ("", [])
  end

  module VERTEX = struct
    type t = string * Graphviz.DotAttributes.vertex list

    let compare (a, _) (b, _) = String.compare a b
    let hash (a, _) = Hashtbl.hash a
    let equal (a, _) (b, _) = String.equal a b
  end

  module G = Imperative.Digraph.ConcreteLabeled (VERTEX) (EDGE)

  let vertex_name v =
    "\"" ^ String.escaped (fst @@ G.V.label v) ^ "\""

  let graph_attributes _ = []
  let default_vertex_attributes _ = []
  let vertex_attributes (_label, attr) = attr
  let default_edge_attributes _ = []

  let edge_attributes (_src, (label, attr), _dst) =
    if String.equal label "" then attr else `Label label :: attr

  let get_subgraph (_ : G.V.t) :
      Graphviz.DotAttributes.subgraph option =
    None

  include G
end

module Dot = Graphviz.Dot (struct
  include G
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 = 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_e 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