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; sender } -> Ping { seq = clamp_int32 seq; 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_msg =
  QCheck.Test.make ~count:1000 ~name:"codec message roundtrip"
    Generators.arb_protocol_msg (fun msg ->
      let msg = normalize_msg msg in
      let size = encoded_size msg + 100 in
      let buf = Cstruct.create size in
      let enc = Encoder.create ~buf in
      encode_msg enc msg;
      let encoded = Encoder.to_cstruct enc in
      let dec = Decoder.create encoded in
      match decode_msg dec with Ok decoded -> decoded = msg | Error _ -> false)

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 =
        4 + 1 + 2
        + String.length packet.cluster
        + 2
        + encoded_size packet.primary
        + List.fold_left (fun acc m -> acc + encoded_size m) 0 packet.piggyback
        + 100
      in
      let buf = Cstruct.create size in
      match encode_packet packet ~buf with
      | Error _ -> false
      | Ok len -> (
          let encoded = Cstruct.sub buf 0 len in
          match decode_packet encoded with
          | Ok decoded -> decoded = packet
          | Error _ -> false))

let test_encoded_size_accurate =
  QCheck.Test.make ~count:1000 ~name:"encoded_size matches actual encoding"
    Generators.arb_protocol_msg (fun msg ->
      let predicted = encoded_size msg in
      let buf = Cstruct.create (predicted + 100) in
      let enc = Encoder.create ~buf in
      encode_msg enc msg;
      let actual = Encoder.pos enc in
      predicted = actual)

