open Swim.Types
module Buffer_pool = Swim.Buffer_pool
module Membership = Swim.Membership
module Member = Membership.Member
module Pending_acks = Swim.Pending_acks

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 test_buffer_pool_acquire_release () =
  let pool = Buffer_pool.create ~size:1024 ~count:4 in
  Alcotest.(check int) "initial available" 4 (Buffer_pool.available pool);
  let buf1 = Buffer_pool.acquire pool in
  Alcotest.(check int) "after acquire" 3 (Buffer_pool.available pool);
  Alcotest.(check int) "buffer size" 1024 (Cstruct.length buf1);
  Buffer_pool.release pool buf1;
  Alcotest.(check int) "after release" 4 (Buffer_pool.available pool)

let test_buffer_pool_with_buffer () =
  let pool = Buffer_pool.create ~size:512 ~count:2 in
  let result =
    Buffer_pool.with_buffer pool (fun buf ->
        Alcotest.(check int) "inside with_buffer" 1 (Buffer_pool.available pool);
        Cstruct.length buf)
  in
  Alcotest.(check int) "returned value" 512 result;
  Alcotest.(check int) "after with_buffer" 2 (Buffer_pool.available pool)

let test_buffer_pool_exception_safe () =
  let pool = Buffer_pool.create ~size:256 ~count:2 in
  (try
     Buffer_pool.with_buffer pool (fun _buf ->
         Alcotest.(check int) "during exception" 1 (Buffer_pool.available pool);
         failwith "test exception")
   with Failure _ -> ());
  Alcotest.(check int) "after exception" 2 (Buffer_pool.available pool)

let test_buffer_pool_all_buffers () =
  let pool = Buffer_pool.create ~size:64 ~count:3 in
  let b1 = Buffer_pool.acquire pool in
  let b2 = Buffer_pool.acquire pool in
  let b3 = Buffer_pool.acquire pool in
  Alcotest.(check int) "all acquired" 0 (Buffer_pool.available pool);
  Buffer_pool.release pool b1;
  Buffer_pool.release pool b2;
  Buffer_pool.release pool b3;
  Alcotest.(check int) "all released" 3 (Buffer_pool.available pool)

let test_membership_add_find () =
  let table = Membership.create () in
  let member = Member.create ~now node1 in
  Membership.add table member;
  Alcotest.(check int) "count" 1 (Membership.count table);
  match Membership.find table node1.id with
  | Some m ->
      Alcotest.(check bool)
        "same node" true
        (equal_node_id (Member.node m).id node1.id)
  | None -> Alcotest.fail "member not found"

let test_membership_remove () =
  let table = Membership.create () in
  let member = Member.create ~now node1 in
  Membership.add table member;
  Alcotest.(check bool)
    "remove existing" true
    (Membership.remove table node1.id);
  Alcotest.(check int) "count after remove" 0 (Membership.count table);
  Alcotest.(check bool)
    "remove non-existing" false
    (Membership.remove table node1.id)

let test_membership_to_list () =
  let table = Membership.create () in
  Membership.add table (Member.create ~now node1);
  Membership.add table (Member.create ~now node2);
  Membership.add table (Member.create ~now node3);
  let members = Membership.to_list table in
  Alcotest.(check int) "list length" 3 (List.length members)

let test_membership_snapshot_consistency () =
  let table = Membership.create () in
  Membership.add table (Member.create ~now node1);
  let snapshots = Membership.snapshot_all table in
  Alcotest.(check int) "snapshot count" 1 (List.length snapshots);
  let snap = List.hd snapshots in
  Alcotest.(check bool)
    "node id matches" true
    (equal_node_id snap.node.id node1.id)

let test_membership_count_accurate () =
  let table = Membership.create () in
  Alcotest.(check int) "empty" 0 (Membership.count table);
  Membership.add table (Member.create ~now node1);
  Alcotest.(check int) "after add 1" 1 (Membership.count table);
  Membership.add table (Member.create ~now node2);
  Alcotest.(check int) "after add 2" 2 (Membership.count table);
  let _ = Membership.remove table node1.id in
  Alcotest.(check int) "after remove" 1 (Membership.count table)

