;; 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 versqion.

;; 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 General Public License for more details.

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

(in-package :eris)

(defun pad (input block-size)
  (declare (type (simple-array (unsigned-byte 8)) input)
           (type integer block-size))
  (let* ((pad-size (- block-size (mod (length input) block-size)))
         (padded-input (adjust-array input (+ pad-size (length input)) :initial-element 0)))
    (replace padded-input input)
    (setf (aref padded-input (length input)) #x80)
    padded-input))

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

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

(eval-when
    (:execute :load-toplevel :compile-toplevel)
  (defclass reference-pair ()
    ((reference :initarg :reference :accessor reference :type (simple-array (unsigned-byte 8) 32))
     (key :initarg :key :accessor key :type (simple-array (unsigned-byte 8) 32)))))

(define-constant null-secret (make-array 32 :element-type '(unsigned-byte 8) :initial-element 0) :test #'equalp)

(defun reference-pair-to-octets (pair buf &optional (start 0))
    (replace buf (reference pair) :start1 start)
    (replace buf (key pair) :start1 (+ 32 start)))

(defun octets-to-reference-pair (octets &optional (start 0))
  (declare (type (simple-array (unsigned-byte 8)) octets))
  (make-instance 'reference-pair :key (subseq octets (+ 32 start) (+ 64 start))
                                 :reference (subseq octets start (+ 32 start))))


(defun compute-reference (block)
  (declare (type (simple-array (unsigned-byte 8)) block))
  (let ((reference (make-array 32 :element-type '(unsigned-byte 8))))
    (ironclad:digest-sequence :blake2/256 block :digest reference)
    reference))

(declaim (inline make-nonce))

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

(defun encrypt-block (input secret reference)
  (declare (type (simple-array (unsigned-byte 8)) input secret reference))
  (let ((mac (ironclad:make-mac :blake2-mac secret :digest-length 32))
        (key (make-array 32 :element-type '(unsigned-byte 8))))
    (ironclad:update-mac mac input)
    (ironclad:produce-mac mac :digest key)
    ;; NOT the IETF chacha. Need to fix this by patching ironclad. Maybe.
    (ironclad:encrypt-in-place (ironclad:make-cipher :chacha :mode :stream :key key
                                                             :initialization-vector null-secret)
                               input)
    (ironclad:digest-sequence :blake2/256 input :digest reference)
    (make-instance 'reference-pair :key key :reference reference)))

(defun encrypt-internal-block (input reference nonce)
  (declare (type (simple-array (unsigned-byte 8)) input secret reference))
  (let ((key (make-array 32 :element-type '(unsigned-byte 8))))
    (ironclad:digest-sequence :blake2/256 input :digest key)
    ;; NOT the IETF chacha. Need to fix this by patching ironclad. Maybe.
    (ironclad:encrypt-in-place (ironclad:make-cipher :chacha :mode :stream :key key
                                                             :initialization-vector nonce)
                               input)
    (ironclad:digest-sequence :blake2/256 input :digest reference)
    (make-instance 'reference-pair :key key :reference reference)))

(defun decrypt-block (input key &optional (nonce null-secret))
  (ironclad:decrypt-in-place
   (ironclad:make-cipher :chacha :mode :stream :key key :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 (unsigned-byte 8))
   (root-reference-pair :initarg :reference-pair :accessor reference-pair)))



(defun read-capability-to-octets (read-capability)
  (declare (type read-capability read-capability))
  (let ((cap (make-array 66 :element-type '(unsigned-byte 8))))
    (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))
    (reference-pair-to-octets (reference-pair read-capability) cap 2)))

(defun octets-to-read-capability (octets)
  (declare (type (simple-array (unsigned-byte 8) (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)
          (octets-to-reference-pair octets 2))
    capability))

(defun read-capability-to-urn (capability)
  (concatenate 'string
               "urn:eris:"
               (bytes-to-base32-unpadded (read-capability-to-octets capability))))

(defun urn-to-read-capability (urn)
  (octets-to-read-capability (base32-to-bytes-unpadded (subseq urn (1+ (position #\: urn :from-end t))))))

(defun reference-to-block-urn (reference)
  (declare (type (simple-array (unsigned-byte 8) (32)) reference))
  (concatenate 'string "urn:blake2b:" (bytes-to-base32-unpadded reference)))

(defun block-urn-to-reference (urn)
  (declare (type string urn))
  (base32-to-bytes-unpadded (subseq urn (1+ (position #\: urn :from-end t)))))

(defvar *output-hashmap* nil)

(defmacro output-block (ref-vector &rest expr)
  `(let ((reference (compute-reference block)))
     (if hash-output
         (unless (gethash reference *output-hashmap*)
           (let ((rk (encrypt-block block secret reference)))
             (vector-push-extend rk ,ref-vector)
             (funcall output-function block (reference rk))
             (setf (gethash reference *output-hashmap*) t)))
         (let ((rk (encrypt-block block secret reference)))
           (vector-push-extend rk ,ref-vector)
           (funcall output-function block (reference rk))))
     ,@expr))

(defmacro output-internal-block (ref-vector nonce &rest expr)
  `(let ((reference (compute-reference block)))
     (if hash-output
         (unless (gethash reference *output-hashmap*)
           (let ((rk (encrypt-internal-block block reference ,nonce)))
             (vector-push-extend rk ,ref-vector)
             (funcall output-function block (reference rk))
             (setf (gethash reference *output-hashmap*) t)))
         (let ((rk (encrypt-internal-block block reference ,nonce)))
           (vector-push-extend rk ,ref-vector)
           (funcall output-function block (reference rk))))
     ,@expr))

(defgeneric eris-encode (input block-size output-function &key secret hash-output)
  (: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.

An optional 32-byte secret can be passed for additional encryption using the
:secret keyword."))

(defmethod eris-encode ((input simple-array) block-size output-function &key (secret null-secret) hash-output)
  (declare (type block-size block-size)
           (type function output-function)
           (type (simple-array (unsigned-byte 8) (32)) secret))

  (setf input (pad input block-size)) 
  
  (let ((reference-vector (make-array 16 :adjustable t :fill-pointer 0))
        (*output-hashmap* (if hash-output (make-hash-table :test #'equalp) nil)))
    (loop for block = (make-array block-size :element-type '(unsigned-byte 8) :initial-element 0)
          for i = 0 then (incf i)
          until (= (length input) (* i block-size)) 
          do (progn
               (replace block input :start2 (* i block-size))
               (output-block reference-vector nil)))
    (eris-create-tree reference-vector block-size output-function :hash-output hash-output)))

(defmethod eris-encode ((input stream) block-size output-function &key (secret null-secret) hash-output)
  (declare (type block-size block-size)
           (type function output-function)
           (type (simple-array (unsigned-byte 8) (32)) secret))
  (let ((reference-vector (make-array 16 :adjustable t :fill-pointer 0))
        (*output-hashmap* (if hash-output (make-hash-table :test #'equalp) nil)))
    (loop for block = (make-array block-size :element-type '(unsigned-byte 8) :initial-element 0)
          for bytes-read = (read-sequence block input)
          for i = 0 then (incf i)
          if (< bytes-read block-size)
            do (setf (aref block bytes-read) #x80)
          do (output-block reference-vector nil)
          until (< bytes-read block-size))
    (eris-create-tree reference-vector block-size output-function :hash-output hash-output)))

(defun eris-create-tree (reference-vector block-size output-function &key hash-output)
  (declare (type block-size block-size)
           (type function output-function)
           (type (simple-array (unsigned-byte 8) (32)) secret))
  (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 '(unsigned-byte 8) :initial-element 0)
                   for rk across reference-vector
                   with i = 0
                   when (eql i block-keys)
                   do (output-internal-block reference-vector-l nonce
                                    (setf block (make-array block-size :element-type '(unsigned-byte 8) :initial-element 0)
                                          i 0))
                   do (progn (reference-pair-to-octets rk block (* 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)))))