open Core
open Qptypes

module StringHashtbl = Hashtbl.Make(String)

type pub_state =
| Waiting
| Running of string
| Stopped

let pub_state_of_string = function
| "Waiting"  -> Waiting
| "Stopped"  -> Stopped
| s -> Running s

let string_of_pub_state = function
| Waiting  -> "Waiting"
| Stopped  -> "Stopped"
| Running s -> s



type t =
{
    queue             : Queuing_system.t ;
    state             : Message.State.t option ;
    address_tcp       : Address.Tcp.t option ; 
    address_inproc    : Address.Inproc.t option ;
    progress_bar      : Progress_bar.t option ;
    running           : bool;
    accepting_clients : bool;
    data              : string StringHashtbl.t;
}



let debug_env =
  match Sys.getenv "QP_TASK_DEBUG" with
  | Some x -> x <> ""
  | None   -> false


let debug str =
  if debug_env then
    Printf.eprintf "TASK : %s%!" str



let zmq_context =
  ZMQ.Context.create ()

let () =
  ZMQ.Context.set_io_threads zmq_context 8


let bind_socket ~socket_type ~socket ~port =
  let rec loop = function
  | 0 -> failwith @@ Printf.sprintf
        "Unable to bind the %s socket to port : %d "
        socket_type port
  | -1 -> ()
  | i -> 
      try
        ZMQ.Socket.bind socket @@ Printf.sprintf "tcp://*:%d" port;
        loop (-1)
      with
      | Unix.Unix_error _ -> (Time.pause @@ Time.Span.of_sec 1. ; loop (i-1) )
      | other_exception -> raise other_exception
  in loop 60


let hostname = lazy (
    try
      Unix.gethostname ()
    with
    | _ -> "localhost"
)


let ip_address = lazy (
  match Sys.getenv "QP_NIC" with
  | None ->
      begin
        try
          Lazy.force hostname
          |> Unix.Inet_addr.of_string_or_getbyname
          |> Unix.Inet_addr.to_string
        with
        | Unix.Unix_error _ ->
            failwith "Unable to find IP address from host name."
      end
  | Some interface ->
      begin
        try
          ok_exn Linux_ext.get_ipv4_address_for_interface interface
        with
        | Unix.Unix_error _ ->
            Lazy.force hostname
            |> Unix.Inet_addr.of_string_or_getbyname
            |> Unix.Inet_addr.to_string
      end
)


let reply_ok rep_socket =
    Message.Ok_msg.create 
    |> Message.Ok_msg.to_string
    |> ZMQ.Socket.send rep_socket 

let reply_wrong_state rep_socket =
    Message.Error_msg.create "Wrong state"
    |> Message.Error_msg.to_string
    |> ZMQ.Socket.send rep_socket 



let stop ~port =
    debug "STOP";
    let req_socket =
      ZMQ.Socket.create zmq_context ZMQ.Socket.req
    and address =
      Printf.sprintf "tcp://localhost:%d" port
    in
    ZMQ.Socket.set_linger_period req_socket 1_000_000;
    ZMQ.Socket.connect req_socket address;

    Message.Terminate (Message.Terminate_msg.create)
    |> Message.to_string
    |> ZMQ.Socket.send req_socket ;

    let msg =
      ZMQ.Socket.recv req_socket
      |> Message.of_string
    in
    let () =
      match msg with
      | Message.Ok _ -> ()
      | _ -> failwith "Problem in termination"
    in
    ZMQ.Socket.set_linger_period req_socket 1_000;
    ZMQ.Socket.close req_socket


let new_job msg program_state rep_socket pair_socket =

    let state = 
        msg.Message.Newjob_msg.state
    in

    let progress_bar =
        Progress_bar.init
          ~start_value:0.
          ~end_value:1.
          ~bar_length:20
          ~title:(Message.State.to_string state)
    in

    let result = 
      { program_state with
        state           = Some state ;
        progress_bar    = Some progress_bar ;
        address_tcp     = Some msg.Message.Newjob_msg.address_tcp;
        address_inproc  = Some msg.Message.Newjob_msg.address_inproc;
        accepting_clients = true;
      }
    in
    reply_ok rep_socket;
    string_of_pub_state Waiting
    |> ZMQ.Socket.send pair_socket ;
    result