let test_membership_no_duplicates () =
  let table = Membership.create () in
  Membership.add table (Member.create ~now node1);
  Membership.add table (Member.create ~now node1);
  Alcotest.(check int) "no duplicate" 1 (Membership.count table)

let test_pending_acks_register_cancel () =
  let pa = Pending_acks.create () in
  let _ = Pending_acks.register pa ~seq:123 in
  Alcotest.(check int) "pending before" 1 (Pending_acks.pending_count pa);
  Pending_acks.cancel pa ~seq:123;
  Alcotest.(check int) "pending after" 0 (Pending_acks.pending_count pa)

let test_pending_acks_complete_not_found () =
  let pa = Pending_acks.create () in
  let completed = Pending_acks.complete pa ~seq:999 ~payload:None in
  Alcotest.(check bool) "not found" false completed

let test_member_state_transitions () =
  let table = Membership.create () in
  let member = Member.create ~now node1 in
  Membership.add table member;
  let updated =
    Membership.update_member table node1.id
      {
        update =
          (fun m ~xt ->
            Member.set_suspect m ~incarnation:(incarnation_of_int 5) ~now ~xt);
      }
  in
  Alcotest.(check bool) "updated" true updated;
  match Membership.find table node1.id with
  | Some m ->
      let snap = Member.snapshot_now m in
      Alcotest.(check string)
        "state suspect" "suspect"
        (member_state_to_string snap.state);
      Alcotest.(check int) "incarnation" 5 (incarnation_to_int snap.incarnation)
  | None -> Alcotest.fail "member not found"

let test_member_set_alive () =
  let table = Membership.create () in
  let member = Member.create ~now node1 in
  Membership.add table member;
  let _ =
    Membership.update_member table node1.id
      {
        update =
          (fun m ~xt ->
            Member.set_suspect m ~incarnation:(incarnation_of_int 1) ~now ~xt);
      }
  in
  let _ =
    Membership.update_member table node1.id
      {
        update =
          (fun m ~xt ->
            Member.set_alive m ~incarnation:(incarnation_of_int 2) ~now ~xt);
      }
  in
  match Membership.find table node1.id with
  | Some m ->
      let snap = Member.snapshot_now m in
      Alcotest.(check string)
        "state alive" "alive"
        (member_state_to_string snap.state)
  | None -> Alcotest.fail "member not found"

let test_member_set_dead () =
  let table = Membership.create () in
  let member = Member.create ~now node1 in
  Membership.add table member;
  let _ =
    Membership.update_member table node1.id
      {
        update =
          (fun m ~xt ->
            Member.set_dead m ~incarnation:(incarnation_of_int 10) ~now ~xt);
      }
  in
  match Membership.find table node1.id with
  | Some m ->
      let snap = Member.snapshot_now m in
      Alcotest.(check string)
        "state dead" "dead"
        (member_state_to_string snap.state)
  | None -> Alcotest.fail "member not found"

let () =
  Eio_main.run @@ fun _env ->
  Alcotest.run "kcas"
    [
      ( "buffer_pool",
        [
          ("acquire_release", `Quick, test_buffer_pool_acquire_release);
          ("with_buffer", `Quick, test_buffer_pool_with_buffer);
          ("exception_safe", `Quick, test_buffer_pool_exception_safe);
          ("all_buffers", `Quick, test_buffer_pool_all_buffers);
        ] );
      ( "membership",
        [
          ("add_find", `Quick, test_membership_add_find);
          ("remove", `Quick, test_membership_remove);
          ("to_list", `Quick, test_membership_to_list);
          ("snapshot_consistency", `Quick, test_membership_snapshot_consistency);
          ("count_accurate", `Quick, test_membership_count_accurate);
          ("no_duplicates", `Quick, test_membership_no_duplicates);
        ] );
      ( "pending_acks",
        [
          ("register_cancel", `Quick, test_pending_acks_register_cancel);
          ("complete_not_found", `Quick, test_pending_acks_complete_not_found);
        ] );
      ( "member_transitions",
        [
          ("state_transitions", `Quick, test_member_state_transitions);
          ("set_alive", `Quick, test_member_set_alive);
          ("set_dead", `Quick, test_member_set_dead);
        ] );
    ]