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

(** Bundle Protocol Version 7 (RFC 9171, CCSDS 734.2-B-1). *)

module Cbor = Cbort.Cbor
module Cbor_rw = Cbort.Rw

(* {1 Endpoint IDs} *)

type eid = Dtn_none | Dtn of string | Ipn of int64 * int64

let pp_eid fmt = function
  | Dtn_none -> Fmt.pf fmt "dtn:none"
  | Dtn s -> Fmt.pf fmt "dtn:%s" s
  | Ipn (node, service) -> Fmt.pf fmt "ipn:%Ld.%Ld" node service

let eid_to_cbor = function
  | Dtn_none -> Cbor.Array [ Cbor.int 1; Cbor.int 0 ]
  | Dtn s -> Cbor.Array [ Cbor.int 1; Cbor.Text s ]
  | Ipn (node, service) ->
      Cbor.Array
        [ Cbor.int 2; Cbor.Array [ Cbor.int64 node; Cbor.int64 service ] ]

let err_unknown_eid_scheme n = Error (Fmt.str "unknown EID scheme: %Ld" n)

let eid_of_cbor cbor =
  match Cbor.to_array cbor with
  | Some [ scheme; ssp ] -> (
      match Cbor.to_int64 scheme with
      | Some 1L -> (
          (* dtn scheme *)
          match Cbor.to_int64 ssp with
          | Some 0L -> Ok Dtn_none
          | _ -> (
              match Cbor.to_text ssp with
              | Some s -> Ok (Dtn s)
              | None -> Error "invalid dtn SSP"))
      | Some 2L -> (
          (* ipn scheme *)
          match Cbor.to_array ssp with
          | Some [ node; service ] -> (
              match (Cbor.to_int64 node, Cbor.to_int64 service) with
              | Some n, Some s -> Ok (Ipn (n, s))
              | _ -> Error "invalid ipn node/service")
          | _ -> Error "ipn SSP must be array of 2")
      | Some n -> err_unknown_eid_scheme n
      | None -> Error "EID scheme must be uint")
  | _ -> Error "EID must be array of 2"

(* {1 CRC Types} *)

type crc_type = No_crc | Crc16 | Crc32c

let pp_crc_type fmt = function
  | No_crc -> Fmt.string fmt "none"
  | Crc16 -> Fmt.string fmt "CRC-16"
  | Crc32c -> Fmt.string fmt "CRC-32C"

let int_of_crc_type = function No_crc -> 0 | Crc16 -> 1 | Crc32c -> 2

let crc_type_of_int = function
  | 0 -> Ok No_crc
  | 1 -> Ok Crc16
  | 2 -> Ok Crc32c
  | n -> Error n

(* {1 Bundle Processing Flags} *)

type flags = {
  is_fragment : bool;
  is_admin_record : bool;
  must_not_fragment : bool;
  ack_requested : bool;
  status_time_requested : bool;
  report_reception : bool;
  report_forwarding : bool;
  report_delivery : bool;
  report_deletion : bool;
}

let default_flags =
  {
    is_fragment = false;
    is_admin_record = false;
    must_not_fragment = false;
    ack_requested = false;
    status_time_requested = false;
    report_reception = false;
    report_forwarding = false;
    report_delivery = false;
    report_deletion = false;
  }

let pp_flags fmt f =
  let flags = [] in
  let flags = if f.is_fragment then "fragment" :: flags else flags in
  let flags = if f.is_admin_record then "admin" :: flags else flags in
  let flags = if f.must_not_fragment then "no-frag" :: flags else flags in
  let flags = if f.ack_requested then "ack" :: flags else flags in
  let flags =
    if f.status_time_requested then "status-time" :: flags else flags
  in
  let flags = if f.report_reception then "report-recv" :: flags else flags in
  let flags = if f.report_forwarding then "report-fwd" :: flags else flags in
  let flags = if f.report_delivery then "report-dlv" :: flags else flags in
  let flags = if f.report_deletion then "report-del" :: flags else flags in
  Fmt.pf fmt "[%s]" (String.concat ", " (List.rev flags))

