(** Shell command utilities. *)

(** Depends: String, Sys, Unix, Stream, Stream2, Channel *)

type elem =
    Stdout of string
  | Stderr of string
  | Exit of int
  | Stopped of int
  | Signaled of int
  | Open

let string_of_elem =
  function
      Stdout s -> "Stdout \"" ^ s ^ "\""
    | Stderr s -> "Stderr \"" ^ s ^ "\""
    | Exit n -> "Exit " ^ string_of_int n
    | Stopped n -> "Stopped " ^ string_of_int n
    | Signaled n -> "Signaled " ^ string_of_int n
    | Open -> "Open"

let to_stream cmd =
  let (output, input, errors) = Unix.open_process_full cmd (Unix.environment ()) in
  let outdescr = Unix.descr_of_in_channel output
  and errdescr = Unix.descr_of_in_channel errors in
  let opendescrs = ref [errdescr; outdescr] in
  (* Calling set_nonblock means we will get Sys_blocked_io exceptions *)
  Unix.set_nonblock outdescr;
  Unix.set_nonblock errdescr;
  let outstrm = Channel.to_stringstream output
  and errstrm = Channel.to_stringstream errors in
  (* Read lines from stderr until eof or it would block, then do the same
     for stdout *)
  Stream.from
    (fun _ ->
       let rec loop () =
	 match !opendescrs with
	   [] -> begin
	     (*Log.put ~v:2 "closing process";*)
	     try
	       match Unix.close_process_full (output, input, errors) with
		 Unix.WEXITED n -> Some (Exit n)
	       | Unix.WSTOPPED n -> Some (Stopped n)
	       | Unix.WSIGNALED n -> Some (Signaled n)
	     with
	       Unix.Unix_error (Unix.EBADF, _, _) -> None
	   end
	 | _ -> begin
	     (* To avoid busy waiting, do a select on the two descriptors *)
	     begin
	       match (Unix.select !opendescrs [] [] (-1.0)) with
		 (descr :: _, _, _) when descr = errdescr -> begin
		   (*Log.v2_eput "stderr ready";*)
		   try Some (Stderr (Stream.next errstrm))
		   with
		     Stream.Failure ->
		       (*Log.put ~v:2 "EOF on stderr";*)
		       opendescrs := List.filter (fun x -> x != errdescr) !opendescrs;
		       loop ()
		   | Sys_blocked_io ->
		       (*Log.put ~v:2 "Sys_blocked_io on stderr";*)
		       loop ()
		 end
	       | (descr :: _, _, _) when descr = outdescr -> begin
		   (*Log.v2_eput "stdout ready";*)
		     ignore (Unix.select [outdescr; errdescr] [] [] (-1.0));
		   try Some (Stdout (Stream.next outstrm))
		   with
		     Stream.Failure ->
		       (*Log.put ~v:2 "EOF on stdout";*)
		       opendescrs := List.filter (fun x -> x != outdescr) !opendescrs;
		       loop ()
		   | Sys_blocked_io ->
		       (*Log.put ~v:2 "Sys_blocked_io on stdout";*)
		       loop ()
		 end
	       | _ ->
		   failwith "internal error"
	     end
	   end in
       loop ())

let to_strings cmd =
  let errout = Buffer.create 16
  and output = Buffer.create 16
  and eos = ref Open in
  Stream.iter
    (function
	 Stderr s -> Buffer.add_string errout s
       | Stdout s -> Buffer.add_string output s
       | final -> eos := final)
    (to_stream cmd);
  (Buffer.contents output, Buffer.contents errout, !eos)

let to_stringstream cmd =
  Stream2.map
    (function
	 Stdout s -> s
       | Stderr s -> ""
       | Exit 0 -> ""
       | elem -> 
	   prerr_endline ("Sh.to_stringstream warning: " ^ cmd ^ " -> " ^ string_of_elem elem);
	   "")
    (to_stream cmd)

let to_linestream cmd =
  Linestream.of_stringstream (to_stringstream cmd)

let to_string cmd =
  let buf = Buffer.create 16
  and eos = ref Open in
  Stream.iter
    (function
	 Stdout s -> Buffer.add_string buf s
       | Stderr s -> ()
       | final -> eos := final)
    (to_stream cmd);
  Buffer.contents buf, !eos

