(*---------------------------------------------------------------------------
   Copyright (c) 2025 Thomas Gazagnaire. All rights reserved.
   SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

(* Test vectors based on RFC 9171 examples *)

let eid_testable = Alcotest.testable Bundle.pp_eid ( = )
let error_testable = Alcotest.testable Bundle.pp_error ( = )

let bundle_testable =
  Alcotest.testable Bundle.pp (fun a b -> Bundle.encode a = Bundle.encode b)

let test_eid_dtn_none () =
  let eid = Bundle.Dtn_none in
  let cbor = Bundle.eid_to_cbor eid in
  let decoded = Bundle.eid_of_cbor cbor in
  Alcotest.(check (result eid_testable string))
    "dtn:none roundtrip" (Ok eid) decoded

let test_eid_dtn () =
  let eid = Bundle.Dtn "//node-1/incoming" in
  let cbor = Bundle.eid_to_cbor eid in
  let decoded = Bundle.eid_of_cbor cbor in
  Alcotest.(check (result eid_testable string))
    "dtn URI roundtrip" (Ok eid) decoded

let test_eid_ipn () =
  let eid = Bundle.Ipn (1L, 1L) in
  let cbor = Bundle.eid_to_cbor eid in
  let decoded = Bundle.eid_of_cbor cbor in
  Alcotest.(check (result eid_testable string))
    "ipn:1.1 roundtrip" (Ok eid) decoded

let test_bundle_flags () =
  let flags =
    Bundle.{ default_flags with ack_requested = true; report_delivery = true }
  in
  let encoded = Bundle.int_of_flags flags in
  let decoded = Bundle.flags_of_int encoded in
  Alcotest.(check bool)
    "ack_requested" flags.ack_requested decoded.ack_requested;
  Alcotest.(check bool)
    "report_delivery" flags.report_delivery decoded.report_delivery;
  Alcotest.(check bool) "is_fragment" flags.is_fragment decoded.is_fragment

let test_block_flags () =
  let flags =
    Bundle.{ block_flags_default with replicate_in_fragment = true }
  in
  let encoded = Bundle.int_of_block_flags flags in
  let decoded = Bundle.block_flags_of_int encoded in
  Alcotest.(check bool)
    "replicate_in_fragment" flags.replicate_in_fragment
    decoded.replicate_in_fragment

let test_crc_type () =
  Alcotest.(check int) "No_crc" 0 (Bundle.int_of_crc_type Bundle.No_crc);
  Alcotest.(check int) "Crc16" 1 (Bundle.int_of_crc_type Bundle.Crc16);
  Alcotest.(check int) "Crc32c" 2 (Bundle.int_of_crc_type Bundle.Crc32c);
  Alcotest.(check (result (of_pp Bundle.pp_crc_type) int))
    "decode 0" (Ok Bundle.No_crc) (Bundle.crc_type_of_int 0);
  Alcotest.(check (result (of_pp Bundle.pp_crc_type) int))
    "decode 2" (Ok Bundle.Crc32c) (Bundle.crc_type_of_int 2);
  Alcotest.(check (result (of_pp Bundle.pp_crc_type) int))
    "decode invalid" (Error 99)
    (Bundle.crc_type_of_int 99)

let test_simple_bundle () =
  let bundle =
    Bundle.v
      ~source:(Bundle.Ipn (1L, 1L))
      ~destination:(Bundle.Ipn (2L, 1L))
      ~creation_timestamp:{ Bundle.time = 0L; seq = 0L }
      ~payload:"Hello, DTN!" ()
  in
  Alcotest.(check (option string))
    "payload" (Some "Hello, DTN!") (Bundle.payload bundle);
  Alcotest.(check int) "version" 7 bundle.primary.version;
  Alcotest.(check (of_pp Bundle.pp_eid))
    "source"
    (Bundle.Ipn (1L, 1L))
    bundle.primary.source;
  Alcotest.(check (of_pp Bundle.pp_eid))
    "destination"
    (Bundle.Ipn (2L, 1L))
    bundle.primary.destination

let test_bundle_roundtrip () =
  let bundle =
    Bundle.v ~source:(Bundle.Dtn "//sender/app")
      ~destination:(Bundle.Dtn "//receiver/app")
      ~creation_timestamp:{ Bundle.time = 1000000L; seq = 42L }
      ~lifetime:3600000L ~payload:"Test payload data" ()
  in
  let encoded = Bundle.encode bundle in
  let decoded = Bundle.decode encoded in
  match decoded with
  | Ok decoded_bundle ->
      Alcotest.(check (option string))
        "payload preserved" (Bundle.payload bundle)
        (Bundle.payload decoded_bundle);
      Alcotest.(check (of_pp Bundle.pp_eid))
        "source preserved" bundle.primary.source decoded_bundle.primary.source;
      Alcotest.(check (of_pp Bundle.pp_eid))
        "destination preserved" bundle.primary.destination
        decoded_bundle.primary.destination;
      Alcotest.(check int64)
        "lifetime preserved" bundle.primary.lifetime
        decoded_bundle.primary.lifetime
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

let test_bundle_no_crc () =
  let bundle =
    Bundle.v ~crc_type:Bundle.No_crc
      ~source:(Bundle.Ipn (1L, 1L))
      ~destination:(Bundle.Ipn (2L, 1L))
      ~creation_timestamp:{ Bundle.time = 0L; seq = 0L }
      ~payload:"No CRC test" ()
  in
  let encoded = Bundle.encode bundle in
  let decoded = Bundle.decode encoded in
  Alcotest.(check (result bundle_testable error_testable))
    "no CRC roundtrip" (Ok bundle) decoded

let test_bundle_crc16 () =
  let bundle =
    Bundle.v ~crc_type:Bundle.Crc16
      ~source:(Bundle.Ipn (1L, 1L))
      ~destination:(Bundle.Ipn (2L, 1L))
      ~creation_timestamp:{ Bundle.time = 0L; seq = 0L }
      ~payload:"CRC-16 test" ()
  in
  let encoded = Bundle.encode bundle in
  let decoded = Bundle.decode encoded in
  match decoded with
  | Ok _ -> ()
  | Error e ->
      Alcotest.fail
        (Fmt.str "CRC-16 bundle decode failed: %a" Bundle.pp_error e)

(* {1 Fragment Tests} *)

let test_fragment_fields () =
  (* RFC 9171 Section 4.3.1: fragment_offset and total_adu_length *)
  let flags = Bundle.{ default_flags with is_fragment = true } in
  let primary =
    {
      Bundle.version = 7;
      flags;
      crc_type = Bundle.Crc32c;
      destination = Bundle.Ipn (2L, 1L);
      source = Bundle.Ipn (1L, 1L);
      report_to = Bundle.Ipn (1L, 1L);
      creation_timestamp = { Bundle.time = 1000L; seq = 0L };
      lifetime = 3600000L;
      fragment_offset = Some 1024L;
      total_adu_length = Some 4096L;
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.Crc32c;
      data = Bundle.Payload_data "fragment data";
    }
  in
  let bundle = { Bundle.primary; blocks = [ payload_block ] } in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check bool) "is_fragment" true decoded.primary.flags.is_fragment;
      Alcotest.(check (option int64))
        "fragment_offset" (Some 1024L) decoded.primary.fragment_offset;
      Alcotest.(check (option int64))
        "total_adu_length" (Some 4096L) decoded.primary.total_adu_length
  | Error e -> Alcotest.failf "fragment decode failed: %a" Bundle.pp_error e

