open Swim.Types
open Swim.Codec

let clamp_int32 n = n land 0x7FFFFFFF

let clamp_incarnation inc =
  incarnation_of_int (clamp_int32 (incarnation_to_int inc))

let normalize_msg msg =
  match msg with
  | Ping { seq; target; sender } ->
      Ping { seq = clamp_int32 seq; target; sender }
  | Ping_req { seq; target; sender } ->
      Ping_req { seq = clamp_int32 seq; target; sender }
  | Ack { seq; responder; payload } ->
      Ack { seq = clamp_int32 seq; responder; payload }
  | Alive { node; incarnation } ->
      Alive { node; incarnation = clamp_incarnation incarnation }
  | Suspect { node; incarnation; suspector } ->
      Suspect { node; incarnation = clamp_incarnation incarnation; suspector }
  | Dead { node; incarnation; declarator } ->
      Dead { node; incarnation = clamp_incarnation incarnation; declarator }
  | User_msg _ as msg -> msg

let normalize_packet packet =
  let primary = normalize_msg packet.primary in
  let piggyback = List.map normalize_msg packet.piggyback in
  { packet with primary; piggyback }

let test_roundtrip_packet =
  QCheck.Test.make ~count:500 ~name:"codec packet roundtrip"
    Generators.arb_packet (fun packet ->
      let packet = normalize_packet packet in
      let size = 8192 in
      let buf = Cstruct.create size in
      match encode_packet packet ~buf with
      | Error _ -> true
      | Ok len -> (
          let encoded = Cstruct.sub buf 0 len in
          match decode_packet encoded with
          | Ok decoded ->
              List.length decoded.piggyback = List.length packet.piggyback
          | Error _ -> true))

let test_empty_piggyback () =
  let node =
    make_node_info
      ~id:(node_id_of_string "node1")
      ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", 7946))
      ~meta:""
  in
  let packet =
    {
      cluster = "test";
      primary =
        Ping { seq = 1; target = node_id_of_string "target"; sender = node };
      piggyback = [];
    }
  in
  let buf = Cstruct.create 1000 in
  match encode_packet packet ~buf with
  | Error _ -> Alcotest.fail "encode failed"
  | Ok len -> (
      let encoded = Cstruct.sub buf 0 len in
      match decode_packet encoded with
      | Ok decoded ->
          Alcotest.(check int)
            "piggyback count" 0
            (List.length decoded.piggyback)
      | Error e -> Alcotest.failf "decode failed: %s" (decode_error_to_string e)
      )

let test_multiple_piggyback () =
  let node =
    make_node_info
      ~id:(node_id_of_string "node1")
      ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", 7946))
      ~meta:""
  in
  let alive_state : member_state = Alive in
  let _ = alive_state in
  let piggyback =
    [
      Alive { node; incarnation = incarnation_of_int 1 };
      Suspect
        {
          node = node_id_of_string "node2";
          incarnation = incarnation_of_int 2;
          suspector = node_id_of_string "node1";
        };
      Dead
        {
          node = node_id_of_string "node3";
          incarnation = incarnation_of_int 3;
          declarator = node_id_of_string "node1";
        };
    ]
  in
  let packet =
    {
      cluster = "test";
      primary =
        Ping { seq = 1; target = node_id_of_string "target"; sender = node };
      piggyback;
    }
  in
  let buf = Cstruct.create 2000 in
  match encode_packet packet ~buf with
  | Error _ -> Alcotest.fail "encode failed"
  | Ok len -> (
      let encoded = Cstruct.sub buf 0 len in
      match decode_packet encoded with
      | Ok decoded ->
          Alcotest.(check int)
            "piggyback count" 3
            (List.length decoded.piggyback)
      | Error e -> Alcotest.failf "decode failed: %s" (decode_error_to_string e)
      )

let test_crc_roundtrip () =
  let data = "hello world" in
  let with_crc = add_crc data in
  match verify_and_strip_crc_string with_crc with
  | Ok stripped -> Alcotest.(check string) "stripped" data stripped
  | Error _ -> Alcotest.fail "CRC verification failed"

let test_crc_corruption_detected () =
  let data = "hello world" in
  let with_crc = add_crc data in
  let corrupted = Bytes.of_string with_crc in
  Bytes.set corrupted 6 '\xFF';
  match verify_and_strip_crc_string (Bytes.to_string corrupted) with
  | Error Invalid_crc -> ()
  | _ -> Alcotest.fail "expected CRC error"

let test_label_roundtrip () =
  let label = "my-label" in
  let data = "payload data" in
  let with_label = add_label label data in
  match strip_label_string with_label with
  | Ok (stripped, extracted_label) ->
      Alcotest.(check string) "payload" data stripped;
      Alcotest.(check string) "label" label extracted_label
  | Error _ -> Alcotest.fail "label extraction failed"

let test_empty_label () =
  let data = "payload data" in
  let with_label = add_label "" data in
  Alcotest.(check string) "no change" data with_label

let test_compound_msg_roundtrip () =
  let msgs = [ "msg1"; "msg2"; "msg3" ] in
  let compound = make_compound_msg msgs in
  let payload = String.sub compound 1 (String.length compound - 1) in
  match decode_compound_msg payload with
  | Ok (decoded, trunc) ->
      Alcotest.(check int) "no truncation" 0 trunc;
      Alcotest.(check int) "msg count" 3 (List.length decoded);
      Alcotest.(check string) "msg1" "msg1" (List.nth decoded 0);
      Alcotest.(check string) "msg2" "msg2" (List.nth decoded 1);
      Alcotest.(check string) "msg3" "msg3" (List.nth decoded 2)
  | Error _ -> Alcotest.fail "compound decode failed"

let qcheck_tests =
  List.map QCheck_alcotest.to_alcotest [ test_roundtrip_packet ]

let unit_tests =
  [
    ("empty_piggyback", `Quick, test_empty_piggyback);
    ("multiple_piggyback", `Quick, test_multiple_piggyback);
    ("crc_roundtrip", `Quick, test_crc_roundtrip);
    ("crc_corruption_detected", `Quick, test_crc_corruption_detected);
    ("label_roundtrip", `Quick, test_label_roundtrip);
    ("empty_label", `Quick, test_empty_label);
    ("compound_msg_roundtrip", `Quick, test_compound_msg_roundtrip);
  ]

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