open Swim.Types
open Swim.Protocol_pure

let node1 =
  make_node_info
    ~id:(node_id_of_string "node1")
    ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", 7946))
    ~meta:""

let node2 =
  make_node_info
    ~id:(node_id_of_string "node2")
    ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\002", 7946))
    ~meta:""

let node3 =
  make_node_info
    ~id:(node_id_of_string "node3")
    ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\003", 7946))
    ~meta:""

let now = Mtime.Span.of_uint64_ns 0L
let alive_state : member_state = Alive
let suspect_state : member_state = Suspect
let dead_state : member_state = Dead

let make_member ?(state = alive_state) ?(incarnation = 0) node =
  {
    node;
    state;
    incarnation = incarnation_of_int incarnation;
    state_change = now;
  }

let test_alive_higher_incarnation_wins () =
  let member = make_member ~incarnation:1 node1 in
  let msg = Alive { node = node1; incarnation = incarnation_of_int 2 } in
  let result = handle_alive ~self:(node_id_of_string "self") member msg ~now in
  Alcotest.(check int)
    "incarnation" 2
    (incarnation_to_int result.new_state.incarnation);
  Alcotest.(check bool) "broadcast" true (List.length result.broadcasts = 1)

let test_alive_lower_incarnation_ignored () =
  let member = make_member ~incarnation:5 node1 in
  let msg = Alive { node = node1; incarnation = incarnation_of_int 3 } in
  let result = handle_alive ~self:(node_id_of_string "self") member msg ~now in
  Alcotest.(check int)
    "incarnation unchanged" 5
    (incarnation_to_int result.new_state.incarnation);
  Alcotest.(check bool) "no broadcast" true (List.length result.broadcasts = 0)

let test_alive_same_incarnation_unsuspects () =
  let member = make_member ~state:suspect_state ~incarnation:3 node1 in
  let msg = Alive { node = node1; incarnation = incarnation_of_int 3 } in
  let result = handle_alive ~self:(node_id_of_string "self") member msg ~now in
  Alcotest.(check string)
    "state alive" "alive"
    (member_state_to_string result.new_state.state)

let test_alive_revives_dead_node () =
  let member = make_member ~state:dead_state ~incarnation:1 node1 in
  let msg = Alive { node = node1; incarnation = incarnation_of_int 5 } in
  let result = handle_alive ~self:(node_id_of_string "self") member msg ~now in
  Alcotest.(check string)
    "state alive" "alive"
    (member_state_to_string result.new_state.state);
  match result.events with
  | [ Join _ ] -> ()
  | _ -> Alcotest.fail "expected Join event"

let test_suspect_triggers_refute_for_self () =
  let member = make_member ~incarnation:1 node1 in
  let msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 1;
        suspector = node_id_of_string "node2";
      }
  in
  let result =
    handle_suspect ~self:(node_id_of_string "node1") member msg ~now
  in
  Alcotest.(check int)
    "incarnation incremented" 2
    (incarnation_to_int result.new_state.incarnation);
  match result.broadcasts with
  | [ Alive { incarnation; _ } ] ->
      Alcotest.(check int)
        "refute incarnation" 2
        (incarnation_to_int incarnation)
  | _ -> Alcotest.fail "expected Alive refute broadcast"

let test_suspect_higher_incarnation_suspects () =
  let member = make_member ~state:alive_state ~incarnation:1 node1 in
  let msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 2;
        suspector = node_id_of_string "node2";
      }
  in
  let result =
    handle_suspect ~self:(node_id_of_string "self") member msg ~now
  in
  Alcotest.(check string)
    "state suspect" "suspect"
    (member_state_to_string result.new_state.state)

let test_suspect_lower_incarnation_ignored () =
  let member = make_member ~state:alive_state ~incarnation:5 node1 in
  let msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 3;
        suspector = node_id_of_string "node2";
      }
  in
  let result =
    handle_suspect ~self:(node_id_of_string "self") member msg ~now
  in
  Alcotest.(check string)
    "state unchanged" "alive"
    (member_state_to_string result.new_state.state);
  Alcotest.(check bool) "no broadcast" true (List.length result.broadcasts = 0)

