;; 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 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)

(defvar *decode-safety-checks* t
  "If set to nil, disable computationally expensive safety checks that ensure the
fetched blocks are what they claim to be. Only disable this if the blocks are
fetched from a trusted party.")

(defmacro hash-check (block hash)
  `(when *decode-safety-checks*
     (let ((hash (ironclad:digest-sequence :blake2/256 ,block)))
       (unless (equalp ,hash hash)
         (error 'hash-mismatch :reference ,hash :hash hash )))))

(defmacro execute-fetch-function (fetch-function &rest args)
  `(restart-case (funcall ,fetch-function ,@args)
     (use-value (value) value)))

(defun key-reference-null? (kr)
  (and (equalp (reference kr) null-secret)
       (equalp (key kr) null-secret)))

(defclass high-block ()
  ((data :initarg :data :accessor data :type octet-vector)
   (position :initarg :position :accessor pos :type integer
             :documentation "Position relative to the higher block.")
   (level :initarg :level :accessor level :type octet)))

(defclass buffer ()
  ((data :initarg :data :accessor data :type octet-vector)
   (pos :initarg :pos :accessor pos :type fixnum)
   (eof :initarg :eof :accessor eof :type fixnum
        :documentation "Either the position of the EOF in the buffer, or the length of the buffer.")))

(defclass eris-decode-stream (fundamental-binary-input-stream)
  ((position :initform 0 :accessor pos :initarg :position :type integer)
   (get-block :accessor get-block :initarg :get-block :type function)
   (block-size :accessor block-size :type block-size :initarg :block-size)
   (root :accessor root :type high-block :initarg :root :documentation "A list of blocks, starting from the root to the level 1 block.")
   (buffer :accessor buffer :initarg :buffer :type buffer)
   (eof :accessor eof :initarg :eof :type integer)
   (nonce-array :accessor nonce-array :initarg :nonce-array :type simple-array))
  (:documentation "Class representing the stream object that decodes an ERIS read capability."))

(defun initialize-nonce-array (level)
  (let ((array (make-array (1+ level))))
    (loop for i from 0 to level
          do (setf (aref array i) (make-nonce i)))
    array))

(defun unpad-block (block)
  (let ((padding (position #x80 block :from-end t)))
    (unless padding
      (error 'padding-error))
    (unless (loop for i across (nsubseq block (1+ padding))
                  always (zerop i))
      (error 'padding-error))
    padding))

(defun unpad-buffer (buffer)
  (with-slots (eof data) buffer
    (setf eof (unpad-block data))))

(defun local-range (level block-size)
  "Given a level and a block-size, determine the range of bytes that it
represents."
  (declare (type fixnum level)
           (type block-size block-size))
  (ecase block-size
    (1024 (expt 2 (+ 10 (* level 4))))
    (32768 (expt 2 (+ 15 (* level 9))))))

(defun find-eof (root get-block block-size level)
  "Find the end of file in a given tree."
  ;; The standard states that:
  ;; "If 64 bytes of zeroes are encountered the rest of the node MUST be checked to be all zeroes"
  ;; This procedure processes the block from right to left so it's not really applicable here

  (flet ((find-last-reference (block)
           (loop for i from (1- (/ block-size 64)) downto 0
                 for key-reference = (octets-to-reference-pair block (* 64 i))
                 unless (key-reference-null? key-reference)
                   return (values key-reference i))))
    (loop with position = 0
          for local-level = level then (1- local-level)
          for block = root then (funcall get-block (reference key-reference) (key key-reference) (make-nonce local-level))
          for (key-reference block-pos) = (multiple-value-list (find-last-reference block))
          if (eql 0 local-level)
            do (setf position (+ position (unpad-block block)))
            and return position
          else
            do (setf position (+ position (* block-pos (local-range (1- local-level) block-size)))))))

(defun advance-next-block (stream)
  "Advance to the next block. 

This function walks the tree in order to get the next block, which is then put
in the buffer object of the STREAM. It sets the EOF indicator of the buffer if
it is necessary and sets the position in the buffer to 0."
  (declare (optimize (speed 3) (debug 0))
           (type eris-decode-stream stream)) 
  (with-slots (buffer position eof block-size root get-block nonce-array) stream
    (declare (type integer position eof)
             (type block-size block-size)
             #+sbcl (sb-ext:muffle-conditions sb-ext:compiler-note))
    (labels ((process-block (blocks mod level) ;; A recursive function that 1+'s each block until the 
               (declare (type fixnum mod))
               (let* ((current (car blocks))
                      (next-pos (setf (pos current)
                                      (mod (the fixnum (1+ (the fixnum (pos current)))) mod))))
                 (declare (type high-block current)
                          (type fixnum next-pos))
                 (when (zerop next-pos) ;; if 1+ mod blocks results in a zero, then update the next block
                   (process-block (cdr blocks) mod (1+ level)))
                 (let ((kr (octets-to-reference-pair
                            (data (cadr blocks))
                            (the fixnum (* next-pos 64)))))
                   (setf (data current)
                         (funcall get-block (reference kr) (key kr) (svref nonce-array level)))))))
      (process-block (reverse root) (/ block-size 64) 0) ;; process the blocks from level 0 up
      ;; setf the buffer data to the level 0 block
      (setf (data buffer) (data (car (last root)))
            (pos buffer) 0)
      ;; if the EOF block is reached, unpad it.
      (when (= (- eof (mod eof block-size)) position)
        (unpad-buffer buffer)))))

(defun read-to-seq (sequence buf &key (start 0) (end (length sequence)) stream)
  (declare (optimize (speed 3) (debug 0))
           (type array-total-size end start)
           #+sbcl (sb-ext:muffle-conditions sb-ext:compiler-note))
  (with-slots (data pos eof) buf
    (declare (type (integer 0 32768) pos eof)
             (type octet-vector data))
    (let ((buffer-bytes (- eof pos))
          (seq-bytes (- end start)))
      (declare (dynamic-extent buffer-bytes seq-bytes)
               (type integer seq-bytes)
               (type (integer 0 32768) buffer-bytes))
      (cond
        ((> buffer-bytes seq-bytes) ;; seq is smaller than buffer case
         (replace sequence data :start1 start :end1 end :start2 pos)
         (setf pos (+ pos seq-bytes)
               (pos stream) (+ (pos stream) seq-bytes))
         (+ start seq-bytes))
        ((eql eof (length data)) ;; seq is larger or equal than buffer case
         (replace sequence data :start1 start :end1 end :start2 pos)
         (setf (pos stream) (+ (pos stream) buffer-bytes))
         (advance-next-block stream) 
         (read-to-seq sequence buf :start (+ start buffer-bytes) :end end :stream stream))
        (t ;; if there is an eof in the buffer
         (replace sequence data :start1 start :end1 end :start2 pos :end2 eof)
         (setf pos (+ pos buffer-bytes)
               (pos stream) (+ (pos stream) buffer-bytes))
         (+ start buffer-bytes))))))


(defun reupdate-block (stream new-pos)
  "Update the blocks of the stream according to new-pos. Used for random access and initialization."
  (unless (root stream)
    (error 'eof :eof new-pos :position (pos stream)))
  (with-slots (block-size root get-block buffer nonce-array eof) stream
    (loop with update-tree = nil
          for blocks = root then (cdr blocks)
          for level = (level (car blocks)) ;; then (decf level)
          for position = new-pos then (mod position (local-range (1- level) block-size)) ;; local position
          for block-position = (floor (/ position (local-range (1- level) block-size))) ;; the position of the block
          for lower-block = (second blocks)
          when (> block-position (/ block-size 64))
            do (error 'eof :eof new-pos :position (pos stream))
          until (eq nil (cdr blocks))
          when (or update-tree (not (eq block-position (pos lower-block))))
            do (let ((rk (octets-to-reference-pair
                          (data (car blocks))
                          (* 64 block-position))))
                 (when (key-reference-null? rk)
                   (error 'eof :eof new-pos :position (pos stream)))
                 (setf (data lower-block) (funcall get-block (reference rk) (key rk) (svref nonce-array (1- level)))
                       (pos lower-block) block-position
                       update-tree t))
          finally (when update-tree ;; if anything changed
                    (setf (pos stream) new-pos
                          (data buffer) (data (car (last root)))
                          (pos buffer) (mod new-pos block-size))
                    (if (< (- eof (mod eof block-size)) new-pos)
                        (unpad-buffer buffer)
                        (setf (eof buffer) block-size))))))

(defun initialize-high-blocks (level)
  (case level
    (-1 nil)
    (t (cons (make-instance 'high-block :level level :position -1 :data (make-array 0 :element-type 'octet))
             (initialize-high-blocks (1- level))))))



(defun eris-decode (read-capability fetch-function &key (cache-capacity 2048))
  "Using the FETCH-FUNCTION, return a stream that decodes the READ-CAPABILITY.
This stream implements the Gray streams protocol.

Fetch-function must be a function with one argument, the reference octet, which
returns a (simple-array (unsigned-byte 8)) containing the block. The block will
be destructively modified, so you MUST provide a fresh array every time. If a
hash-table is used, a (copy-seq) needs to be done on the return value of
gethash.

The keyword argument CACHE-CAPACITY indicates the amount of blocks stored in the
cache. It may be NIL to turn off caching entirely.

The conditions ERIS:MISSING-BLOCK, ERIS:PADDING, ERIS:HASH-MISMATCH and ERIS:EOF
may be signaled by operations on the stream. Any condition signaled from within
FETCH-FUNCTION is not handled."
  (declare (type read-capability read-capability)
           (type function fetch-function)
           (type (or integer null) cache-capacity))
  (with-slots (level block-size root-reference-pair) read-capability
    (let* ((get-block (if cache-capacity ;; handle caching
                          (cached-lambda (:cache-class 'lru-cache 
                                                       :capacity cache-capacity
                                                       :table (make-hash-table :size (1+ cache-capacity) :test #'equalp))
                                         (reference key nonce)
                                         (declare (type octet-vector reference key nonce))
                                         (let ((block (execute-fetch-function fetch-function reference)))
                                           (declare (type octet-vector block))
                                           (hash-check block reference)
                                           (decrypt-block block key nonce)))
                        (lambda (reference key nonce)
                          (declare (type octet-vector reference key nonce)) 
                          (let ((block (execute-fetch-function fetch-function reference)))
                            (declare (type octet-vector block))
                            (hash-check block reference)
                            (decrypt-block block key nonce)))))
           (root (funcall get-block (reference root-reference-pair) (key root-reference-pair) (make-nonce level))))
      ;; "Implementations MUST verify the key appearing in the read capability
      ;; if level of encoded content is larger than 0."
      (when (> level 0) (hash-check root (key root-reference-pair)))
      (case level
        ;; Treat level 0 blocks specially, since those are just a single buffer.
        (0 (let ((stream
                   (make-instance 'eris-decode-stream
                                  :buffer (make-instance 'buffer
                                                         :data (make-array block-size :element-type 'octet)
                                                         :pos 0
                                                         :eof -1)
                                  :get-block (lambda (reference key)
                                               (declare (ignore key))
                                               (error 'missing-block :reference reference))
                                  :block-size block-size
                                  :root nil
                                  :eof (find-eof root get-block block-size level)
                                  :nonce-array (initialize-nonce-array 0))))
             (replace (data (buffer stream)) root)
             (unpad-buffer (buffer stream))
             stream))
        (t (make-instance 'eris-decode-stream
                          :buffer (make-instance 'buffer
                                                 :data (make-array block-size :element-type 'octet)
                                                 :pos -1
                                                 :eof block-size)
                          :get-block get-block
                          :block-size block-size
                          :root (cons (make-instance 'high-block :level level
                                                                 :data root
                                                                 :position 0)
                                      (initialize-high-blocks (1- level)))
                          :eof (find-eof root get-block block-size level)
                          :nonce-array (initialize-nonce-array level)))))))

(defmethod stream-file-position ((stream eris-decode-stream))
  "Provides the file position of the stream. This method is setf-able in order to
change the position."
  (pos stream))

(defmethod stream-read-sequence ((stream eris-decode-stream) seq start end &key)
  (when (minusp (pos (buffer stream)))
    ;; initializes the buffer
    (reupdate-block stream (pos stream)))
  (with-slots (buffer position) stream
    (read-to-seq seq buffer :start start :end end :stream stream)))

(defmethod stream-read-byte ((stream eris-decode-stream))
  (when (minusp (pos (buffer stream)))
    ;; initializes the buffer
    (reupdate-block stream (pos stream)))
  (with-slots (position buffer block-size) stream
    (with-slots (pos eof data) buffer
      (cond
        ((eql pos block-size)
         (advance-next-block stream)
         (stream-read-byte stream))
        ((eql pos eof)
         :eof)
        (t (prog1 (aref data pos)
             (incf pos)
             (incf position)))))))

(defmethod stream-element-type ((stream eris-decode-stream))
  '(unsigned-byte 8))

(defun eris-file-length (stream)
  "This is the equivalent of \"file-length\" for eris-decode-stream."
  (eof stream))

(defmethod (setf stream-file-position) (set-position (stream eris-decode-stream))
  (with-slots (position block-size buffer eof) stream
    (when set-position
      (case set-position
        (:end (setf set-position eof))
        (:start (setf set-position 0)))
      (let ((buffer-pos (mod set-position block-size)))
        (if (< set-position eof)
            (cond
              ;; If the pos is within the buffer (and initialized): 
              ((and (<= 0 (- set-position (- position (pos buffer))) (1- block-size))
                    (not (minusp (pos buffer))))
               (setf (pos buffer) buffer-pos
                     position set-position))
              
              (t (reupdate-block stream set-position)
                 (setf (pos buffer) buffer-pos)))
            (error 'eof :eof eof :position position))))
    set-position))