let int_of_flags f =
  let n = 0 in
  let n = if f.is_fragment then n lor 0x000001 else n in
  let n = if f.is_admin_record then n lor 0x000002 else n in
  let n = if f.must_not_fragment then n lor 0x000004 else n in
  let n = if f.ack_requested then n lor 0x000020 else n in
  let n = if f.status_time_requested then n lor 0x000040 else n in
  let n = if f.report_reception then n lor 0x004000 else n in
  let n = if f.report_forwarding then n lor 0x010000 else n in
  let n = if f.report_delivery then n lor 0x020000 else n in
  let n = if f.report_deletion then n lor 0x040000 else n in
  n

let flags_of_int n =
  {
    is_fragment = n land 0x000001 <> 0;
    is_admin_record = n land 0x000002 <> 0;
    must_not_fragment = n land 0x000004 <> 0;
    ack_requested = n land 0x000020 <> 0;
    status_time_requested = n land 0x000040 <> 0;
    report_reception = n land 0x004000 <> 0;
    report_forwarding = n land 0x010000 <> 0;
    report_delivery = n land 0x020000 <> 0;
    report_deletion = n land 0x040000 <> 0;
  }

(* {1 Block Processing Flags} *)

type block_flags = {
  replicate_in_fragment : bool;
  report_if_unprocessed : bool;
  delete_if_unprocessed : bool;
  discard_if_unprocessed : bool;
}

let block_flags_default =
  {
    replicate_in_fragment = false;
    report_if_unprocessed = false;
    delete_if_unprocessed = false;
    discard_if_unprocessed = false;
  }

let pp_block_flags fmt f =
  let flags = [] in
  let flags = if f.replicate_in_fragment then "replicate" :: flags else flags in
  let flags = if f.report_if_unprocessed then "report" :: flags else flags in
  let flags = if f.delete_if_unprocessed then "delete" :: flags else flags in
  let flags = if f.discard_if_unprocessed then "discard" :: flags else flags in
  Fmt.pf fmt "[%s]" (String.concat ", " (List.rev flags))

let int_of_block_flags f =
  let n = 0 in
  let n = if f.replicate_in_fragment then n lor 0x01 else n in
  let n = if f.report_if_unprocessed then n lor 0x02 else n in
  let n = if f.delete_if_unprocessed then n lor 0x04 else n in
  let n = if f.discard_if_unprocessed then n lor 0x10 else n in
  n

let block_flags_of_int n =
  {
    replicate_in_fragment = n land 0x01 <> 0;
    report_if_unprocessed = n land 0x02 <> 0;
    delete_if_unprocessed = n land 0x04 <> 0;
    discard_if_unprocessed = n land 0x10 <> 0;
  }

(* {1 Creation Timestamp} *)

type timestamp = { time : int64; seq : int64 }

let pp_timestamp fmt ts = Fmt.pf fmt "{ time=%Ld; seq=%Ld }" ts.time ts.seq

(* {1 Primary Block} *)

type primary_block = {
  version : int;
  flags : flags;
  crc_type : crc_type;
  destination : eid;
  source : eid;
  report_to : eid;
  creation_timestamp : timestamp;
  lifetime : int64;
  fragment_offset : int64 option;
  total_adu_length : int64 option;
}

let pp_primary_block fmt pb =
  Fmt.pf fmt
    "@[<v 2>primary {@ version=%d;@ flags=%a;@ crc=%a;@ dest=%a;@ src=%a;@ \
     report_to=%a;@ timestamp=%a;@ lifetime=%Ld%a%a@ }@]"
    pb.version pp_flags pb.flags pp_crc_type pb.crc_type pp_eid pb.destination
    pp_eid pb.source pp_eid pb.report_to pp_timestamp pb.creation_timestamp
    pb.lifetime
    (fun fmt -> function
      | Some off -> Fmt.pf fmt ";@ fragment_offset=%Ld" off | None -> ())
    pb.fragment_offset
    (fun fmt -> function
      | Some len -> Fmt.pf fmt ";@ total_adu_length=%Ld" len | None -> ())
    pb.total_adu_length