let test_suspect_dead_node_ignored () =
  let member = make_member ~state:dead_state ~incarnation:1 node1 in
  let msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 5;
        suspector = node_id_of_string "node2";
      }
  in
  let result =
    handle_suspect ~self:(node_id_of_string "self") member msg ~now
  in
  Alcotest.(check string)
    "state dead" "dead"
    (member_state_to_string result.new_state.state)

let test_dead_marks_node_dead () =
  let member = make_member ~state:alive_state ~incarnation:1 node1 in
  let msg =
    Dead
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 2;
        declarator = node_id_of_string "node2";
      }
  in
  let result = handle_dead member msg ~now in
  Alcotest.(check string)
    "state dead" "dead"
    (member_state_to_string result.new_state.state);
  match result.events with
  | [ Leave _ ] -> ()
  | _ -> Alcotest.fail "expected Leave event"

let test_dead_already_dead_ignored () =
  let member = make_member ~state:dead_state ~incarnation:5 node1 in
  let msg =
    Dead
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 10;
        declarator = node_id_of_string "node2";
      }
  in
  let result = handle_dead member msg ~now in
  Alcotest.(check bool) "no events" true (List.length result.events = 0)

let test_dead_lower_incarnation_ignored () =
  let member = make_member ~state:alive_state ~incarnation:10 node1 in
  let msg =
    Dead
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 5;
        declarator = node_id_of_string "node2";
      }
  in
  let result = handle_dead member msg ~now in
  Alcotest.(check string)
    "state alive" "alive"
    (member_state_to_string result.new_state.state)

let test_invalidates_dead_beats_all () =
  let dead_msg =
    Dead
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 1;
        declarator = node_id_of_string "node2";
      }
  in
  let alive_msg = Alive { node = node1; incarnation = incarnation_of_int 5 } in
  let suspect_msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 5;
        suspector = node_id_of_string "node2";
      }
  in
  Alcotest.(check bool)
    "dead invalidates alive" true
    (invalidates ~newer:dead_msg ~older:alive_msg);
  Alcotest.(check bool)
    "dead invalidates suspect" true
    (invalidates ~newer:dead_msg ~older:suspect_msg)

let test_invalidates_alive_beats_suspect_same_inc () =
  let alive_msg = Alive { node = node1; incarnation = incarnation_of_int 5 } in
  let suspect_msg =
    Suspect
      {
        node = node_id_of_string "node1";
        incarnation = incarnation_of_int 5;
        suspector = node_id_of_string "node2";
      }
  in
  Alcotest.(check bool)
    "alive beats suspect" true
    (invalidates ~newer:alive_msg ~older:suspect_msg)

let test_invalidates_higher_incarnation_wins () =
  let alive_old = Alive { node = node1; incarnation = incarnation_of_int 1 } in
  let alive_new = Alive { node = node1; incarnation = incarnation_of_int 2 } in
  Alcotest.(check bool)
    "higher inc wins" true
    (invalidates ~newer:alive_new ~older:alive_old);
  Alcotest.(check bool)
    "lower inc doesnt" false
    (invalidates ~newer:alive_old ~older:alive_new)

let test_invalidates_different_nodes_false () =
  let alive1 = Alive { node = node1; incarnation = incarnation_of_int 10 } in
  let alive2 = Alive { node = node2; incarnation = incarnation_of_int 1 } in
  Alcotest.(check bool)
    "different nodes" false
    (invalidates ~newer:alive1 ~older:alive2)

let test_merge_dead_local_wins () =
  let local = make_member ~state:dead_state ~incarnation:1 node1 in
  let remote = make_member ~state:alive_state ~incarnation:10 node1 in
  let result = merge_member_state ~local ~remote in
  Alcotest.(check string)
    "local dead wins" "dead"
    (member_state_to_string result.state)

let test_merge_remote_dead_higher_inc_wins () =
  let local = make_member ~state:alive_state ~incarnation:1 node1 in
  let remote = make_member ~state:dead_state ~incarnation:5 node1 in
  let result = merge_member_state ~local ~remote in
  Alcotest.(check string)
    "remote dead wins" "dead"
    (member_state_to_string result.state)