let test_fragment_offset_zero () =
  (* First fragment has offset 0 *)
  let flags = Bundle.{ default_flags with is_fragment = true } in
  let primary =
    {
      Bundle.version = 7;
      flags;
      crc_type = Bundle.No_crc;
      destination = Bundle.Ipn (2L, 1L);
      source = Bundle.Ipn (1L, 1L);
      report_to = Bundle.Ipn (1L, 1L);
      creation_timestamp = { Bundle.time = 0L; seq = 0L };
      lifetime = 3600000L;
      fragment_offset = Some 0L;
      total_adu_length = Some 2048L;
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Payload_data "first fragment";
    }
  in
  let bundle = { Bundle.primary; blocks = [ payload_block ] } in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option int64))
        "fragment_offset" (Some 0L) decoded.primary.fragment_offset
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

(* {1 Additional Block Type Tests} *)

let test_hop_count_block () =
  (* RFC 9171 Section 4.4.4: Hop Count block type 10 *)
  let primary =
    {
      Bundle.version = 7;
      flags = Bundle.default_flags;
      crc_type = Bundle.No_crc;
      destination = Bundle.Ipn (2L, 1L);
      source = Bundle.Ipn (1L, 1L);
      report_to = Bundle.Ipn (1L, 1L);
      creation_timestamp = { Bundle.time = 0L; seq = 0L };
      lifetime = 3600000L;
      fragment_offset = None;
      total_adu_length = None;
    }
  in
  let hop_count_block =
    {
      Bundle.block_type = Bundle.Hop_count;
      block_number = 2;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Hop_count_data { limit = 30; count = 5 };
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Payload_data "test";
    }
  in
  let bundle =
    { Bundle.primary; blocks = [ hop_count_block; payload_block ] }
  in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option (pair int int)))
        "hop count"
        (Some (30, 5))
        (Bundle.hop_count decoded)
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