(* {1 Canonical Blocks} *)

type block_type =
  | Payload
  | Previous_node
  | Bundle_age
  | Hop_count
  | Other of int

let pp_block_type fmt = function
  | Payload -> Fmt.string fmt "payload"
  | Previous_node -> Fmt.string fmt "previous-node"
  | Bundle_age -> Fmt.string fmt "bundle-age"
  | Hop_count -> Fmt.string fmt "hop-count"
  | Other n -> Fmt.pf fmt "other(%d)" n

let int_of_block_type = function
  | Payload -> 1
  | Previous_node -> 6
  | Bundle_age -> 7
  | Hop_count -> 10
  | Other n -> n

let block_type_of_int = function
  | 1 -> Payload
  | 6 -> Previous_node
  | 7 -> Bundle_age
  | 10 -> Hop_count
  | n -> Other n

type block_data =
  | Payload_data of string
  | Previous_node_data of eid
  | Bundle_age_data of int64
  | Hop_count_data of { limit : int; count : int }
  | Unknown_data of string

let pp_block_data fmt = function
  | Payload_data b -> Fmt.pf fmt "payload[%d bytes]" (String.length b)
  | Previous_node_data eid -> Fmt.pf fmt "previous-node=%a" pp_eid eid
  | Bundle_age_data age -> Fmt.pf fmt "age=%Ld ms" age
  | Hop_count_data { limit; count } -> Fmt.pf fmt "hop-count=%d/%d" count limit
  | Unknown_data b -> Fmt.pf fmt "unknown[%d bytes]" (String.length b)

type canonical_block = {
  block_type : block_type;
  block_number : int;
  flags : block_flags;
  crc_type : crc_type;
  data : block_data;
}

let pp_canonical_block fmt cb =
  Fmt.pf fmt "@[<v 2>block#%d {@ type=%a;@ flags=%a;@ crc=%a;@ data=%a@ }@]"
    cb.block_number pp_block_type cb.block_type pp_block_flags cb.flags
    pp_crc_type cb.crc_type pp_block_data cb.data

(* {1 Bundle} *)

type t = { primary : primary_block; blocks : canonical_block list }

let pp fmt bundle =
  Fmt.pf fmt "@[<v 2>bundle {@ %a" pp_primary_block bundle.primary;
  List.iter (fun b -> Fmt.pf fmt "@ %a" pp_canonical_block b) bundle.blocks;
  Fmt.pf fmt "@ }@]"

(* {1 Errors} *)

type error =
  | Invalid_version of int
  | Invalid_crc_type of int
  | Invalid_eid_scheme of int
  | Invalid_block_type of int
  | Missing_payload_block
  | Crc_mismatch of { expected : string; computed : string }
  | Cbor_error of string
  | Truncated of string

let pp_hex fmt s = String.iter (fun c -> Fmt.pf fmt "%02x" (Char.code c)) s

let pp_error fmt = function
  | Invalid_version v -> Fmt.pf fmt "invalid version: %d (expected 7)" v
  | Invalid_crc_type n -> Fmt.pf fmt "invalid CRC type: %d" n
  | Invalid_eid_scheme n -> Fmt.pf fmt "invalid EID scheme: %d" n
  | Invalid_block_type n -> Fmt.pf fmt "invalid block type: %d" n
  | Missing_payload_block -> Fmt.pf fmt "missing payload block"
  | Crc_mismatch { expected; computed } ->
      Fmt.pf fmt "CRC mismatch: expected %a, computed %a" pp_hex expected pp_hex
        computed
  | Cbor_error msg -> Fmt.pf fmt "CBOR error: %s" msg
  | Truncated msg -> Fmt.pf fmt "truncated: %s" msg

open Result.Syntax

(* Helper to convert Option to Result with Cbor_error *)
let require msg = function Some v -> Ok v | None -> Error (Cbor_error msg)

(* {1 CRC Computation} *)

let crc16_x25 = Crc.crc16_x25
let crc32c = Crc.crc32c

