open Core
open Qputils

(* Environment variables :

   QP_PREFIX=gdb   : to run gdb (or valgrind, or whatever)
   QP_MPIRUN=mpirun: to run mpi slaves
   QP_TASK_DEBUG=1 : debug task server

*)

let print_list () =
  Lazy.force Qpackage.executables 
  |> List.iter ~f:(fun (x,_) -> Printf.printf " * %s\n" x) 

let () = 
  Random.self_init ()

let run slave mpi exe ezfio_file =


  (** Check availability of the ports *)
  let port_number = 
    let zmq_context =
      ZMQ.Context.create ()
    in
    let dummy_socket = 
      ZMQ.Socket.create zmq_context ZMQ.Socket.rep
    in
    let rec try_new_port port_number =
      try 
        List.iter [ 0;1;2;3;4;5;6;7;8;9 ] ~f:(fun i ->
            let address = 
              Printf.sprintf "tcp://%s:%d" (Lazy.force TaskServer.ip_address) (port_number+i)
            in
            ZMQ.Socket.bind dummy_socket address;
            ZMQ.Socket.unbind dummy_socket address;
        );
        port_number
      with
      | Unix.Unix_error _ -> try_new_port (port_number+100)
    in
    let result = 
      try_new_port 41279
    in
    ZMQ.Socket.close dummy_socket;
    ZMQ.Context.terminate zmq_context;
    result
  in
  let time_start = 
    Time.now ()
  in

  if (not (Sys.file_exists_exn ezfio_file)) then
    failwith ("EZFIO directory "^ezfio_file^" not found");

  let executables = Lazy.force Qpackage.executables in
  if (not (List.exists ~f:(fun (x,_) -> x = exe) executables)) then
    begin
        Printf.printf "\nPossible choices:\n";
        List.iter executables ~f:(fun (x,_) -> Printf.printf "* %s\n%!" x);
        failwith ("Executable "^exe^" not found")
    end;

  Printf.printf "%s\n" (Time.to_string time_start);
  Printf.printf "===============\nQuantum Package\n===============\n\n";
  Printf.printf "Git Commit: %s\n" Git.message;
  Printf.printf "Git Date  : %s\n" Git.date;
  Printf.printf "Git SHA1  : %s\n" Git.sha1;
  Printf.printf "\n\n%!";


  (** Check input *)
  begin
    match (Sys.command ("qp_edit -c "^ezfio_file)) with
    | 0 -> ()
    | i -> failwith "Error: Input inconsistent\n"
  end;

  let qp_run_address_filename = 
   Filename.concat (Qpackage.ezfio_work ezfio_file) "qp_run_address"
  in

  let () = 
    if slave then
      try
        let address = 
          In_channel.read_all qp_run_address_filename
          |> String.strip
        in
        Unix.putenv ~key:"QP_RUN_ADDRESS_MASTER" ~data:address
      with Sys_error _ -> failwith "No master is not running"
  in
       
  (** Start task server *)
  let task_thread =
     let thread = 
      Thread.create ( fun () -> 
         TaskServer.run port_number )
     in
     thread ();
  in
  let address =
    Printf.sprintf "tcp://%s:%d" (Lazy.force TaskServer.ip_address) port_number
  in
  Unix.putenv ~key:"QP_RUN_ADDRESS" ~data:address;
  let () = 
    if (not slave) then
      Out_channel.with_file qp_run_address_filename  ~f:(
        fun oc -> Out_channel.output_lines oc [address]) 
  in


  (** Run executable *)
  let prefix = 
    match Sys.getenv "QP_PREFIX" with
    | Some x -> x^" "
    | None -> ""
  and mpirun = 
    match (mpi, Sys.getenv "QP_MPIRUN") with
    | (true, None) -> "mpirun "
    | (true, Some x) -> x^" "
    | _ -> ""
  and exe =
    match (List.find ~f:(fun (x,_) -> x = exe) executables) with
    | Some (_,x) -> x^" "
    | None -> assert false
  in
  let exit_code = 
    match (Sys.command (mpirun^prefix^exe^ezfio_file)) with
    | 0 -> 0
    | i -> (Printf.printf "Program exited with code %d.\n%!" i; i)
  in

  TaskServer.stop ~port:port_number;
  Thread.join task_thread;
  if (not slave) then
    Sys.remove qp_run_address_filename;

  let duration = Time.diff (Time.now()) time_start 
  |> Time.Span.to_string in
  Printf.printf "Wall time : %s\n\n" duration;
  if (exit_code <> 0) then
    exit exit_code

let spec = 
  let open Command.Spec in
  empty
  +> flag "slave" no_arg
     ~doc:(" Required for slave tasks")
  +> flag "mpi" no_arg
     ~doc:(" Required for MPI slaves")
  +> anon ("executable" %: string)
  +> anon ("ezfio_file" %: string)
;;



let () =
  Command.basic
  ~summary: "Quantum Package command"
  ~readme:( fun () -> "
Executes a Quantum Package binary file among these:\n\n"
^ (Lazy.force Qpackage.executables
    |> List.map ~f:(fun (x,_) -> Printf.sprintf " * %s" x )
    |> String.concat ~sep:"\n" 
    )
  )
  spec
  (fun slave mpi exe ezfio_file () ->
    run slave mpi exe ezfio_file
  )
  |> Command.run   ~version: Git.sha1   ~build_info: Git.message