let change_pub_state msg program_state rep_socket pair_socket =
  let msg = 
    match msg with
    | `Waiting -> Waiting 
    | `Stopped -> Stopped
    | `Running ->
      begin
	let state =
	   match program_state.state with
	   | Some x -> x
	   | None  -> failwith "Trying to change pub state while no job is ready"
        in
        Running (Message.State.to_string state)
      end
  in
  reply_ok rep_socket;
  string_of_pub_state msg
  |> ZMQ.Socket.send pair_socket ;

  program_state

let force_state = 
  Message.State.of_string "force"

let end_job msg program_state rep_socket pair_socket =

    let failure () =
        reply_wrong_state rep_socket;
        program_state

    and success () =
        reply_ok rep_socket;
        { 
          queue             = Queuing_system.create ();
          state             = None ;
          progress_bar      = Progress_bar.clear ();
          address_tcp       = None;
          address_inproc    = None;
          running           = true;
          accepting_clients = false;
          data = StringHashtbl.create ();
        }

    and wait n =
      Printf.sprintf "waiting for %d slaves..." n
      |> Message.Error_msg.create 
      |> Message.Error_msg.to_string
      |> ZMQ.Socket.send rep_socket ;
      program_state
    in

    match program_state.state with
    | None -> failure ()
    | Some state -> 
      begin
        if (msg.Message.Endjob_msg.state = force_state) then
          begin
            string_of_pub_state Waiting
            |> ZMQ.Socket.send pair_socket ;
              success ()
          end
        else if (msg.Message.Endjob_msg.state = state) then
          begin
            string_of_pub_state Waiting
            |> ZMQ.Socket.send pair_socket ;
            if (Queuing_system.number_of_clients program_state.queue = 0) then
              success ()
            else
              wait (Queuing_system.number_of_clients program_state.queue)
          end
        else
          failure ()
      end


let connect msg program_state rep_socket =

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    in

    if (not program_state.accepting_clients) then
      failure ()
    else
      match program_state.state with
      | None -> failure ()
      | Some state -> 
          let push_address =
              match msg with
              | Message.Connect_msg.Tcp    -> 
                begin
                    match program_state.address_tcp  with
                    | Some address -> Address.Tcp address
                    | None -> failwith "Error: No TCP address"
                end
              | Message.Connect_msg.Inproc -> 
                begin
                    match program_state.address_inproc with
                    | Some address -> Address.Inproc address
                    | None -> failwith "Error: No inproc address"
                end
              | Message.Connect_msg.Ipc    -> assert false
          in
          
          let new_queue, client_id =
              Queuing_system.add_client program_state.queue
          in
          Message.ConnectReply (Message.ConnectReply_msg.create
              ~state:state ~client_id ~push_address)
          |> Message.to_string
          |> ZMQ.Socket.send rep_socket ;
          { program_state with
            queue = new_queue 
          }


let disconnect msg program_state rep_socket =

    let state, client_id =
      msg.Message.Disconnect_msg.state, 
      msg.Message.Disconnect_msg.client_id
    in

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    
    and success () = 

        let new_program_state = 
            { program_state with
              queue = Queuing_system.del_client ~client_id program_state.queue
            }
        in
        Message.DisconnectReply (Message.DisconnectReply_msg.create ~state)
        |> Message.to_string
        |> ZMQ.Socket.send rep_socket ;
        new_program_state

    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
      begin
        if (state = state') then
          success ()
        else
          failure ()
      end
    
let del_task msg program_state rep_socket =

    let state, task_ids =
      msg.Message.DelTask_msg.state, 
      msg.Message.DelTask_msg.task_ids
    in

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    
    and success () = 

        let queue = 
           List.fold ~f:(fun queue task_id -> Queuing_system.del_task ~task_id queue)
                      ~init:program_state.queue task_ids
        in
        let accepting_clients = 
            (Queuing_system.number_of_queued queue > Queuing_system.number_of_clients queue) 
        in
        let new_program_state = 
            { program_state with
              accepting_clients ;
              queue ;
            }
        in
        let more = 
            (Queuing_system.number_of_tasks queue > 0)
        in
        Message.DelTaskReply (Message.DelTaskReply_msg.create ~task_ids ~more)
        |> Message.to_string
        |> ZMQ.Socket.send ~block:true rep_socket ; (** /!\ Has to be blocking *)
        new_program_state

    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
      begin
        if (state = state') then
          success ()
        else
          failure ()
      end



let add_task msg program_state rep_socket =

    let state, tasks =
        msg.Message.AddTask_msg.state,
        msg.Message.AddTask_msg.tasks
    in

    let increment_progress_bar = function
      | Some bar -> Some (Progress_bar.increment_end bar)
      | None -> None
    in

    let result = 
        let new_queue, new_bar = 
          List.fold ~f:(fun (queue, bar) task ->
              Queuing_system.add_task ~task queue,
              increment_progress_bar bar)
            ~init:(program_state.queue, program_state.progress_bar) tasks 
        in
        { program_state with
          queue = new_queue;
          progress_bar = new_bar
        }
    in
    reply_ok rep_socket;
    result



let get_task msg program_state rep_socket pair_socket =

    let state, client_id =
      msg.Message.GetTask_msg.state, 
      msg.Message.GetTask_msg.client_id
    in

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    
    and success () = 

        let queue, task_id, task =
            Queuing_system.pop_task ~client_id program_state.queue
        in

        let accepting_clients = 
            (Queuing_system.number_of_queued queue >
             Queuing_system.number_of_clients queue) 
        in

        let no_task =
          Queuing_system.number_of_queued queue = 0
        in

        if no_task then
          string_of_pub_state Waiting 
          |> ZMQ.Socket.send pair_socket
        else
          string_of_pub_state (Running (Message.State.to_string state))
          |> ZMQ.Socket.send pair_socket;

        let new_program_state = 
            { program_state with
              queue ;
              accepting_clients;
            }
        in

        Message.GetTaskReply (Message.GetTaskReply_msg.create ~task ~task_id)
        |> Message.to_string
        |> ZMQ.Socket.send rep_socket ;
        new_program_state

    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
      begin
        if (state = state') then
          success ()
        else
          failure ()
      end



let get_tasks msg program_state rep_socket pair_socket =

    let state, client_id, n_tasks =
      msg.Message.GetTasks_msg.state, 
      msg.Message.GetTasks_msg.client_id,
      Strictly_positive_int.to_int msg.Message.GetTasks_msg.n_tasks
    in

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    
    and success () = 

        let rec build_list accu queue = function
        | 0 -> queue, (List.rev accu)
        | n -> 
            let new_queue, task_id, task =
               Queuing_system.pop_task ~client_id queue
            in
            match (task_id, task) with
            | Some task_id, Some task -> 
              build_list ( (Some task_id, task)::accu ) new_queue (n-1)
            | _ -> build_list ( (None, "terminate")::accu ) queue 0
        in

        let new_queue, result =
          build_list [] program_state.queue (n_tasks)
        in

        let no_task =
          Queuing_system.number_of_queued new_queue = 0
        in

        let accepting_clients = 
            (Queuing_system.number_of_queued new_queue >
             Queuing_system.number_of_clients new_queue) 
        in

        if no_task then
          string_of_pub_state Waiting 
          |> ZMQ.Socket.send pair_socket
        else
          string_of_pub_state (Running (Message.State.to_string state))
          |> ZMQ.Socket.send pair_socket;
        
        let new_program_state = 
            { program_state with
              queue = new_queue;
              accepting_clients;
            }
        in

        Message.GetTasksReply (Message.GetTasksReply_msg.create result)
        |> Message.to_string_list
        |> ZMQ.Socket.send_all rep_socket ;
        new_program_state
    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
      begin
        if (state = state') then
          success ()
        else
          failure ()
      end



let task_done msg program_state rep_socket =

    let state, client_id, task_ids =
        msg.Message.TaskDone_msg.state,
        msg.Message.TaskDone_msg.client_id,
        msg.Message.TaskDone_msg.task_ids
    in

    let increment_progress_bar = function
      | Some bar -> Some (Progress_bar.increment_cur bar)
      | None -> None
    in

    let failure () =
        reply_wrong_state rep_socket;
        program_state
    
    and success () = 
        let new_queue, new_bar = 
          List.fold ~f:(fun (queue, bar) task_id -> 
              Queuing_system.end_task ~task_id ~client_id queue,
              increment_progress_bar bar)
            ~init:(program_state.queue, program_state.progress_bar) task_ids 
        in

        let accepting_clients = 
            (Queuing_system.number_of_queued new_queue >
             Queuing_system.number_of_clients new_queue) 
        in

        let result = 
          { program_state with
            queue = new_queue;
            progress_bar = new_bar;
            accepting_clients
          }
        in
        reply_ok rep_socket;
        result
    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
      begin
        if (state = state') then
          success ()
        else
          failure ()
      end



let put_data msg rest_of_msg program_state rep_socket =

    debug (Message.PutData_msg.to_string msg);
    let state, key, value =
      msg.Message.PutData_msg.state,
      msg.Message.PutData_msg.key, 
      match rest_of_msg with
      | [ x ] -> x
      | _ -> failwith "Badly formed put_data message"
    in

    let success ()  =
      StringHashtbl.set program_state.data ~key ~data:value ;
      Message.PutDataReply (Message.PutDataReply_msg.create ())
      |> Message.to_string
      |> ZMQ.Socket.send rep_socket;
      program_state

    and failure () =
        reply_wrong_state rep_socket;
        program_state
    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
        if (state = state') then
          success ()
        else
          failure ()


let get_data msg program_state rep_socket =

    debug (Message.GetData_msg.to_string msg);
    let state, key =
        msg.Message.GetData_msg.state,
        msg.Message.GetData_msg.key
    in

    let success () = 
      let value = 
        match StringHashtbl.find program_state.data key with
        | Some value -> value
        | None -> ""
      in
      Message.GetDataReply (Message.GetDataReply_msg.create ~value)
      |> Message.to_string_list
      |> ZMQ.Socket.send_all rep_socket;
      program_state

    and failure () =
        reply_wrong_state rep_socket;
        program_state
    in

    match program_state.state with
    | None -> assert false
    | Some state' -> 
        if (state = state') then
          success ()
        else
          failure ()


let terminate program_state rep_socket =
    reply_ok rep_socket;
    { program_state with
      address_tcp = None;
      address_inproc = None;
      running = false
    }

  
let abort program_state rep_socket =
    let queue, client_id =
        Queuing_system.add_client program_state.queue
    in
    let rec aux accu queue = function
    | 0 -> (queue, accu)
    | rest ->
      let new_queue, task_id, _ =
        Queuing_system.pop_task ~client_id queue
      in
      let new_accu = 
        match task_id with
        | Some task_id -> task_id::accu
        | None -> accu
      in
      Queuing_system.number_of_queued new_queue
      |> aux new_accu new_queue
    in
    let queue, tasks =
      aux [] queue 1
    in
    let queue = 
      List.fold ~f:(fun queue task_id -> 
                  Queuing_system.end_task ~task_id ~client_id queue)
                  ~init:queue tasks
    in
    let queue = 
      List.fold ~f:(fun queue task_id -> Queuing_system.del_task ~task_id queue)
                          ~init:queue tasks
    in
    let queue =
      Queuing_system.del_client ~client_id  queue
    in
    reply_ok rep_socket;

    { program_state with
      queue ; 
      accepting_clients = false;
    }

  
let error msg program_state rep_socket =
    Message.Error (Message.Error_msg.create msg)
    |> Message.to_string
    |> ZMQ.Socket.send rep_socket ;
    program_state

let start_pub_thread ~port = 
    Thread.create (fun () ->
      let timeout =
        1000
      in

      let pair_socket =
        ZMQ.Socket.create zmq_context ZMQ.Socket.pair
      and address = 
        "inproc://pair" 
      in
      ZMQ.Socket.connect pair_socket address;

      let pub_socket =
        ZMQ.Socket.create zmq_context ZMQ.Socket.pub
      in
      bind_socket ~socket_type:"PUB" ~socket:pub_socket ~port;

      let pollitem =
        ZMQ.Poll.mask_of 
          [| (pair_socket, ZMQ.Poll.In) |]
      in

      let rec run state = 
        let new_state =
          let polling = 
            ZMQ.Poll.poll ~timeout pollitem
          in
          if (polling.(0) = Some ZMQ.Poll.In) then
            ZMQ.Socket.recv ~block:false pair_socket
            |> pub_state_of_string
          else
            state
        in
        ZMQ.Socket.send pub_socket @@ string_of_pub_state new_state;
        match state with
        | Stopped -> ()
        | _ -> run new_state
      in
      run Waiting;
      ZMQ.Socket.set_linger_period pair_socket 1000 ;
      ZMQ.Socket.close pair_socket;
      ZMQ.Socket.set_linger_period pub_socket 1000 ;
      ZMQ.Socket.close pub_socket;
    )

let run ~port =

    (** Bind inproc socket for changing state of pub *)
    let pair_socket =
      ZMQ.Socket.create zmq_context ZMQ.Socket.pair
    and address = 
      "inproc://pair" 
    in
    ZMQ.Socket.bind pair_socket address;

    let pub_thread = 
      start_pub_thread ~port:(port+1) ()
    in

    (** Bind REP socket *)
    let rep_socket =
      ZMQ.Socket.create zmq_context ZMQ.Socket.rep
    in
    ZMQ.Socket.set_linger_period rep_socket 1_000_000;
    bind_socket "REP" rep_socket port;

    let initial_program_state =
    {   queue = Queuing_system.create () ;
        running = true ;
        state = None;
        address_tcp = None;
        address_inproc = None;
        progress_bar = None ;
        accepting_clients = false;
        data = StringHashtbl.create ();
    }
    in

    (** ZMR polling item *)
    let pollitem =
      ZMQ.Poll.mask_of
        [| (rep_socket, ZMQ.Poll.In) |]
    in

    let address =
      Printf.sprintf "tcp://%s:%d" (Lazy.force ip_address) port
    in
    Printf.printf "Task server running : %s\n%!" address;


    (** Main loop *)
    let rec main_loop program_state = function
    | false -> ()
    | true ->
        let polling =
            ZMQ.Poll.poll ~timeout:1000 pollitem
        in
        if (polling.(0) <> Some ZMQ.Poll.In) then
          main_loop program_state true
        else
          begin
              let program_state = 
                  match program_state.progress_bar with
                  | None -> program_state
                  | Some bar -> 
                    if bar.Progress_bar.dirty then 
                        { program_state with
                          progress_bar = Some (Progress_bar.display bar)
                        }
                    else
                        program_state
              in

              (** Extract message *)
              let raw_message, rest =
                  match ZMQ.Socket.recv_all rep_socket with
                  | x :: rest -> x, rest
                  | [] -> failwith "Badly formed message"
              in
              let message =
                  Message.of_string raw_message
              in

              (** Debug input *)
              let () = 
                if debug_env then
                  begin
                    Printf.sprintf "q:%d  r:%d  n:%d  c:%d : %s\n%!"
                    (Queuing_system.number_of_queued program_state.queue)
                    (Queuing_system.number_of_running program_state.queue)
                    (Queuing_system.number_of_tasks program_state.queue)
                    (Queuing_system.number_of_clients program_state.queue)
                    (Message.to_string message)
                    |> debug
                  end
              in

              let new_program_state = 
                try
                  match program_state.state, message with
                  | _     , Message.Terminate   _ -> terminate program_state rep_socket
                  | _     , Message.Abort       _ -> abort program_state rep_socket
                  | _     , Message.PutData     x -> put_data x rest program_state rep_socket
                  | _     , Message.GetData     x -> get_data x program_state rep_socket
                  | None  , Message.Newjob      x -> new_job x program_state rep_socket pair_socket
                  | _     , Message.Newjob      _ -> error "A job is already running" program_state rep_socket
                  | Some _, Message.Endjob      x -> end_job x program_state rep_socket pair_socket
                  | Some _, Message.SetRunning    -> change_pub_state `Running program_state rep_socket pair_socket
                  | _, Message.SetWaiting    -> change_pub_state `Waiting program_state rep_socket pair_socket
                  | _, Message.SetStopped    -> change_pub_state `Stopped program_state rep_socket pair_socket
                  | None  ,                     _ -> error "No job is running" program_state rep_socket
                  | Some _, Message.Connect     x -> connect x program_state rep_socket
                  | Some _, Message.Disconnect  x -> disconnect x program_state rep_socket
                  | Some _, Message.AddTask     x -> add_task x program_state rep_socket
                  | Some _, Message.DelTask     x -> del_task x program_state rep_socket
                  | Some _, Message.GetTask     x -> get_task x program_state rep_socket pair_socket
                  | Some _, Message.GetTasks    x -> get_tasks x program_state rep_socket pair_socket
                  | Some _, Message.TaskDone    x -> task_done x program_state rep_socket
                  | _     , _                     ->
                    error ("Invalid message : "^(Message.to_string message))  program_state rep_socket
                with
                | Failure f -> 
                    error (f^" : "^raw_message) program_state rep_socket
                | Assert_failure (f,i,j) ->
                    error (Printf.sprintf "%s:%d:%d : %s" f i j raw_message) program_state rep_socket

              in
              main_loop new_program_state new_program_state.running
          end
    in main_loop initial_program_state true;

    ZMQ.Socket.send pair_socket @@ string_of_pub_state Stopped;
    Thread.join pub_thread;
    ZMQ.Socket.close rep_socket