let compute_crc crc_type data =
  match crc_type with
  | No_crc -> ""
  | Crc16 ->
      let crc = crc16_x25 data in
      let buf = Bytes.create 2 in
      Bytes.set_uint16_be buf 0 crc;
      Bytes.to_string buf
  | Crc32c ->
      let crc = crc32c data in
      let buf = Bytes.create 4 in
      Bytes.set_int32_be buf 0 (Int32.of_int crc);
      Bytes.to_string buf

(* {1 CBOR Helpers} *)

let cbor_encode (v : Cbor.t) : string =
  let buf = Buffer.create 256 in
  let writer = Bytesrw.Bytes.Writer.of_buffer buf in
  let enc = Cbor_rw.encoder writer in
  Cbor_rw.write_cbor enc v;
  Cbor_rw.flush_encoder enc;
  Buffer.contents buf

let cbor_decode (s : string) : Cbor.t option =
  try
    let reader = Bytesrw.Bytes.Reader.of_string s in
    let dec = Cbor_rw.decoder reader in
    Some (Cbor_rw.read_cbor dec)
  with End_of_file | Failure _ | Invalid_argument _ -> None

(* {1 Encoding} *)

let encode_timestamp ts = Cbor.Array [ Cbor.int64 ts.time; Cbor.int64 ts.seq ]

let encode_block_data = function
  | Payload_data b -> Cbor.Bytes b
  | Previous_node_data eid ->
      let cbor = eid_to_cbor eid in
      Cbor.Bytes (cbor_encode cbor)
  | Bundle_age_data age -> Cbor.Bytes (cbor_encode (Cbor.int64 age))
  | Hop_count_data { limit; count } ->
      let cbor = Cbor.Array [ Cbor.int limit; Cbor.int count ] in
      Cbor.Bytes (cbor_encode cbor)
  | Unknown_data b -> Cbor.Bytes b

let encode_primary pb =
  let elements =
    [
      Cbor.int pb.version;
      Cbor.int (int_of_flags pb.flags);
      Cbor.int (int_of_crc_type pb.crc_type);
      eid_to_cbor pb.destination;
      eid_to_cbor pb.source;
      eid_to_cbor pb.report_to;
      encode_timestamp pb.creation_timestamp;
      Cbor.int64 pb.lifetime;
    ]
  in
  let elements =
    match pb.fragment_offset with
    | Some off -> elements @ [ Cbor.int64 off ]
    | None -> elements
  in
  let elements =
    match pb.total_adu_length with
    | Some len -> elements @ [ Cbor.int64 len ]
    | None -> elements
  in
  (* Add CRC placeholder if needed *)
  let elements =
    match pb.crc_type with
    | No_crc -> elements
    | Crc16 -> elements @ [ Cbor.Bytes (String.make 2 '\x00') ]
    | Crc32c -> elements @ [ Cbor.Bytes (String.make 4 '\x00') ]
  in
  let arr = Cbor.Array elements in
  (* Compute CRC if needed *)
  match pb.crc_type with
  | No_crc -> arr
  | _ -> (
      let encoded = cbor_encode arr in
      let crc = compute_crc pb.crc_type encoded in
      let crc_len = String.length crc in
      (* Replace the last crc_len bytes with computed CRC *)
      let encoded_bytes = Bytes.of_string encoded in
      Bytes.blit_string crc 0 encoded_bytes
        (Bytes.length encoded_bytes - crc_len)
        crc_len;
      (* Re-decode to get proper CBOR structure *)
      match cbor_decode (Bytes.to_string encoded_bytes) with
      | Some v -> v
      | None -> arr)

let encode_canonical cb =
  let elements =
    [
      Cbor.int (int_of_block_type cb.block_type);
      Cbor.int cb.block_number;
      Cbor.int (int_of_block_flags cb.flags);
      Cbor.int (int_of_crc_type cb.crc_type);
      encode_block_data cb.data;
    ]
  in
  let elements =
    match cb.crc_type with
    | No_crc -> elements
    | Crc16 -> elements @ [ Cbor.Bytes (String.make 2 '\x00') ]
    | Crc32c -> elements @ [ Cbor.Bytes (String.make 4 '\x00') ]
  in
  let arr = Cbor.Array elements in
  match cb.crc_type with
  | No_crc -> arr
  | _ -> (
      let encoded = cbor_encode arr in
      let crc = compute_crc cb.crc_type encoded in
      let crc_len = String.length crc in
      let encoded_bytes = Bytes.of_string encoded in
      Bytes.blit_string crc 0 encoded_bytes
        (Bytes.length encoded_bytes - crc_len)
        crc_len;
      match cbor_decode (Bytes.to_string encoded_bytes) with
      | Some v -> v
      | None -> arr)

