;; This file is part of eris-cl.
;; Copyright (C) 2022 Piotr Szarmański

;; eris-cl is free software: you can redistribute it and/or modify it under the
;; terms of the GNU Lesser General Public License as published by the Free
;; Software Foundation, either version 3 of the License, or (at your option) any
;; later version.

;; eris-cl is distributed in the hope that it will be useful, but WITHOUT ANY
;; WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS FOR
;; A PARTICULAR PURPOSE. See the GNU Lesser General Public LIcense for more details.

;; You should have received a copy of the GNU Lesser General Public LIcense along with
;; eris-cl. If not, see <https://www.gnu.org/licenses/>.

(in-package :eris)

(deftype block-size ()
        `(member 1024 32768))

(defconstant 32kib 32768)
(defconstant 1kib 1024)

(define-constant null-secret (make-array 32 :element-type 'octet :initial-element 0)
  :test #'equalp
  :documentation
  "32-byte null vector.")

(defun make-nonce (level)
  (let ((nonce (make-array 12 :element-type 'octet :initial-element 0)))
    (setf (aref nonce 0) level)
    nonce))


(defun encrypt-block (input secret)
  (declare (type octet-vector input secret))
  (let ((mac (ironclad:make-mac :blake2-mac secret :digest-length 32))
        (rk (make-array 64 :element-type 'octet))) ;; reference-key pair
    (ironclad:update-mac mac input)
    (ironclad:produce-mac mac :digest rk :digest-start 32) ;; get key
    (ironclad:encrypt-in-place (ironclad:make-cipher :chacha :mode :stream :key (subseq rk 32 64)
                                                             :initialization-vector null-secret)
                               input) ;; encrypt block
    (ironclad:digest-sequence :blake2/256 input :digest rk) ;; get reference
    rk))


(defun encrypt-internal-block (input nonce)
  (declare (type octet-vector input))
  (let ((rk (make-array 64 :element-type 'octet))) ;; reference-key pair
    (ironclad:digest-sequence :blake2/256 input :digest rk :digest-start 32);; get key
    (ironclad:encrypt-in-place (ironclad:make-cipher :chacha :mode :stream :key (subseq rk 32 64)
                                                             :initialization-vector nonce)
                               input) ;; encrypt block
    (ironclad:digest-sequence :blake2/256 input :digest rk) ;; get reference
    rk))

(defun decrypt-block (input rk &optional (nonce null-secret))
  (declare (type octet-vector input rk nonce))
  (ironclad:decrypt-in-place
   (ironclad:make-cipher :chacha :mode :stream :key (subseq rk 32 64) :initialization-vector nonce)
   input)
  input)


(defclass read-capability ()
  ((block-size :initarg :block-size
               :accessor block-size
               :type block-size
               :documentation "A value of either 1024 or 1kb blocks.") 
   (level :initarg :level :accessor level :type octet)
   (root-reference-pair :initarg :reference-pair :accessor reference-pair))
  (:documentation "Class representing the concept of an ERIS read capability."))

(-> read-capability-to-octets (read-capability) (octet-vector 66))
(defun read-capability-to-octets (read-capability)
  "Convert a read-capability object to its standard binary representation. Returns
a (simple-array (unsigned-byte 8)) object."
  (declare (type read-capability read-capability))
  (let ((cap (make-array 66 :element-type 'octet)))
    (case (block-size read-capability) ;; This depends on the version of the standard
      (1024 (setf (aref cap 0) #x0a))
      (32768 (setf (aref cap 0) #x0f)))
    (setf (aref cap 1) (level read-capability))
    (replace cap (reference-pair read-capability) :start1 2)
    cap))

(-> octets-to-read-capability ((octet-vector 66)) (values read-capability &optional))
(defun octets-to-read-capability (octets)
  "Convert the standard binary representation for ERIS read capabilities into a
read-capability object. Returns the read-capability.

An ERIS:VERSION-MISMATCH condition may be signaled if the corresponding
versioning bytes are not supported by eris-cl."
  (declare (type (octet-vector 66) octets))
  (let ((capability (make-instance 'read-capability)))
    (setf (block-size capability)
          (case (aref octets 0)
            (#x0a 1kib)
            (#x0f 32kib)
            (t (error 'version-mismatch))))
    (setf (level capability)
          (aref octets 1))
    (setf (reference-pair capability)
          (let ((kr (make-array 64 :element-type 'octet)))
            (replace kr octets :start2 2)
            kr)) ;; TODO CHECK CORRECTNESS
    capability))

(-> read-capability-to-urn (read-capability) string)
(defun read-capability-to-urn (capability)
  "Convert a read-capability object into a URN string."
  (declare (type read-capability capability))
  (concatenate 'string
               "urn:eris:"
               (bytes-to-base32-unpadded (read-capability-to-octets capability))))

(-> urn-to-read-capability (string) (values read-capability &optional))
(defun urn-to-read-capability (urn)
  "Convert a urn:eris URN string into a read-capability object."
  (declare (type string urn))
  (octets-to-read-capability (base32-to-bytes-unpadded (subseq urn (1+ (position #\: urn :from-end t))))))

(-> reference-to-block-urn ((octet-vector 32)) string)
(defun reference-to-block-urn (reference)
  "Convert a 32-byte block reference into a URN string."
  (declare (type (octet-vector 32) reference))
  (concatenate 'string "urn:blake2b:" (bytes-to-base32-unpadded reference)))

(-> block-urn-to-reference (string) (values (octet-vector 32) &optional))
(defun block-urn-to-reference (urn)
  "Convert a urn:blake2b URN string into a 32-byte block reference vector."
  (declare (type string urn))
  (base32-to-bytes-unpadded (subseq urn (1+ (position #\: urn :from-end t)))))

;; This macro assumes that there are variables BLOCK, SECRET and OUTPUT-FUNCTION
;; in the lexenv.

(defmacro output-block (rks i)
  `(let ((rk (encrypt-block block secret)))
     (setf (svref ,rks ,i) rk)
     (funcall output-function block (subseq rk 0 32))))


