module Noun = Nock_lib.Noun
module Serial = Nock_lib.Serial

let bytes_of_list ints =
  let len = List.length ints in
  let buf = Bytes.create len in
  List.iteri
    (fun i byte ->
      if byte < 0 || byte > 0xFF then invalid_arg "byte out of range";
      Bytes.set buf i (Char.chr byte))
    ints;
  buf

let hex_of_bytes bytes =
  let buf = Buffer.create (Bytes.length bytes * 2) in
  Bytes.iter
    (fun ch -> Buffer.add_string buf (Printf.sprintf "%02x" (Char.code ch)))
    bytes;
  Buffer.contents buf

let atom_int n = Noun.atom (Z.of_int n)

let atom_of_bytes ints =
  let rec loop acc shift = function
    | [] -> acc
    | byte :: rest ->
        let part = Z.shift_left (Z.of_int byte) (8 * shift) in
        loop (Z.add acc part) (shift + 1) rest
  in
  Noun.atom (loop Z.zero 0 ints)

let atom_of_ascii s =
  let bytes = List.init (String.length s) (fun i -> Char.code s.[i]) in
  atom_of_bytes bytes

let pair a b = Noun.cell a b
let trel a b c = Noun.cell a (pair b c)
let qual a b c d = Noun.cell a (trel b c d)

type case = {
  name : string;
  noun : Noun.noun;
  expected : bytes;
}

let case name noun bytes =
  { name; noun; expected = bytes_of_list bytes }

let deep_case =
  (* Direct translation of C: [[[[1 [[2 [[3 [[4 [[5 [6 [7 [[8 0] 0]]]] 0]] 0]] 0]] 0]] 0] 0] *)
  case "deep"
     (pair
       (pair
         (pair
           (atom_int 1)
           (pair
             (pair
               (atom_int 2)
               (pair
                 (pair
                   (atom_int 3)
                   (pair
                     (pair
                       (atom_int 4)
                       (pair
                         (trel
                           (atom_int 5)
                           (atom_int 6)
                           (pair
                             (atom_int 7)
                             (pair
                               (pair (atom_int 8) (atom_int 0))
                               (atom_int 0)
                              )
                            )
                          )
                        (atom_int 0))
                      )
                    (atom_int 0))
                  )
                (atom_int 0))
              )
            (atom_int 0))
          )
        (atom_int 0))
      (atom_int 0))
    
    
    [ 0x15; 0x17; 0xb2; 0xd0; 0x85; 0x59; 0xb8; 0x61;
      0x87; 0x5f; 0x10; 0x54; 0x55; 0x05 ]

let wide_case =
  let inp =
    [ 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x01 ]
  in
  case "wide" (atom_of_bytes inp)
    [ 0x00; 0x0c; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x00; 0x00; 0x08 ]

let date_case =
  let inp =
    [ 0x00; 0x00; 0x00; 0x00; 0x00; 0x00; 0x0c; 0xa8;
      0xab; 0x60; 0xef; 0x2d; 0x0d; 0x00; 0x00; 0x80 ]
  in
  case "date" (atom_of_bytes inp)
    [ 0x00; 0x02; 0x00; 0x00; 0x00; 0x00; 0x00; 0x00;
      0x18; 0x50; 0x57; 0xc1; 0xde; 0x5b; 0x1a; 0x00;
      0x00; 0x00; 0x01 ]

let alpha_case =
  let a = atom_of_ascii "abcdefjhijklmnopqrstuvwxyz" in
  let noun = qual a (atom_int 2) (atom_int 3) a in
  case "alpha" noun
    [ 0x01; 0xf8; 0x0c; 0x13; 0x1b; 0x23; 0x2b; 0x33;
      0x53; 0x43; 0x4b; 0x53; 0x5b; 0x63; 0x6b; 0x73;
      0x7b; 0x83; 0x8b; 0x93; 0x9b; 0xa3; 0xab; 0xb3;
      0xbb; 0xc3; 0xcb; 0xd3; 0x87; 0x0c; 0x3d; 0x09 ]

let cases =
  [
    case "0" (atom_int 0) [ 0x02 ];
    case "1" (atom_int 1) [ 0x0c ];
    case "2" (atom_int 2) [ 0x48 ];
    case "%fast" (atom_of_ascii "fast")
      [ 0xc0; 0x37; 0x0b; 0x9b; 0xa3; 0x03 ];
    case "%full" (atom_of_ascii "full")
      [ 0xc0; 0x37; 0xab; 0x63; 0x63; 0x03 ];
    case "[0 0]"
      (pair (atom_int 0) (atom_int 0))
      [ 0x29 ];
    case "[1 1]"
      (pair (atom_int 1) (atom_int 1))
      [ 0x31; 0x03 ];
    case "[1 2]"
      (pair (atom_int 1) (atom_int 2))
      [ 0x31; 0x12 ];
    case "[2 3]"
      (pair (atom_int 2) (atom_int 3))
      [ 0x21; 0xd1 ];
    case "[%fast %full]"
      (pair (atom_of_ascii "fast") (atom_of_ascii "full"))
      [ 0x01; 0xdf; 0x2c; 0x6c; 0x8e; 0x0e; 0x7c; 0xb3; 0x3a; 0x36; 0x36 ];
    case "[1 1 1]"
      (pair (atom_int 1) (pair (atom_int 1) (atom_int 1)))
      [ 0x71; 0xcc ];
    case "[1 2 3]"
      (trel (atom_int 1) (atom_int 2) (atom_int 3))
      [ 0x71; 0x48; 0x34 ];
    case "[%fast %full %fast]"
      (let fast = atom_of_ascii "fast" in
       let full = atom_of_ascii "full" in
       pair fast (pair full fast))
      [ 0x01; 0xdf; 0x2c; 0x6c; 0x8e; 0x1e; 0xf0; 0xcd;
        0xea; 0xd8; 0xd8; 0x93 ];
    case "[[1 2] 3]"
      (pair (pair (atom_int 1) (atom_int 2)) (atom_int 3))
      [ 0xc5; 0x48; 0x34 ];
    case "[[1 2] [1 2] [1 2]]"
      (let one = atom_int 1 in
       let two = atom_int 2 in
       let sub = pair one two in
       trel sub sub sub)
      [ 0xc5; 0xc8; 0x26; 0x27; 0x01 ];
    case "[[0 0] [[0 0] 1 1] 1 1]"
      (pair
        (pair (atom_int 0) (atom_int 0))
        (pair
          (pair
            (pair (atom_int 0) (atom_int 0))
            (pair (atom_int 1) (atom_int 1)))
          (pair (atom_int 1) (atom_int 1))))
      [ 0xa5; 0x35; 0x19; 0xf3; 0x18; 0x05 ];
    deep_case;
    wide_case;
    date_case;
    alpha_case;
  ]

let check_case { name; noun; expected } () =
  Printf.printf "case %s:\n" name;
  let jammed = Serial.jam noun in
  let jam_hex = hex_of_bytes jammed in
  let expect_hex = hex_of_bytes expected in
  Printf.printf "  jam    : 0x%s\n" jam_hex;
  Printf.printf "  expect : 0x%s\n%!" expect_hex;
  Alcotest.(check string)
    (name ^ " jam bytes")
    expect_hex
    jam_hex;
  let cued = Serial.cue expected in
  Alcotest.(check bool)
    (name ^ " cue roundtrip")
    true
    (Noun.equal noun cued)

let tests =
  List.map (fun case -> case.name, `Quick, check_case case) cases

let () = Alcotest.run ~verbose:true "serial-vectors" [ "cases", tests ]