(************************************************************************)
(*  v      *   The Coq Proof Assistant  /  The Coq Development Team     *)
(* <O___,, *   INRIA - CNRS - LIX - LRI - PPS - Copyright 1999-2015     *)
(*   \VV/  **************************************************************)
(*    //   *      This file is distributed under the terms of the       *)
(*         *       GNU Lesser General Public License Version 2.1        *)
(************************************************************************)

let debug = false

type request =
  | Hello of Flags.priority
  | Get of int
  | TryGet of int
  | GiveBack of int
  | Ping

type response =
  | Tokens of int
  | Noluck
  | Pong of int * int * int

exception ParseError

(* make it work with telnet: strip trailing \r  *)
let strip_r s =
  let len = String.length s in
  if s.[len - 1] <> '\r' then s else String.sub s 0 (len - 1)

let positive_int_of_string n =
  try
    let n = int_of_string n in
    if n <= 0 then raise ParseError else n
  with Invalid_argument _ | Failure _ -> raise ParseError

let parse_request s =
  if debug then Printf.eprintf "parsing '%s'\n" s;
  match Str.split (Str.regexp " ") (strip_r s) with
  | [ "HELLO"; "LOW" ] -> Hello Flags.Low
  | [ "HELLO"; "HIGH" ] -> Hello Flags.High
  | [ "GET"; n ] -> Get (positive_int_of_string n)
  | [ "TRYGET"; n ] -> TryGet (positive_int_of_string n)
  | [ "GIVEBACK"; n ] -> GiveBack (positive_int_of_string n)
  | [ "PING" ] -> Ping
  | _ -> raise ParseError

let parse_response s =
  if debug then Printf.eprintf "parsing '%s'\n" s;
  match Str.split (Str.regexp " ") (strip_r s) with
  | [ "TOKENS"; n ] -> Tokens (positive_int_of_string n)
  | [ "NOLUCK" ] -> Noluck
  | [ "PONG"; n; m; p ] ->
      let n = try int_of_string n with _ -> raise ParseError in
      let m = try int_of_string m with _ -> raise ParseError in
      let p = try int_of_string p with _ -> raise ParseError in
      Pong (n,m,p)
  | _ -> raise ParseError
       
let print_request = function
  | Hello Flags.Low -> "HELLO LOW\n"
  | Hello Flags.High -> "HELLO HIGH\n"
  | Get n -> Printf.sprintf "GET %d\n" n
  | TryGet n -> Printf.sprintf "TRYGET %d\n" n
  | GiveBack n -> Printf.sprintf "GIVEBACK %d\n" n
  | Ping -> "PING\n"

let print_response = function
  | Tokens n -> Printf.sprintf "TOKENS %d\n" n
  | Noluck -> "NOLUCK\n"
  | Pong (n,m,p) -> Printf.sprintf "PONG %d %d %d\n" n m p

let connect s =
  try
    match Str.split (Str.regexp ":") s with
    | [ h; p ] ->
        let open Unix in
        let s = socket PF_INET SOCK_STREAM 0 in
        connect s (ADDR_INET (inet_addr_of_string h,int_of_string p));
        Some s
    | _ -> None
  with Unix.Unix_error _ -> None

let manager = ref None

let option_map f = function None -> None | Some x -> Some (f x)

let init p =
  try
    let sock = Sys.getenv "COQWORKMGR_SOCK" in
    manager := option_map (fun s -> 
      let cout = Unix.out_channel_of_descr s in
      set_binary_mode_out cout true;
      let cin = Unix.in_channel_of_descr s in
      set_binary_mode_in cin true;
      output_string cout (print_request (Hello p)); flush cout;
      cin, cout) (connect sock)
  with Not_found | End_of_file -> ()

let with_manager f g =
  try
    match !manager with
    | None -> f ()
    | Some (cin, cout) -> g cin cout
  with
  | ParseError | End_of_file -> manager := None; f ()

let get n =
  with_manager
  (fun () ->
    min n (min !Flags.async_proofs_n_workers !Flags.async_proofs_n_tacworkers))
  (fun cin cout ->
    output_string cout (print_request (Get n));
    flush cout;
    let l = input_line cin in
    match parse_response l with
    | Tokens m -> m
    | _ -> raise (Failure "coqworkmgr protocol error"))

let tryget n =
  with_manager
  (fun () ->
    Some
     (min n
       (min !Flags.async_proofs_n_workers !Flags.async_proofs_n_tacworkers)))
  (fun cin cout ->
    output_string cout (print_request (TryGet n));
    flush cout;
    let l = input_line cin in
    match parse_response l with
    | Tokens m -> Some m
    | Noluck -> None
    | _ -> raise (Failure "coqworkmgr protocol error"))

let giveback n =
  with_manager
  (fun () -> ())
  (fun cin cout ->
    output_string cout (print_request (GiveBack n));
    flush cout)