let test_merge_higher_incarnation_wins () =
  let local = make_member ~state:alive_state ~incarnation:1 node1 in
  let remote = make_member ~state:alive_state ~incarnation:5 node1 in
  let result = merge_member_state ~local ~remote in
  Alcotest.(check int) "higher inc" 5 (incarnation_to_int result.incarnation)

let test_merge_suspect_beats_alive_higher_inc () =
  let local = make_member ~state:alive_state ~incarnation:1 node1 in
  let remote = make_member ~state:suspect_state ~incarnation:5 node1 in
  let result = merge_member_state ~local ~remote in
  Alcotest.(check string)
    "suspect" "suspect"
    (member_state_to_string result.state)

let test_merge_alive_beats_suspect_same_or_higher_inc () =
  let local = make_member ~state:suspect_state ~incarnation:3 node1 in
  let remote = make_member ~state:alive_state ~incarnation:3 node1 in
  let result = merge_member_state ~local ~remote in
  Alcotest.(check string) "alive" "alive" (member_state_to_string result.state)

let test_suspicion_timeout_increases_with_nodes () =
  let config = default_config in
  let t1 = suspicion_timeout config ~node_count:10 in
  let t2 = suspicion_timeout config ~node_count:100 in
  Alcotest.(check bool) "more nodes = longer timeout" true (t2 > t1)

let test_suspicion_timeout_bounded () =
  let config = { default_config with suspicion_max_timeout = 30.0 } in
  let t = suspicion_timeout config ~node_count:1000000 in
  Alcotest.(check bool) "bounded" true (t <= 30.0)

let test_suspicion_timeout_zero_nodes () =
  let config = default_config in
  let t = suspicion_timeout config ~node_count:0 in
  Alcotest.(check bool) "handles zero" true (t >= 0.0)

let test_retransmit_limit_increases_with_nodes () =
  let config = default_config in
  let r1 = retransmit_limit config ~node_count:10 in
  let r2 = retransmit_limit config ~node_count:100 in
  Alcotest.(check bool) "more nodes = higher limit" true (r2 > r1)

let test_next_probe_empty_list () =
  let result =
    next_probe_target ~self:(node_id_of_string "self") ~probe_index:0
      ~members:[]
  in
  match result with None -> () | Some _ -> Alcotest.fail "expected None"

let test_next_probe_skips_self () =
  let self = node_id_of_string "node1" in
  let result =
    next_probe_target ~self ~probe_index:0 ~members:[ node1; node2 ]
  in
  match result with
  | Some (target, _) ->
      Alcotest.(check bool) "not self" false (equal_node_id target.id self)
  | None -> Alcotest.fail "expected Some"

let test_next_probe_wraps_around () =
  let self = node_id_of_string "self" in
  let result =
    next_probe_target ~self ~probe_index:5 ~members:[ node1; node2 ]
  in
  match result with Some _ -> () | None -> Alcotest.fail "expected Some"

let test_next_probe_all_self_returns_none () =
  let self = node_id_of_string "node1" in
  let result = next_probe_target ~self ~probe_index:0 ~members:[ node1 ] in
  match result with
  | None -> ()
  | Some _ -> Alcotest.fail "expected None when only self"

let test_select_indirect_targets_excludes_self_and_target () =
  let self = node_id_of_string "node1" in
  let exclude = node_id_of_string "node2" in
  let result =
    select_indirect_targets ~self ~exclude ~count:10
      ~members:[ node1; node2; node3 ]
  in
  Alcotest.(check int) "only node3" 1 (List.length result);
  Alcotest.(check bool)
    "is node3" true
    (equal_node_id (List.hd result).id node3.id)

let test_select_indirect_targets_limits_count () =
  let self = node_id_of_string "self" in
  let exclude = node_id_of_string "exclude" in
  let result =
    select_indirect_targets ~self ~exclude ~count:1
      ~members:[ node1; node2; node3 ]
  in
  Alcotest.(check int) "limited to 1" 1 (List.length result)

let clamp_incarnation inc =
  incarnation_of_int (incarnation_to_int inc land 0x7FFFFFFF)

