module P = Extend_protocol
module R = P.Reader

module Description = struct
  type t = P.description

  let make_v0 ~name ~version = { P.name; version }
end

module Reader = struct
  type t = (module R.V0)
  let make_v0 (x : (module R.V0)) : t = x

  module Make (V : R.V0) = struct
    open P.Reader

    let buffer = ref None

    let get_buffer () =
      match !buffer with
      | None -> invalid_arg "No buffer loaded"
      | Some buffer -> buffer

    let exec = function
      | Req_load buf ->
        buffer := Some (V.load buf);
        Res_loaded
      | Req_parse -> Res_parse (V.parse (get_buffer ()))
      | Req_parse_line (pos, str) ->
        Res_parse (V.parse_line (get_buffer ()) pos str)
      | Req_parse_for_completion pos ->
        let info, tree = V.for_completion (get_buffer ()) pos in
        Res_parse_for_completion (info, tree)
      | Req_get_ident_at pos ->
        Res_get_ident_at (V.ident_at (get_buffer ()) pos)
      | Req_print_outcome trees ->
        let print t =
          V.print_outcome Format.str_formatter t;
          Format.flush_str_formatter ()
        in
        let trees = List.rev_map print trees in
        Res_print_outcome (List.rev trees)
      | Req_pretty_print p ->
        V.pretty_print Format.str_formatter p;
        Res_pretty_print (Format.flush_str_formatter ())
  end
end

module Utils = struct
  (* Postpone messages until ready *)
  let send, set_ready =
    let is_ready = ref false in
    let postponed = ref [] in
    let really_send msg = output_value stdout msg in
    let set_ready () =
      is_ready := true;
      let postponed' = List.rev !postponed in
      postponed := [];
      List.iter really_send postponed'
    in
    let send msg =
      if !is_ready then really_send msg else postponed := msg :: !postponed
    in
    (send, set_ready)

  let notify msg = send (P.Notify msg)
  let debug msg = send (P.Debug msg)
end

module Handshake = struct
  let magic_number : string = "MERLINEXTEND002"

  type versions =
    { ast_impl_magic_number : string;
      ast_intf_magic_number : string;
      cmi_magic_number : string;
      cmt_magic_number : string
    }

  let versions =
    Config.
      { ast_impl_magic_number;
        ast_intf_magic_number;
        cmi_magic_number;
        cmt_magic_number
      }

  let negotiate (capabilities : P.capabilities) =
    output_string stdout magic_number;
    output_value stdout versions;
    output_value stdout capabilities;
    flush stdout;
    Utils.set_ready ();
    match input_value stdin with
    | exception End_of_file -> exit 0
    | P.Start_communication -> ()
    | _ ->
      prerr_endline "Unexpected value after handshake.";
      exit 1

  exception Error of string

  let () =
    Printexc.register_printer (function
      | Error msg -> Some (Printf.sprintf "Extend_main.Handshake.Error %S" msg)
      | _ -> None)

  let negotiate_driver ext_name i o =
    let magic' = really_input_string i (String.length magic_number) in
    (if magic' <> magic_number then
       let msg =
         Printf.sprintf
           "Extension %s has incompatible protocol version %S (expected %S)"
           ext_name magic' magic_number
       in
       raise (Error msg));
    let versions' : versions = input_value i in
    let check_v prj name =
      if prj versions <> prj versions' then
        let msg =
          Printf.sprintf
            "Extension %s %s has incompatible version %S (expected %S)" ext_name
            name (prj versions') (prj versions)
        in
        raise (Error msg)
    in
    check_v (fun x -> x.ast_impl_magic_number) "implementation AST";
    check_v (fun x -> x.ast_intf_magic_number) "interface AST";
    check_v (fun x -> x.cmi_magic_number) "compiled interface (CMI)";
    check_v (fun x -> x.cmt_magic_number) "typedtree (CMT)";
    output_value o P.Start_communication;
    flush o;
    let capabilities : P.capabilities = input_value i in
    capabilities
end

(** The main entry point of an extension. *)
let extension_main ?reader desc =
  (* Check if invoked from Merlin *)
  begin
    match Sys.getenv "__MERLIN_MASTER_PID" with
    | exception Not_found ->
      Printf.eprintf
        "This is %s merlin extension, version %s.\n\
         This binary should be invoked from merlin and cannot be used directly.\n\
         %!"
        desc.P.name desc.P.version;
      exit 1
    | _ -> ()
  end;
  (* Communication happens on stdin/stdout. *)
  Handshake.negotiate { P.reader = reader <> None };
  let reader =
    match reader with
    | None -> fun _ -> failwith "No reader"
    | Some (module R : R.V0) ->
      let module M = Reader.Make (R) in
      M.exec
  in
  let respond f =
    match f () with
    | (r : P.response) -> Utils.send r
    | exception exn ->
      let name = Printexc.exn_slot_name exn in
      let desc = Printexc.to_string exn in
      Utils.send (P.Exception (name, desc))
  in
  let rec loop () =
    flush stdout;
    match input_value stdin with
    | exception End_of_file -> exit 0
    | P.Start_communication ->
      prerr_endline "Unexpected message.";
      exit 2
    | P.Reader_request request ->
      respond (fun () -> P.Reader_response (reader request));
      loop ()
  in
  loop ()
