open Std

type parsetree =
  [ `Interface of Parsetree.signature | `Implementation of Parsetree.structure ]

type comment = string * Location.t

type result =
  { lexer_keywords : string list;
    lexer_errors : exn list;
    parser_errors : exn list;
    comments : comment list;
    parsetree : parsetree;
    no_labels_for_completion : bool
  }

(* Normal entry point *)

let normal_parse ?for_completion config source =
  let kind =
    let filename = Mconfig.(config.query.filename) in
    let extension =
      match String.rindex filename '.' with
      | exception Not_found -> ""
      | pos -> String.sub ~pos ~len:(String.length filename - pos) filename
    in
    Logger.log ~section:"Mreader" ~title:"run" "extension(%S) = %S" filename
      extension;
    if
      List.exists
        ~f:(fun (_impl, intf) -> intf = extension)
        Mconfig.(config.merlin.suffixes)
    then Mreader_parser.MLI
    else Mreader_parser.ML
  in
  let lexer =
    let keywords = Extension.keywords Mconfig.(config.merlin.extensions) in
    Mreader_lexer.make Mconfig.(config.ocaml.warnings) keywords config source
  in
  let no_labels_for_completion, lexer =
    match for_completion with
    | None -> (false, lexer)
    | Some pos ->
      let pos =
        Msource.get_lexing_pos source ~filename:(Mconfig.filename config) pos
      in
      Mreader_lexer.for_completion lexer pos
  in
  let parser = Mreader_parser.make Mconfig.(config.ocaml.warnings) lexer kind in
  let lexer_keywords = Mreader_lexer.keywords lexer
  and lexer_errors = Mreader_lexer.errors lexer
  and parser_errors = Mreader_parser.errors parser
  and parsetree = Mreader_parser.result parser
  and comments = Mreader_lexer.comments lexer in
  { lexer_keywords;
    lexer_errors;
    parser_errors;
    comments;
    parsetree;
    no_labels_for_completion
  }

(* Pretty-printing *)

type pretty_parsetree = Extend_protocol.Reader.pretty_parsetree
type outcometree = Extend_protocol.Reader.outcometree

let ambient_reader = ref None

let instantiate_reader spec config source =
  match spec with
  | [] -> (lazy None, ignore)
  | name :: args -> (
    let reader = lazy (Mreader_extend.start name args config source) in
    ( reader,
      fun () ->
        if Lazy.is_val reader then
          match Lazy.force reader with
          | None -> ()
          | Some reader -> Mreader_extend.stop reader ))

let get_reader config =
  let rec find_reader assocsuffixes =
    match assocsuffixes with
    | [] -> []
    | (suffix, reader) :: t ->
      if Filename.check_suffix Mconfig.(config.query.filename) suffix then
        [ reader ]
      else find_reader t
  in
  match Mconfig.(config.merlin.reader) with
  (* if a reader flag exists then this is explicitly used disregarding suffix association *)
  | [] -> find_reader Mconfig.(config.merlin.extension_to_reader)
  | x -> x

let mocaml_printer reader ppf otree =
  let str =
    match reader with
    | (lazy (Some reader)) -> Mreader_extend.print_outcome otree reader
    | _ -> None
  in
  match str with
  | Some str -> Format.pp_print_string ppf str
  | None -> Mocaml.default_printer ppf otree

let with_ambient_reader config source f =
  let ambient_reader' = !ambient_reader in
  let reader_spec = get_reader config in
  let reader, stop = instantiate_reader reader_spec config source in
  ambient_reader := Some (reader, reader_spec, source);
  Misc.try_finally
    (fun () -> Mocaml.with_printer (mocaml_printer reader) f)
    ~always:(fun () ->
      ambient_reader := ambient_reader';
      stop ())

let try_with_reader config source f =
  let reader_spec = get_reader config in
  let (lazy reader), stop =
    match !ambient_reader with
    | Some (reader, reader_spec', source')
      when compare reader_spec reader_spec' = 0 && compare source source' = 0 ->
      (reader, ignore)
    | _ -> instantiate_reader reader_spec config source
  in
  match reader with
  | None ->
    stop ();
    None
  | Some reader -> Misc.try_finally (fun () -> f reader) ~always:stop

let print_pretty config source tree =
  match try_with_reader config source (Mreader_extend.print_pretty tree) with
  | Some result -> result
  | None ->
    let ppf, to_string = Std.Format.to_string () in
    let open Extend_protocol.Reader in
    begin
      match tree with
      | Pretty_case_list x -> Pprintast.case_list ppf x
      | Pretty_core_type x -> Pprintast.core_type ppf x
      | Pretty_expression x -> Pprintast.expression ppf x
      | Pretty_pattern x -> Pprintast.pattern ppf x
      | Pretty_signature x -> Pprintast.signature ppf x
      | Pretty_structure x -> Pprintast.structure ppf x
      | Pretty_toplevel_phrase x -> Pprintast.toplevel_phrase ppf x
    end;
    to_string ()

let default_print_outcome tree =
  Mocaml.default_printer Format.str_formatter tree;
  Format.flush_str_formatter ()

let print_outcome config source tree =
  match try_with_reader config source (Mreader_extend.print_outcome tree) with
  | Some result -> result
  | None -> default_print_outcome tree

let print_batch_outcome config source tree =
  match try_with_reader config source (Mreader_extend.print_outcomes tree) with
  | Some result -> result
  | None -> List.map ~f:default_print_outcome tree

let reconstruct_identifier config source pos =
  match
    try_with_reader config source (Mreader_extend.reconstruct_identifier pos)
  with
  | None | Some [] -> Mreader_lexer.reconstruct_identifier config source pos
  | Some result -> result

(* Entry point *)

let parse ?for_completion config = function
  | source, None -> begin
    match
      try_with_reader config source (Mreader_extend.parse ?for_completion)
    with
    | Some (`No_labels no_labels_for_completion, parsetree) ->
      let lexer_errors, parser_errors, comments = ([], [], []) in
      let lexer_keywords =
        []
        (* TODO? *)
      in
      { lexer_keywords;
        lexer_errors;
        parser_errors;
        comments;
        parsetree;
        no_labels_for_completion
      }
    | None -> normal_parse ?for_completion config source
  end
  | _, Some parsetree ->
    let lexer_errors, parser_errors, comments = ([], [], []) in
    let lexer_keywords = [] in
    { lexer_keywords;
      lexer_errors;
      parser_errors;
      comments;
      parsetree;
      no_labels_for_completion = false
    }
