open Swim.Types
module Cluster = Swim.Cluster

external env_cast : 'a -> 'b = "%identity"

let make_config ~port ~name =
  {
    default_config with
    bind_addr = "\127\000\000\001";
    bind_port = port;
    node_name = Some name;
    protocol_interval = 0.1;
    probe_timeout = 0.05;
    suspicion_mult = 2;
    secret_key = String.make 16 'k';
    cluster_name = "test-cluster";
  }

let test_cluster_create_start_shutdown sw env () =
  let config = make_config ~port:17946 ~name:"test-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error `Invalid_key -> Alcotest.fail "invalid key"
  | Ok cluster ->
      Cluster.start cluster;
      let local = Cluster.local_node cluster in
      Alcotest.(check string)
        "node name" "test-node"
        (node_id_to_string local.id);
      Eio.Time.sleep env#clock 0.05;
      Cluster.shutdown cluster

let test_cluster_stats sw env () =
  let config = make_config ~port:17947 ~name:"stats-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      Eio.Time.sleep env#clock 0.05;
      let stats = Cluster.stats cluster in
      Alcotest.(check int) "initial nodes_alive" 0 stats.nodes_alive;
      Cluster.shutdown cluster

let test_cluster_add_member sw env () =
  let config = make_config ~port:17948 ~name:"add-member-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      let fake_node =
        make_node_info
          ~id:(node_id_of_string "fake-node")
          ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\002", 7946))
          ~meta:"test-meta"
      in
      Cluster.add_member cluster fake_node;
      Eio.Time.sleep env#clock 0.05;
      Alcotest.(check int) "member count" 1 (Cluster.member_count cluster);
      Cluster.shutdown cluster

let test_cluster_remove_member sw env () =
  let config = make_config ~port:17949 ~name:"remove-member-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      let fake_node =
        make_node_info
          ~id:(node_id_of_string "fake-node")
          ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\002", 7946))
          ~meta:"test-meta"
      in
      Cluster.add_member cluster fake_node;
      Eio.Time.sleep env#clock 0.02;
      Alcotest.(check int) "before remove" 1 (Cluster.member_count cluster);
      let removed =
        Cluster.remove_member cluster (node_id_of_string "fake-node")
      in
      Alcotest.(check bool) "removed" true removed;
      Alcotest.(check int) "after remove" 0 (Cluster.member_count cluster);
      Cluster.shutdown cluster

let test_cluster_members_list sw env () =
  let config = make_config ~port:17950 ~name:"members-list-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      for i = 1 to 3 do
        let node =
          make_node_info
            ~id:(node_id_of_string (Printf.sprintf "node-%d" i))
            ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", 7946 + i))
            ~meta:""
        in
        Cluster.add_member cluster node
      done;
      Eio.Time.sleep env#clock 0.02;
      let members = Cluster.members cluster in
      Alcotest.(check int) "members count" 3 (List.length members);
      Cluster.shutdown cluster

let test_cluster_is_healthy sw env () =
  let config = make_config ~port:17951 ~name:"health-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      Alcotest.(check bool)
        "unhealthy without members" false
        (Cluster.is_healthy cluster);
      Cluster.shutdown cluster

let test_cluster_broadcast sw env () =
  let config = make_config ~port:17952 ~name:"broadcast-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      Cluster.broadcast cluster ~topic:"test-topic" ~payload:"hello world";
      Eio.Time.sleep env#clock 0.02;
      Cluster.shutdown cluster

let test_cluster_on_message sw env () =
  let config = make_config ~port:17953 ~name:"message-handler-node" in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error _ -> Alcotest.fail "create failed"
  | Ok cluster ->
      Cluster.start cluster;
      let received = ref false in
      Cluster.on_message cluster (fun _topic _payload _origin ->
          received := true);
      Eio.Time.sleep env#clock 0.02;
      Cluster.shutdown cluster

let test_cluster_invalid_key sw env () =
  let config =
    { (make_config ~port:17954 ~name:"bad-key") with secret_key = "short" }
  in
  let env_wrap = { stdenv = env; sw } in
  match Cluster.create ~sw ~env:env_wrap ~config with
  | Error `Invalid_key -> ()
  | Ok _ -> Alcotest.fail "expected invalid key error"

let () =
  Eio_main.run @@ fun env ->
  let env = env_cast env in
  Eio.Switch.run @@ fun sw ->
  Alcotest.run "integration"
    [
      ( "cluster",
        [
          ( "create_start_shutdown",
            `Quick,
            test_cluster_create_start_shutdown sw env );
          ("stats", `Quick, test_cluster_stats sw env);
          ("add_member", `Quick, test_cluster_add_member sw env);
          ("remove_member", `Quick, test_cluster_remove_member sw env);
          ("members_list", `Quick, test_cluster_members_list sw env);
          ("is_healthy", `Quick, test_cluster_is_healthy sw env);
          ("broadcast", `Quick, test_cluster_broadcast sw env);
          ("on_message", `Quick, test_cluster_on_message sw env);
          ("invalid_key", `Quick, test_cluster_invalid_key sw env);
        ] );
    ]