let test_previous_node_block () =
  (* RFC 9171 Section 4.4.2: Previous Node block type 6 *)
  let prev_eid = Bundle.Ipn (3L, 0L) in
  let primary =
    {
      Bundle.version = 7;
      flags = Bundle.default_flags;
      crc_type = Bundle.No_crc;
      destination = Bundle.Ipn (2L, 1L);
      source = Bundle.Ipn (1L, 1L);
      report_to = Bundle.Ipn (1L, 1L);
      creation_timestamp = { Bundle.time = 0L; seq = 0L };
      lifetime = 3600000L;
      fragment_offset = None;
      total_adu_length = None;
    }
  in
  let prev_node_block =
    {
      Bundle.block_type = Bundle.Previous_node;
      block_number = 2;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Previous_node_data prev_eid;
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Payload_data "test";
    }
  in
  let bundle =
    { Bundle.primary; blocks = [ prev_node_block; payload_block ] }
  in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option (of_pp Bundle.pp_eid)))
        "previous node" (Some prev_eid)
        (Bundle.previous_node decoded)
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

let test_bundle_age_block () =
  (* RFC 9171 Section 4.4.3: Bundle Age block type 7 *)
  let primary =
    {
      Bundle.version = 7;
      flags = Bundle.default_flags;
      crc_type = Bundle.No_crc;
      destination = Bundle.Ipn (2L, 1L);
      source = Bundle.Ipn (1L, 1L);
      report_to = Bundle.Ipn (1L, 1L);
      creation_timestamp = { Bundle.time = 0L; seq = 0L };
      lifetime = 3600000L;
      fragment_offset = None;
      total_adu_length = None;
    }
  in
  let age_block =
    {
      Bundle.block_type = Bundle.Bundle_age;
      block_number = 2;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Bundle_age_data 12345L;
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Payload_data "test";
    }
  in
  let bundle = { Bundle.primary; blocks = [ age_block; payload_block ] } in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option int64))
        "bundle age" (Some 12345L) (Bundle.age decoded)
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

(* {1 do_not_fragment Flag Tests} *)

let test_do_not_fragment_flag () =
  (* RFC 9171 Section 4.2.3: must_not_fragment flag *)
  let flags = Bundle.{ default_flags with must_not_fragment = true } in
  let encoded_flags = Bundle.int_of_flags flags in
  Alcotest.(check bool) "bit 2 set" true (encoded_flags land 0x04 <> 0);
  let decoded_flags = Bundle.flags_of_int encoded_flags in
  Alcotest.(check bool) "must_not_fragment" true decoded_flags.must_not_fragment;
  (* Roundtrip through a full bundle *)
  let bundle =
    Bundle.v ~flags
      ~source:(Bundle.Ipn (1L, 1L))
      ~destination:(Bundle.Ipn (2L, 1L))
      ~creation_timestamp:{ Bundle.time = 0L; seq = 0L }
      ~payload:"no fragment" ()
  in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check bool)
        "must_not_fragment preserved" true
        decoded.primary.flags.must_not_fragment
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

(* {1 CRC-32C Tests} *)

let test_crc32c_roundtrip () =
  (* Verify CRC-32C encoding and decoding *)
  let bundle =
    Bundle.v ~crc_type:Bundle.Crc32c
      ~source:(Bundle.Ipn (1L, 1L))
      ~destination:(Bundle.Ipn (2L, 1L))
      ~creation_timestamp:{ Bundle.time = 500000L; seq = 1L }
      ~payload:"CRC-32C payload" ()
  in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option string))
        "payload preserved" (Some "CRC-32C payload") (Bundle.payload decoded);
      Alcotest.(check (of_pp Bundle.pp_crc_type))
        "crc_type preserved" Bundle.Crc32c decoded.primary.crc_type
  | Error e -> Alcotest.failf "CRC-32C decode failed: %a" Bundle.pp_error e

let test_crc32c_compute () =
  (* Verify compute_crc produces correct length for CRC-32C *)
  let data = "test data for CRC-32C" in
  let crc = Bundle.compute_crc Bundle.Crc32c data in
  Alcotest.(check int) "CRC-32C is 4 bytes" 4 (String.length crc);
  (* No_crc should produce empty string *)
  let no_crc = Bundle.compute_crc Bundle.No_crc data in
  Alcotest.(check int) "No_crc is 0 bytes" 0 (String.length no_crc);
  (* CRC-16 should produce 2 bytes *)
  let crc16 = Bundle.compute_crc Bundle.Crc16 data in
  Alcotest.(check int) "CRC-16 is 2 bytes" 2 (String.length crc16)

(* {1 Multiple Extension Blocks Test} *)

