(*
  Copyright (C) 2019 TG x Thoth

  This program is free software: you can redistribute it and/or modify
  it under the terms of the GNU Affero General Public License as published by
  the Free Software Foundation, either version 3 of the License.

  This program is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  GNU Affero General Public License for more details.

  You should have received a copy of the GNU Affero General Public License
  along with this program.  If not, see <https://www.gnu.org/licenses/>.
*)

open Stdint
open OUnit2

let u64 = Uint64.of_int64
let pf = Fmt.pf
let out = Fmt.stdout
let sent_msgs = Hashtbl.create 3

module Node_id = P2p.Node_id
module Node = P2p.Node.Make (Node_id)
module View = P2p.View.Make (Node_id) (Node)
module Msg_id = P2p.Msg_id
module Ringcast = P2p_ringcast.Make (Node_id) (Node) (View) (Msg_id)

module Io = struct

  type t = {
      node_id: Node_id.t;
      in_chan : Lwt_io.input_channel;
      out_chan : Lwt_io.output_channel;
    }

  let init node_id in_chan out_chan =
    { node_id; in_chan; out_chan }

  (** [initiate_gossip t node xchg]
      sends [xchg] entries to node [dst]
      and returns response *)
  let initiate_gossip t dst xchg =
    pf out "%a # INITIATE_GOSSIP to node %a\n" Node_id.pp t.node_id Node.pp dst;
    pf out "xchg to send:\n%a\n" View.pp xchg;
    flush stdout;
    let%lwt _ = Lwt_io.write_value t.out_chan xchg in
    Lwt_io.read_value t.in_chan

  (** [respond_gossip t node xchg]
      sends [xchg] entries in response to node [dst] *)
  let respond_gossip t dst xchg =
    pf out "%a # RESPOND_GOSSIP to node %a\n" Node_id.pp t.node_id Node.pp dst;
    pf out "xchg to send:\n%a\n" View.pp xchg;
    flush stdout;
    let%lwt _ = Lwt_io.write_value t.out_chan xchg in
    Lwt.return_unit

  (** [gossip_recvd t node view recvd]
      is called after entries are received during a gossip exchange;
      allows rewriting [recvd] entries with the returned value. *)
  let gossip_recvd _t _src recvd _view =
    Lwt.return recvd

  (** [view_updated node view]
      is called when [view] has been updated after a gossip exchange *)
  let view_updated t node view =
    pf out "%a # VIEW_UPDATED of node %a\n%a\n" Node_id.pp t.node_id Node.pp node View.pp view;
    flush stdout;
    Lwt.return_unit

  let send_msg t dst mid msg =
    pf out "%a # SEND_MSG %a to node %a%s\n"
      Node_id.pp t.node_id
      Msg_id.pp mid
      Node.pp dst
      (Cstruct.to_string msg);
    flush stdout;
    let n =
      match Hashtbl.find_opt sent_msgs mid with
      | Some n -> n
      | None -> 0 in
    Hashtbl.add sent_msgs mid (n + 1);
    Lwt.return_unit

  let get_xview _t =
    View.empty
end

module Ringcast_lwt =
  P2p_ringcast_lwt.Make (Node_id) (Node) (View) (Msg_id) (Ringcast) (Io)

let rec read_chan ch rc node rnode =
  let%lwt recvd = Lwt_io.read_value ch in
  pf out "%a # READ_CHAN\n" Node_id.pp (Node.id node);
  pf out "recvd:\n%a\n" View.pp recvd;
  flush stdout;
  let%lwt view = Ringcast_lwt.respond rc rnode recvd in
  pf out "recvd:\n%a\n" View.pp recvd;
  pf out "view:\n%a\n" View.pp view;
  flush stdout;
  read_chan ch rc node rnode

let forward delay rc src msgid msg =
  let%lwt _ = Lwt_unix.sleep delay in
  Ringcast_lwt.forward rc src msgid (Cstruct.of_string msg)

let check_sent_msgs delay =
  let%lwt _ = Lwt_unix.sleep delay in
  let s1, s2, s3 =
    Hashtbl.find sent_msgs (Uint64.of_int 1),
    Hashtbl.find sent_msgs (Uint64.of_int 2),
    Hashtbl.find sent_msgs (Uint64.of_int 3) in
  let%lwt _ = Lwt_io.printf "check_recvd_msgs: %d %d %d\n"
                s1 s2 s3 in
  assert_equal s1 3;
  assert_equal s2 3;
  assert_equal s3 4;
  Lwt.return_unit

let _ = Nocrypto_entropy_lwt.initialize ()

let () =
  let view_len = 8 in
  let xchg_len = 4 in
  let period = 1.0 in
  let fanout = 5 in
  let seen_len = 10 in

  let (in_ch1, out_ch2) = Lwt_io.pipe () in
  let (in_ch2, out_ch1) = Lwt_io.pipe () in

  let node1 = Node.init (u64 100L) in
  let io1 = Io.init (Node.id node1) in_ch1 out_ch1 in
  let view1 =
    View.add (Node.init (u64 110L))
      (View.add (Node.init (u64 120L))
         (View.add (Node.init (u64 130L))
            (View.add (Node.init (u64 140L))
               (View.add (Node.init (u64 150L))
                  (View.add (Node.init (u64 160L))
                     (View.add (Node.init (u64 170L))
                        View.empty)))))) in
  let rc1 = Ringcast_lwt.init ~me:node1 ~view:view1 ~view_len ~xchg_len
              ~period ~fanout ~seen_len ~io:io1 in

  let node2 = Node.init (u64 200L) in
  let io2 = Io.init (Node.id node2) in_ch2 out_ch2 in
  let view2 =
    View.add (Node.init (u64 210L))
      (View.add (Node.init (u64 220L))
         (View.add (Node.init (u64 230L))
            (View.add (Node.init (u64 240L))
               (View.add (Node.init (u64 250L))
                  (View.add (Node.init (u64 260L))
                     (View.add (Node.init (u64 270L))
                        View.empty)))))) in
  let rc2 = Ringcast_lwt.init ~me:node2 ~view:view2 ~view_len ~xchg_len ~period ~fanout ~seen_len ~io:io2 in

  let timeout = Lwt_unix.sleep 5.5 in
  Lwt_main.run @@
    Lwt.join [
        Lwt.choose [ Ringcast_lwt.run rc1;
                     Ringcast_lwt.run rc2;
                     read_chan in_ch1 rc1 node1 node2;
                     read_chan in_ch2 rc2 node2 node1;
                     Lwt.map (fun () -> Ringcast_lwt.shutdown rc1;
                                        Ringcast_lwt.shutdown rc2) timeout ];
        forward 1.0 rc1 node1 (Uint64.of_int 1)
          (Printf.sprintf "first message from %s" (Node.to_string node1));
        forward 1.5 rc1 node2 (Uint64.of_int 1)
          (Printf.sprintf "first message #2 from %s" (Node.to_string node2));
        forward 2.0 rc1 node1 (Uint64.of_int 2)
          (Printf.sprintf "second message from %s" (Node.to_string node1));
        forward 3.0 rc1 node2 (Uint64.of_int 3)
          (Printf.sprintf "third message from %s" (Node.to_string node2));
        forward 4.0 rc1 node2 (Uint64.of_int 1)
          (Printf.sprintf "first message #3 from %s" (Node.to_string node2));
        check_sent_msgs 5.3
      ]