(defmacro output-internal-block (ref-vector nonce)
  `(let ((rk (encrypt-internal-block block ,nonce)))
     (vector-push-extend rk ,ref-vector)
     (funcall output-function block (subseq rk 0 32))))


;; These CHUNK- functions are written in order to allow processing files in
;; parallel. 

(defun chunk-array (array block-size output-function secret &key pad (start 0) (end (length array)))
  "Split (SIMPLE-ARRAY (UNSIGNED-BYTE 8) that is a multiple of BLOCK-SIZE into
chunks, output them and collect references. Returns a vector of references.
START and END behave as expected.

Pass PAD as T if the output should be padded."
  (declare (type block-size block-size)
           (type octet-vector array))
  (when (and (not pad) (zerop (- end start))) ;; need this because of the loop unrolling 
    (return-from chunk-array (make-array 0 :element-type 'octet-vector)))
  
  (let ((length (- end start)))
    (let ((blocks (if pad
                      (/ (+ length (- block-size (mod length block-size))) block-size)
                      (/ length block-size))))
      (let ((block (make-octet-vector block-size))
            (rks (make-array blocks :element-type 'octet-vector :initial-element null-secret)))
        (loop for i from 0 below (1- blocks)
              do (progn
                   (replace block array :start2 (+ start (* block-size i)))
                   (setf block (output-block rks i))))
        ;; handle last block
        (replace block array :start2 (+ start (* block-size (1- blocks))))
        (when pad
          (setf (aref block (mod length block-size)) #x80)
          (fill block 0 :start (1+ (mod length block-size))))
        (output-block rks (1- blocks))
        rks)))) 


;; Implementation note: This is CHUNK-ARRAY but copypasted with (LENGTH ARRAY)
;; changed to LENGTH and REPLACE changed to READ-SEQUENCE. It is, however, more
;; memory-efficient than reading a file into an array and then chunking it.

(defun chunk-stream (stream block-size output-function length secret &key pad)
  "Like CHUNK-ARRAY, but with streams. LENGTH indicates the amount of bytes to
read and should be a multiple of BLOCK-SIZE unless PAD is T."
  (declare (type block-size block-size)
           (type integer length))
  (when (and (not pad) (zerop length)) ;; need this because of the loop unrolling 
    (return-from chunk-stream (make-array 0 :element-type 'octet-vector)))
  (let ((blocks (if pad
                    (/ (+ length  (- block-size (mod length block-size))) block-size)
                    (/ length  block-size))))
    (let ((block (make-octet-vector block-size))
          ;; initialize with null-secret to please SBCL
          (rks (make-array blocks :element-type 'octet-vector :initial-element null-secret)))
      (loop for i from 0 below (1- blocks)
            do (progn
                 (read-sequence block stream )
                 (setf block (output-block rks i))))
      ;; handle last block
      (read-sequence block stream)
      (when pad
        (setf (aref block (mod length block-size)) #x80)
        (fill block 0 :start (1+ (mod length block-size))))
      (output-block rks (1- blocks))
      rks)))

(defgeneric eris-encode (input block-size output-function &key secret)
  (:documentation
   "Encode an INPUT into BLOCK-SIZE (32kib or 1kib) blocks, that are output using
the function OUTPUT-FUNCTION. This function wil be called with two arguments: an
encoded block and a 32-byte reference octet vector, and it MUST return
a (SIMPLE-ARRAY (UNSIGNED-BYTE 8)) of equal size to the one given, which will be
destructively modified. Returns a read-capability object.

A SECRET can be provided to use with encryption; otherwise the null secret (* 32 0x0)
is used."))

(defmethod eris-encode ((input vector) block-size output-function &key (secret null-secret))
  (declare (type block-size block-size)
           (type function output-function)
           (type (octet-vector 32) secret))
  (eris-create-tree
   (chunk-array input block-size output-function secret :pad t)
   block-size output-function))

(defmethod eris-encode ((input pathname) block-size output-function &key (secret null-secret))
  (declare (type block-size block-size)
           (type function output-function)
           (type (octet-vector 32) secret))
  (with-open-file (f input :element-type 'octet)
    (eris-create-tree
     (chunk-stream f block-size output-function (file-length f) secret :pad t)
     block-size output-function)))


(defmethod eris-encode ((input file-stream) block-size output-function &key (secret null-secret))
  (declare (type block-size block-size)
           (type function output-function)
           (type (octet-vector 32) secret))
  (eris-create-tree
   (chunk-stream input block-size output-function
                 (- (file-length input) (file-position input))
                 secret :pad t)
   block-size output-function))

;; This is the odd one out because it is not possible to determine the length of
;; a non-file stream (modulo broadcast and synonym streams).
(defmethod eris-encode ((input stream) block-size output-function &key (secret null-secret))
  (declare (type block-size block-size)
           (type function output-function)
           (type (octet-vector 32) secret))
  (let ((reference-vector (make-array 16 :adjustable t :fill-pointer 0))
        (block (make-array block-size :element-type 'octet :initial-element 0)))
    (declare (type octet-vector block))
    (loop for bytes-read = (read-sequence block input)
          for i = 0 then (incf i)
          if (< bytes-read block-size)
            do (progn (setf (aref block bytes-read) #x80)
                      (fill block 0 :start (1+ bytes-read)))
          do (progn (setf block (let ((rk (encrypt-block block secret)))
                                  (vector-push-extend rk reference-vector)
                                  (funcall output-function block (subseq rk 0 32)))))
          until (< bytes-read block-size))
    (eris-create-tree reference-vector block-size output-function)))

(defun eris-create-tree (reference-vector block-size output-function)
  (declare (type block-size block-size)
           (type function output-function))
  (loop with block-keys = (/ block-size 64)
        with level = 0
        with reference-vector-l = (make-array 16 :adjustable t :fill-pointer 0)
        for nonce = (make-nonce (1+ level))
        when (eql (length reference-vector) 1)
          do (return (make-instance 'read-capability
                                    :reference-pair (aref reference-vector 0)
                                    :level level
                                    :block-size block-size)) 
        do (progn
             (incf level)
             ;; loop across the key-reference vector and build the tree
             (loop with block = (make-array block-size :element-type 'octet :initial-element 0)
                   for rk across reference-vector
                   with i = 0 
                   when (eql i block-keys)
                     do (progn (setf block (output-internal-block reference-vector-l nonce))
                               (setf i 0)
                               (fill block 0))
                   do (progn (replace block rk :start1 (* 64 i))
                             (incf i))
                   finally (unless (zerop i)
                             ;; If i is zero, then the amount of blocks is just
                             ;; right. Otherwise add a final unfinished block.
                             
                             (output-internal-block reference-vector-l nonce)))
             (setf reference-vector reference-vector-l)
             (setf reference-vector-l (make-array 16 :adjustable t :fill-pointer 0)))))