let test_multiple_extension_blocks () =
  (* Bundle with hop count, previous node, bundle age, and payload *)
  let primary =
    {
      Bundle.version = 7;
      flags = Bundle.default_flags;
      crc_type = Bundle.Crc32c;
      destination = Bundle.Dtn "//receiver/app";
      source = Bundle.Dtn "//sender/app";
      report_to = Bundle.Dtn "//sender/app";
      creation_timestamp = { Bundle.time = 1000000L; seq = 0L };
      lifetime = 7200000L;
      fragment_offset = None;
      total_adu_length = None;
    }
  in
  let hop_count_block =
    {
      Bundle.block_type = Bundle.Hop_count;
      block_number = 2;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Hop_count_data { limit = 20; count = 0 };
    }
  in
  let prev_node_block =
    {
      Bundle.block_type = Bundle.Previous_node;
      block_number = 3;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Previous_node_data (Bundle.Ipn (5L, 0L));
    }
  in
  let age_block =
    {
      Bundle.block_type = Bundle.Bundle_age;
      block_number = 4;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.No_crc;
      data = Bundle.Bundle_age_data 500L;
    }
  in
  let payload_block =
    {
      Bundle.block_type = Bundle.Payload;
      block_number = 1;
      flags = Bundle.block_flags_default;
      crc_type = Bundle.Crc32c;
      data = Bundle.Payload_data "multi-block bundle";
    }
  in
  let bundle =
    {
      Bundle.primary;
      blocks = [ hop_count_block; prev_node_block; age_block; payload_block ];
    }
  in
  let encoded = Bundle.encode bundle in
  match Bundle.decode encoded with
  | Ok decoded ->
      Alcotest.(check (option (pair int int)))
        "hop count"
        (Some (20, 0))
        (Bundle.hop_count decoded);
      Alcotest.(check (option (of_pp Bundle.pp_eid)))
        "previous node"
        (Some (Bundle.Ipn (5L, 0L)))
        (Bundle.previous_node decoded);
      Alcotest.(check (option int64))
        "bundle age" (Some 500L) (Bundle.age decoded);
      Alcotest.(check (option string))
        "payload" (Some "multi-block bundle") (Bundle.payload decoded)
  | Error e -> Alcotest.failf "decode failed: %a" Bundle.pp_error e

let test_block_type_codes () =
  (* RFC 9171 Section 4.4: verify block type integer codes *)
  Alcotest.(check int) "Payload" 1 (Bundle.int_of_block_type Bundle.Payload);
  Alcotest.(check int)
    "Previous_node" 6
    (Bundle.int_of_block_type Bundle.Previous_node);
  Alcotest.(check int)
    "Bundle_age" 7
    (Bundle.int_of_block_type Bundle.Bundle_age);
  Alcotest.(check int)
    "Hop_count" 10
    (Bundle.int_of_block_type Bundle.Hop_count);
  Alcotest.(check int)
    "Other 200" 200
    (Bundle.int_of_block_type (Bundle.Other 200))

let suite =
  ( "bundle",
    [
      Alcotest.test_case "EID dtn:none" `Quick test_eid_dtn_none;
      Alcotest.test_case "EID dtn URI" `Quick test_eid_dtn;
      Alcotest.test_case "EID ipn" `Quick test_eid_ipn;
      Alcotest.test_case "bundle flags" `Quick test_bundle_flags;
      Alcotest.test_case "block flags" `Quick test_block_flags;
      Alcotest.test_case "CRC type" `Quick test_crc_type;
      Alcotest.test_case "simple bundle" `Quick test_simple_bundle;
      Alcotest.test_case "roundtrip" `Quick test_bundle_roundtrip;
      Alcotest.test_case "no CRC" `Quick test_bundle_no_crc;
      Alcotest.test_case "CRC-16" `Quick test_bundle_crc16;
      Alcotest.test_case "fragment fields" `Quick test_fragment_fields;
      Alcotest.test_case "fragment offset zero" `Quick test_fragment_offset_zero;
      Alcotest.test_case "hop count block" `Quick test_hop_count_block;
      Alcotest.test_case "previous node block" `Quick test_previous_node_block;
      Alcotest.test_case "bundle age block" `Quick test_bundle_age_block;
      Alcotest.test_case "do_not_fragment flag" `Quick test_do_not_fragment_flag;
      Alcotest.test_case "CRC-32C roundtrip" `Quick test_crc32c_roundtrip;
      Alcotest.test_case "CRC-32C compute" `Quick test_crc32c_compute;
      Alcotest.test_case "multiple extension blocks" `Quick
        test_multiple_extension_blocks;
      Alcotest.test_case "block type codes" `Quick test_block_type_codes;
    ] )