let to_linelist cmd =
  Stream2.to_list (to_linestream cmd)

let line1 cmd =
  let chan = Unix.open_process_in cmd in
  let line = try Some (input_line chan) with _ -> None in
  ignore (Unix.close_process_in chan);
  line

(** Wrapper around Sys.command.  Print and execute a command, raise a
  Failure exception if its result code is nonzero. *)
let run cmd =
  match Sys.command ("set -e; " ^ cmd) with
    0 -> ()
  | n -> raise (Failure (cmd ^ " -> " ^ (string_of_int n)))

(** Version that prints the command to be run (ala shell's -x flag) *)
let vrun cmd = prerr_endline ("-> " ^ cmd); run cmd

(** Send a stream's contents to a command *)
let stream_in cmd strm =
  let chan = Unix.open_process_out cmd in
  Stream.iter (fun line -> output_string chan (line ^ "\n")) strm;
  ignore (Unix.close_process_out chan)

(*type line = Stdout of string | Stderr of string*)

(** capture stdout and stderr of a command, raise exception on
  failure.  FIXME: it busywaits - I think the select fixes this. *)
let filter_old errfn outfn s =
  let (output, input, errors) =
    Unix.open_process_full s (Unix.environment ()) in
  let outdescr = Unix.descr_of_in_channel output
  and errdescr = Unix.descr_of_in_channel errors in
  Unix.set_nonblock outdescr;
  Unix.set_nonblock errdescr;
  (* Read lines from stderr until eof or it would block, then do the same
     for stdout *)
  let stderreof = ref false
  and stdouteof = ref false in
  while not !stderreof || not !stdouteof do
    ignore (Unix.select [outdescr; errdescr] [] [] (-1.0));
    if not !stderreof then begin
      try while true do errfn (input_line errors) done
      with
	End_of_file -> stderreof := true
      | Sys_blocked_io -> ()
    end;
    if not !stdouteof then begin
      try while true do outfn (input_line output) done
      with
	End_of_file -> stdouteof := true
      | Sys_blocked_io -> ()
    end;
  done;
  match Unix.close_process_full (output, input, errors) with
    Unix.WEXITED 0 -> ()
  | Unix.WEXITED n -> raise (Failure (s ^ " -> " ^ (string_of_int n)))
  | Unix.WSTOPPED n -> raise (Failure (s ^ " stopped: " ^ (string_of_int n)))
  | Unix.WSIGNALED n -> raise (Failure (s ^ " signaled: " ^ (string_of_int n)))

let filter errfn outfn cmd =
  let result = ref Open in
  Stream.iter
    (function
	 Stdout s -> outfn s
       | Stderr s -> errfn s
       | other -> result := other)
    (to_stream cmd);
  !result

(** Send the text to the command's standard input and return its
    standard output.  We need to write the input string to the process
    standard input and read from its output, but if we write too much
    at once our write call will block.  How is this done in a nice
    functional fashion? *)

let filter2 cmd text =

  let bufsize = 4096 in
  let buf = String.create bufsize in
  let (ichan, ochan) = Unix.open_process cmd in
  let (idescr, odescr) = Unix.descr_of_in_channel ichan, Unix.descr_of_out_channel ochan in

  Unix.set_nonblock idescr;
  Unix.set_nonblock odescr;

  let olist = ref [odescr] in

  let length = String.length text in
  let result = ref "" in

  let rec loop pos =
    let count = length - pos in
    match Unix.select [idescr] !olist [] (-1.0) with
      (idescr :: _), _, _ -> begin
	match Unix.read idescr buf 0 bufsize with
	  m when m > 0 ->
	    result := !result ^ String.sub buf 0 m;
	    loop pos
	| 0 ->
	    ignore (Unix.close_process (ichan, ochan));
	    !result
	| m ->
	    ignore (Unix.close_process (ichan, ochan));
	    !result
      end
    | [], (odescr :: _), _ -> begin
	if count > 0
	then begin
	  match Unix.single_write odescr text pos count with
	    n when n > 0 ->
	      loop (pos + n);
	  | n ->
	      ignore (Unix.close_process (ichan, ochan));
	      !result
	end
	else begin
	  Unix.close odescr;
	  olist := [];
	  loop pos
	end
      end
    | _ -> loop pos in 
  loop 0

(*
#use "topfind";;
#directory "+threads"
#require "unix"
#load "threads.cma"
*)

(* Run a process and manage the input and output using threads.  A
   thread is created to read input strings from INSTREAM and write
   them to the process's standard input.  Two more threads are created
   to read strings from the standard output and standard error
   channels and pass those strings to OUTFN and ERRFN.  the default
   input stream is the empty stream, and the default output functions
   write the strings to the console. *)

(* This function requires the thread library. *)
(*
type process_result =
    Exit of int
  | Stopped of int
  | Signaled of int

exception ExitExn of int
exception StoppedExn of int
exception SignaledExn of int

let filter3 =

  let apply_channel chan f =
    let buf = String.create 4096 in
    let eof = ref false in
    while not !eof do
      match input chan buf 0 4096 with
	0 -> eof := true
      | n -> f buf 0 n
    done in

  fun ?(instream=(Stream.of_list []))
      ?(outfn=output stdout)
      ?(errfn=output stderr)
      ?(onexit=fun n -> if n <> 0 then raise (ExitExn n) else Exit 0)
      ?(onstop=fun n -> raise (StoppedExn n))
      ?(onsignal=fun n -> raise (SignaledExn n))
      cmd ->
	let (std_output, std_input, std_error) = Unix.open_process_full cmd (Unix.environment ()) in
	let stdin_writer = Thread.create (Stream.iter (output_string std_input)) instream
	and stdout_reader = Thread.create (apply_channel std_output) outfn
	and stderr_reader = Thread.create (apply_channel std_error) errfn in
	Thread.join stdin_writer;
	Thread.join stdout_reader;
	Thread.join stderr_reader;
	match Unix.close_process_full (std_output, std_input, std_error) with
	  Unix.WEXITED n -> onexit n
	| Unix.WSTOPPED n -> onstop n
	| Unix.WSIGNALED n -> onsignal n
*)

(*
  let outfn c = prerr_endline (Printf.sprintf "output: '%c'" c)
  let errfn c = prerr_endline (Printf.sprintf "errout: '%c'" c)

  let count_test =
  fun cmd ->
  let stdout_chars = ref 0
  and stderr_chars = ref 0 in
  filter3
  ~outfn:(fun s off len -> stdout_chars := !stdout_chars + len)
  ~errfn:(fun s off len -> stderr_chars := !stderr_chars + len)
  ~onexit:(fun n -> Printf.printf "Exit code %d\n" n)
  cmd;
  Printf.printf "stdout: %d chars, stderr: %d chars\n" !stdout_chars !stderr_chars
*)

module Old = struct

  type closer = Strict | Normal | Lax

  (*
    let strict_closer chan =
    match Unix.close_process_in chan with
    Unix.WEXITED 0 -> ()
    | Unix.WEXITED n ->
    raise (Failure ("Exit code " ^ (string_of_int n)))
    | Unix.WSTOPPED n ->
    raise (Failure ("Stopped: " ^ (string_of_int n)))
    | Unix.WSIGNALED n ->
    raise (Failure ("Signaled: " ^ (string_of_int n)))

    let closer chan =
    match Unix.close_process_in chan with
    Unix.WEXITED n -> ()
    | Unix.WSTOPPED n ->
    raise (Failure ("Stopped: " ^ (string_of_int n)))
    | Unix.WSIGNALED n ->
    raise (Failure ("Signaled: " ^ (string_of_int n)))

    let lax_closer chan = ignore (Unix.close_process_in chan)  
  *)

  (** Return a string containing a command's output *)
  (* FIXME: the closer mechanism is awkward, default closer ignores exit code. *)
  let string ?(closer=Normal) cmd =
    match closer, to_string cmd with
      _, (s, Exit 0) -> s
    | Strict, (s, Exit n) -> failwith ("Exit code" ^ (string_of_int n))
    | _, (s, Exit n) -> s
    | Lax, (s, _) -> s
    | _, (s, Stopped n) -> failwith ("Stopped: " ^ (string_of_int n))
    | _, (s, Signaled n) -> failwith ("Signaled: " ^ (string_of_int n))
    | _, (s, Open) -> failwith ("Open")
    | _, (s, Stdout _) -> failwith ("internal")
    | _, (s, Stderr _) -> failwith ("internal")

end