let to_cbor bundle =
  let primary_cbor = encode_primary bundle.primary in
  let blocks_cbor = List.map encode_canonical bundle.blocks in
  (* BPv7 uses indefinite-length array, but cbort doesn't have Iarray,
     so we use regular Array which should be compatible *)
  Cbor.Array (primary_cbor :: blocks_cbor)

let encode bundle = cbor_encode (to_cbor bundle)
let write enc bundle = Cbor_rw.write_cbor enc (to_cbor bundle)

(* {1 Decoding} *)

let decode_timestamp cbor =
  match Cbor.to_array cbor with
  | Some [ time_cbor; seq_cbor ] -> (
      match (Cbor.to_int64 time_cbor, Cbor.to_int64 seq_cbor) with
      | Some time, Some seq -> Ok { time; seq }
      | _ -> Error (Cbor_error "timestamp fields must be uint"))
  | _ -> Error (Cbor_error "timestamp must be array of 2")

let decode_block_data block_type data_bytes =
  let require_cbor msg = cbor_decode data_bytes |> require msg in
  match block_type with
  | Payload -> Ok (Payload_data data_bytes)
  | Previous_node ->
      let* cbor = require_cbor "failed to decode previous node CBOR" in
      let* eid =
        eid_of_cbor cbor |> Result.map_error (fun msg -> Cbor_error msg)
      in
      Ok (Previous_node_data eid)
  | Bundle_age ->
      let* cbor = require_cbor "failed to decode bundle age CBOR" in
      let* age = Cbor.to_int64 cbor |> require "bundle age must be uint" in
      Ok (Bundle_age_data age)
  | Hop_count -> (
      let* cbor = require_cbor "failed to decode hop count CBOR" in
      let* elements =
        Cbor.to_array cbor |> require "hop count must be array of 2"
      in
      match elements with
      | [ limit_cbor; count_cbor ] ->
          let* limit =
            Cbor.to_int64 limit_cbor |> require "hop count fields must be uint"
          in
          let* count =
            Cbor.to_int64 count_cbor |> require "hop count fields must be uint"
          in
          Ok
            (Hop_count_data
               { limit = Int64.to_int limit; count = Int64.to_int count })
      | _ -> Error (Cbor_error "hop count must be array of 2"))
  | Other _ -> Ok (Unknown_data data_bytes)

let decode_primary cbor =
  let* elements = Cbor.to_array cbor |> require "primary block must be array" in
  if List.length elements < 8 then
    Error (Truncated "primary block needs at least 8 elements")
  else
    let get i = List.nth elements i in
    let* version = Cbor.to_int64 (get 0) |> require "version must be uint" in
    if version <> 7L then Error (Invalid_version (Int64.to_int version))
    else
      let* flags_int = Cbor.to_int64 (get 1) |> require "flags must be uint" in
      let flags = flags_of_int (Int64.to_int flags_int) in
      let* crc_int = Cbor.to_int64 (get 2) |> require "CRC type must be uint" in
      let* crc_type =
        crc_type_of_int (Int64.to_int crc_int)
        |> Result.map_error (fun n -> Invalid_crc_type n)
      in
      let* destination =
        eid_of_cbor (get 3) |> Result.map_error (fun msg -> Cbor_error msg)
      in
      let* source =
        eid_of_cbor (get 4) |> Result.map_error (fun msg -> Cbor_error msg)
      in
      let* report_to =
        eid_of_cbor (get 5) |> Result.map_error (fun msg -> Cbor_error msg)
      in
      let* creation_timestamp = decode_timestamp (get 6) in
      let* lifetime =
        Cbor.to_int64 (get 7) |> require "lifetime must be uint"
      in
      let fragment_offset, total_adu_length =
        if flags.is_fragment && List.length elements >= 10 then
          (Cbor.to_int64 (get 8), Cbor.to_int64 (get 9))
        else (None, None)
      in
      Ok
        {
          version = 7;
          flags;
          crc_type;
          destination;
          source;
          report_to;
          creation_timestamp;
          lifetime;
          fragment_offset;
          total_adu_length;
        }