let test_merge_converges =
  QCheck.Test.make ~count:200
    ~name:"merge converges (applying twice yields same result)"
    (QCheck.pair Generators.arb_member_snapshot Generators.arb_member_snapshot)
    (fun (a, b) ->
      let a =
        { a with node = node1; incarnation = clamp_incarnation a.incarnation }
      in
      let b =
        { b with node = node1; incarnation = clamp_incarnation b.incarnation }
      in
      let ab = merge_member_state ~local:a ~remote:b in
      let ab2 = merge_member_state ~local:ab ~remote:b in
      ab.state = ab2.state && ab.incarnation = ab2.incarnation)

let test_merge_idempotent =
  QCheck.Test.make ~count:200 ~name:"merge is idempotent"
    Generators.arb_member_snapshot (fun a ->
      let result = merge_member_state ~local:a ~remote:a in
      result.state = a.state && result.incarnation = a.incarnation)

let qcheck_tests =
  List.map QCheck_alcotest.to_alcotest
    [ test_merge_converges; test_merge_idempotent ]

let unit_tests =
  [
    ("alive_higher_incarnation_wins", `Quick, test_alive_higher_incarnation_wins);
    ( "alive_lower_incarnation_ignored",
      `Quick,
      test_alive_lower_incarnation_ignored );
    ( "alive_same_incarnation_unsuspects",
      `Quick,
      test_alive_same_incarnation_unsuspects );
    ("alive_revives_dead_node", `Quick, test_alive_revives_dead_node);
    ( "suspect_triggers_refute_for_self",
      `Quick,
      test_suspect_triggers_refute_for_self );
    ( "suspect_higher_incarnation_suspects",
      `Quick,
      test_suspect_higher_incarnation_suspects );
    ( "suspect_lower_incarnation_ignored",
      `Quick,
      test_suspect_lower_incarnation_ignored );
    ("suspect_dead_node_ignored", `Quick, test_suspect_dead_node_ignored);
    ("dead_marks_node_dead", `Quick, test_dead_marks_node_dead);
    ("dead_already_dead_ignored", `Quick, test_dead_already_dead_ignored);
    ( "dead_lower_incarnation_ignored",
      `Quick,
      test_dead_lower_incarnation_ignored );
    ("invalidates_dead_beats_all", `Quick, test_invalidates_dead_beats_all);
    ( "invalidates_alive_beats_suspect_same_inc",
      `Quick,
      test_invalidates_alive_beats_suspect_same_inc );
    ( "invalidates_higher_incarnation_wins",
      `Quick,
      test_invalidates_higher_incarnation_wins );
    ( "invalidates_different_nodes_false",
      `Quick,
      test_invalidates_different_nodes_false );
    ("merge_dead_local_wins", `Quick, test_merge_dead_local_wins);
    ( "merge_remote_dead_higher_inc_wins",
      `Quick,
      test_merge_remote_dead_higher_inc_wins );
    ("merge_higher_incarnation_wins", `Quick, test_merge_higher_incarnation_wins);
    ( "merge_suspect_beats_alive_higher_inc",
      `Quick,
      test_merge_suspect_beats_alive_higher_inc );
    ( "merge_alive_beats_suspect_same_or_higher_inc",
      `Quick,
      test_merge_alive_beats_suspect_same_or_higher_inc );
    ( "suspicion_timeout_increases_with_nodes",
      `Quick,
      test_suspicion_timeout_increases_with_nodes );
    ("suspicion_timeout_bounded", `Quick, test_suspicion_timeout_bounded);
    ("suspicion_timeout_zero_nodes", `Quick, test_suspicion_timeout_zero_nodes);
    ( "retransmit_limit_increases_with_nodes",
      `Quick,
      test_retransmit_limit_increases_with_nodes );
    ("next_probe_empty_list", `Quick, test_next_probe_empty_list);
    ("next_probe_skips_self", `Quick, test_next_probe_skips_self);
    ("next_probe_wraps_around", `Quick, test_next_probe_wraps_around);
    ( "next_probe_all_self_returns_none",
      `Quick,
      test_next_probe_all_self_returns_none );
    ( "select_indirect_targets_excludes",
      `Quick,
      test_select_indirect_targets_excludes_self_and_target );
    ( "select_indirect_targets_limits",
      `Quick,
      test_select_indirect_targets_limits_count );
  ]

let () =
  Alcotest.run "protocol_pure"
    [ ("property", qcheck_tests); ("unit", unit_tests) ]