let make_valid_packet_buf () =
  let node =
    make_node_info ~id:(node_id_of_string "n1")
      ~addr:(`Udp (Eio.Net.Ipaddr.of_raw "\127\000\000\001", 7946))
      ~meta:""
  in
  let packet =
    {
      cluster = "test";
      primary = Ping { seq = 1; sender = node };
      piggyback = [];
    }
  in
  let buf = Cstruct.create 1000 in
  match encode_packet packet ~buf with
  | Ok len -> Cstruct.sub buf 0 len
  | Error _ -> failwith "encode failed"

let test_invalid_magic_rejected () =
  let buf = make_valid_packet_buf () in
  Cstruct.blit_from_string "FAIL" 0 buf 0 4;
  match decode_packet buf with
  | Error Invalid_magic -> ()
  | _ -> Alcotest.fail "expected Invalid_magic error"

let test_unsupported_version_rejected () =
  let buf = make_valid_packet_buf () in
  Cstruct.set_uint8 buf 4 0x99;
  match decode_packet buf with
  | Error (Unsupported_version 0x99) -> ()
  | Error (Unsupported_version v) ->
      Alcotest.failf "expected version 0x99 but got %d" v
  | _ -> Alcotest.fail "expected Unsupported_version error"

let test_invalid_tag_rejected () =
  let buf = Cstruct.create 100 in
  let enc = Encoder.create ~buf in
  Encoder.write_bytes enc (Cstruct.of_string "SWIM");
  Encoder.write_byte enc 1;
  Encoder.write_string enc "default";
  Encoder.write_int16_be enc 1;
  Encoder.write_byte enc 0xFF;
  let encoded = Encoder.to_cstruct enc in
  match decode_packet encoded with
  | Error (Invalid_tag 0xFF) -> ()
  | Error (Invalid_tag t) -> Alcotest.failf "expected tag 0xFF but got %d" t
  | _ -> Alcotest.fail "expected Invalid_tag error"

let test_encoder_write_byte () =
  let buf = Cstruct.create 10 in
  let enc = Encoder.create ~buf in
  Encoder.write_byte enc 0x42;
  Encoder.write_byte enc 0xFF;
  let result = Encoder.to_cstruct enc in
  Alcotest.(check int) "length" 2 (Cstruct.length result);
  Alcotest.(check int) "byte 0" 0x42 (Cstruct.get_uint8 result 0);
  Alcotest.(check int) "byte 1" 0xFF (Cstruct.get_uint8 result 1)

let test_encoder_write_int16_be () =
  let buf = Cstruct.create 10 in
  let enc = Encoder.create ~buf in
  Encoder.write_int16_be enc 0x1234;
  let result = Encoder.to_cstruct enc in
  Alcotest.(check int) "length" 2 (Cstruct.length result);
  Alcotest.(check int) "value" 0x1234 (Cstruct.BE.get_uint16 result 0)

let test_encoder_write_int32_be () =
  let buf = Cstruct.create 10 in
  let enc = Encoder.create ~buf in
  Encoder.write_int32_be enc 0x12345678l;
  let result = Encoder.to_cstruct enc in
  Alcotest.(check int) "length" 4 (Cstruct.length result);
  Alcotest.(check int32) "value" 0x12345678l (Cstruct.BE.get_uint32 result 0)

let test_encoder_write_string () =
  let buf = Cstruct.create 100 in
  let enc = Encoder.create ~buf in
  Encoder.write_string enc "hello";
  let result = Encoder.to_cstruct enc in
  Alcotest.(check int) "length" 7 (Cstruct.length result);
  Alcotest.(check int) "str_len" 5 (Cstruct.BE.get_uint16 result 0);
  Alcotest.(check string)
    "content" "hello"
    (Cstruct.to_string ~off:2 ~len:5 result)

let test_encoder_write_empty_string () =
  let buf = Cstruct.create 10 in
  let enc = Encoder.create ~buf in
  Encoder.write_string enc "";
  let result = Encoder.to_cstruct enc in
  Alcotest.(check int) "length" 2 (Cstruct.length result);
  Alcotest.(check int) "str_len" 0 (Cstruct.BE.get_uint16 result 0)

let test_decoder_read_byte () =
  let buf = Cstruct.of_string "\x42\xFF" in
  let dec = Decoder.create buf in
  Alcotest.(check int) "byte 0" 0x42 (Decoder.read_byte dec);
  Alcotest.(check int) "byte 1" 0xFF (Decoder.read_byte dec)

let test_decoder_read_int16_be () =
  let buf = Cstruct.create 2 in
  Cstruct.BE.set_uint16 buf 0 0x1234;
  let dec = Decoder.create buf in
  Alcotest.(check int) "value" 0x1234 (Decoder.read_int16_be dec)

let test_decoder_read_int32_be () =
  let buf = Cstruct.create 4 in
  Cstruct.BE.set_uint32 buf 0 0x12345678l;
  let dec = Decoder.create buf in
  Alcotest.(check int32) "value" 0x12345678l (Decoder.read_int32_be dec)

let test_decoder_read_string () =
  let buf = Cstruct.create 10 in
  Cstruct.BE.set_uint16 buf 0 5;
  Cstruct.blit_from_string "hello" 0 buf 2 5;
  let dec = Decoder.create buf in
  Alcotest.(check string) "value" "hello" (Decoder.read_string dec)

let test_decoder_remaining () =
  let buf = Cstruct.create 10 in
  let dec = Decoder.create buf in
  Alcotest.(check int) "initial" 10 (Decoder.remaining dec);
  let _ = Decoder.read_byte dec in
  Alcotest.(check int) "after byte" 9 (Decoder.remaining dec);
  let _ = Decoder.read_int32_be dec in
  Alcotest.(check int) "after int32" 5 (Decoder.remaining dec)

let test_decoder_is_empty () =
  let buf = Cstruct.create 1 in
  let dec = Decoder.create buf in
  Alcotest.(check bool) "not empty" false (Decoder.is_empty dec);
  let _ = Decoder.read_byte dec in
  Alcotest.(check bool) "empty" true (Decoder.is_empty dec)

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; 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; 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 qcheck_tests =
  List.map QCheck_alcotest.to_alcotest
    [ test_roundtrip_msg; test_roundtrip_packet; test_encoded_size_accurate ]

let unit_tests =
  [
    ("invalid_magic_rejected", `Quick, test_invalid_magic_rejected);
    ("unsupported_version_rejected", `Quick, test_unsupported_version_rejected);
    ("invalid_tag_rejected", `Quick, test_invalid_tag_rejected);
    ("encoder_write_byte", `Quick, test_encoder_write_byte);
    ("encoder_write_int16_be", `Quick, test_encoder_write_int16_be);
    ("encoder_write_int32_be", `Quick, test_encoder_write_int32_be);
    ("encoder_write_string", `Quick, test_encoder_write_string);
    ("encoder_write_empty_string", `Quick, test_encoder_write_empty_string);
    ("decoder_read_byte", `Quick, test_decoder_read_byte);
    ("decoder_read_int16_be", `Quick, test_decoder_read_int16_be);
    ("decoder_read_int32_be", `Quick, test_decoder_read_int32_be);
    ("decoder_read_string", `Quick, test_decoder_read_string);
    ("decoder_remaining", `Quick, test_decoder_remaining);
    ("decoder_is_empty", `Quick, test_decoder_is_empty);
    ("empty_piggyback", `Quick, test_empty_piggyback);
    ("multiple_piggyback", `Quick, test_multiple_piggyback);
  ]

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