let decode_canonical cbor =
  let* elements =
    Cbor.to_array cbor |> require "canonical block must be array"
  in
  if List.length elements < 5 then
    Error (Truncated "canonical block needs at least 5 elements")
  else
    let get i = List.nth elements i in
    let* type_int =
      Cbor.to_int64 (get 0) |> require "block type must be uint"
    in
    let block_type = block_type_of_int (Int64.to_int type_int) in
    let* block_number =
      Cbor.to_int64 (get 1) |> require "block number must be uint"
    in
    let* flags_int = Cbor.to_int64 (get 2) |> require "flags must be uint" in
    let flags = block_flags_of_int (Int64.to_int flags_int) in
    let* crc_int = Cbor.to_int64 (get 3) |> require "CRC type must be uint" in
    let* crc_type =
      crc_type_of_int (Int64.to_int crc_int)
      |> Result.map_error (fun n -> Invalid_crc_type n)
    in
    let* data_bytes =
      Cbor.to_bytes (get 4) |> require "block data must be bytes"
    in
    let* data = decode_block_data block_type data_bytes in
    Ok
      {
        block_type;
        block_number = Int64.to_int block_number;
        flags;
        crc_type;
        data;
      }

let of_cbor cbor =
  let* elements = Cbor.to_array cbor |> require "bundle must be array" in
  match elements with
  | [] -> Error (Truncated "bundle must have at least primary block")
  | primary_cbor :: blocks_cbor ->
      let* primary = decode_primary primary_cbor in
      let rec decode_blocks acc = function
        | [] -> Ok (List.rev acc)
        | b :: rest ->
            let* block = decode_canonical b in
            decode_blocks (block :: acc) rest
      in
      let* blocks = decode_blocks [] blocks_cbor in
      let has_payload = List.exists (fun b -> b.block_type = Payload) blocks in
      if has_payload then Ok { primary; blocks }
      else Error Missing_payload_block

let decode s =
  match cbor_decode s with
  | Some cbor -> of_cbor cbor
  | None -> Error (Cbor_error "failed to decode CBOR")

let read dec =
  try
    let cbor = Cbor_rw.read_cbor dec in
    of_cbor cbor
  with e -> Error (Cbor_error (Printexc.to_string e))

(* {1 Constructors} *)

let v ?(flags = default_flags) ?(crc_type = Crc32c) ?report_to
    ?(lifetime = 86400000L) ~source ~destination ~creation_timestamp ~payload ()
    =
  let report_to = match report_to with Some e -> e | None -> source in
  let primary =
    {
      version = 7;
      flags;
      crc_type;
      destination;
      source;
      report_to;
      creation_timestamp;
      lifetime;
      fragment_offset = None;
      total_adu_length = None;
    }
  in
  let payload_block =
    {
      block_type = Payload;
      block_number = 1;
      flags = block_flags_default;
      crc_type;
      data = Payload_data payload;
    }
  in
  { primary; blocks = [ payload_block ] }

let payload bundle =
  List.find_map
    (fun b -> match b.data with Payload_data p -> Some p | _ -> None)
    bundle.blocks

let previous_node bundle =
  List.find_map
    (fun b ->
      match b.data with Previous_node_data eid -> Some eid | _ -> None)
    bundle.blocks

let age bundle =
  List.find_map
    (fun b -> match b.data with Bundle_age_data age -> Some age | _ -> None)
    bundle.blocks

let hop_count bundle =
  List.find_map
    (fun b ->
      match b.data with
      | Hop_count_data { limit; count } -> Some (limit, count)
      | _ -> None)
    bundle.blocks