;;; Chickadee Game Toolkit
;;; Copyright © 2021 David Thompson <dthompson2@worcester.edu>
;;;
;;; Chickadee is free software: you can redistribute it and/or modify
;;; it under the terms of the GNU General Public License as published
;;; by the Free Software Foundation, either version 3 of the License,
;;; or (at your option) any later version.
;;;
;;; Chickadee 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 this program.  If not, see
;;; <http://www.gnu.org/licenses/>.

;;; Commentary:
;;
;; PNG image loading.
;;
;;; Code:

(define-module (chickadee image png)
  #:use-module (chickadee config)
  #:use-module (ice-9 binary-ports)
  #:use-module (ice-9 exceptions)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-9)
  #:use-module (srfi srfi-11)
  #:use-module (system foreign)
  #:export (load-png))


;;;
;;; Enums
;;;

(define PNG_TRANSFORM_IDENTITY 0)

(define PNG_COLOR_MASK_PALETTE 1)
(define PNG_COLOR_MASK_COLOR 2)
(define PNG_COLOR_MASK_ALPHA 4)

(define PNG_COLOR_TYPE_GRAY 0)
(define PNG_COLOR_TYPE_PALETTE (logior PNG_COLOR_MASK_COLOR PNG_COLOR_MASK_PALETTE))
(define PNG_COLOR_TYPE_RGB PNG_COLOR_MASK_COLOR)
(define PNG_COLOR_TYPE_RGB_ALPHA (logior PNG_COLOR_MASK_COLOR PNG_COLOR_MASK_ALPHA))
(define PNG_COLOR_TYPE_GRAY_ALPHA PNG_COLOR_MASK_ALPHA)


;;;
;;; Low-level bindings
;;;

(define libpng-func
  (let ((lib (dynamic-link* %libpng)))
    (lambda (return-type function-name arg-types)
      (pointer->procedure return-type
                          (dynamic-func function-name lib)
                          arg-types))))

(define-syntax-rule (define-foreign name return-type func-name arg-types)
  (define name
    (libpng-func return-type func-name arg-types)))

(define-foreign png-sig-cmp
  int "png_sig_cmp" (list '* size_t size_t))

(define-foreign png-create-read-struct
  '* "png_create_read_struct" '(* * * *))

(define-foreign png-destroy-read-struct
  void "png_destroy_read_struct" '(* * *))

(define-foreign png-create-info-struct
  '* "png_create_info_struct" '(*))

(define-foreign png-set-read-fn
  void "png_set_read_fn" '(* * *))

(define-foreign png-set-sig-bytes
  void "png_set_sig_bytes" (list '* int))

(define-foreign png-read-png
  void "png_read_png" (list '* '* int '*))

(define-foreign png-get-rows
  '* "png_get_rows" '(* *))

(define-foreign png-get-rowbytes
  size_t "png_get_rowbytes" '(* *))

(define-foreign png-get-image-width
  unsigned-int "png_get_image_width" '(* *))

(define-foreign png-get-image-height
  unsigned-int "png_get_image_height" '(* *))

(define-foreign png-get-color-type
  int8 "png_get_color_type" '(* *))

(define-foreign png-get-bit-depth
  int8 "png_get_bit_depth" '(* *))


;;;
;;; High-level bindings
;;;

(define (display-png-read-struct handle port)
  (display "#<png-read-struct>" port))

(define-wrapped-pointer-type <png-read-struct> png-read-struct?
  wrap-png-read-struct unwrap-png-read-struct display-png-read-struct)

(define (display-png-info-struct handle port)
  (display "#<png-info-struct>" port))

(define-wrapped-pointer-type <png-info-struct> png-info-struct?
  wrap-png-info-struct unwrap-png-info-struct display-png-info-struct)

(define (on-error read-struct message)
  (display (current-error-port) message))

(define (on-warning read-struct message)
  (display (current-error-port) message))

(define (create-read-struct)
  (let ((ptr (png-create-read-struct (string->pointer %libpng-version)
                                     %null-pointer
                                     (procedure->pointer void on-error '(* *))
                                     (procedure->pointer void on-warning '(* *)))))
    (if (null-pointer? ptr)
        (raise-exception
         (make-exception-with-message
          "could not create png read struct"))
        (wrap-png-read-struct ptr))))

(define (create-info-struct read-struct)
  (let ((ptr (png-create-info-struct (unwrap-png-read-struct read-struct))))
    (if (null-pointer? ptr)
        (raise-exception
         (make-exception-with-message
          "could not create png info struct"))
        (wrap-png-info-struct ptr))))

(define (destroy-read-struct read-struct info-struct)
  (png-destroy-read-struct (unwrap-png-read-struct read-struct)
                           (unwrap-png-info-struct info-struct)
                           %null-pointer))

(define (image-width read-struct info-struct)
  (png-get-image-width (unwrap-png-read-struct read-struct)
                       (unwrap-png-info-struct info-struct)))

(define (image-height read-struct info-struct)
  (png-get-image-height (unwrap-png-read-struct read-struct)
                        (unwrap-png-info-struct info-struct)))

(define (color-type read-struct info-struct)
  (png-get-color-type (unwrap-png-read-struct read-struct)
                      (unwrap-png-info-struct info-struct)))

(define (bit-depth read-struct info-struct)
  (png-get-bit-depth (unwrap-png-read-struct read-struct)
                     (unwrap-png-info-struct info-struct)))

(define (bytes-per-row read-struct info-struct)
  (png-get-rowbytes (unwrap-png-read-struct read-struct)
                    (unwrap-png-info-struct info-struct)))

(define (rows read-struct info-struct)
  (png-get-rows (unwrap-png-read-struct read-struct)
                (unwrap-png-info-struct info-struct)))

(define (set-sig-bytes! read-struct n)
  (png-set-sig-bytes (unwrap-png-read-struct read-struct) n))

(define (set-read-function! read-struct func-ptr)
  (png-set-read-fn (unwrap-png-read-struct read-struct)
                   %null-pointer
                   func-ptr))

(define (read-png read-struct info-struct)
  (png-read-png (unwrap-png-read-struct read-struct)
                (unwrap-png-info-struct info-struct)
                PNG_TRANSFORM_IDENTITY
                %null-pointer))

(define (png-header? header)
  (zero? (png-sig-cmp (bytevector->pointer header) 0 8)))

(define* (dereference-pointer pointer #:optional (offset 0))
  (make-pointer
   (bytevector-uint-ref (pointer->bytevector pointer
                                             (sizeof uintptr_t)
                                             (* offset (sizeof uintptr_t)))
                        0
                        (native-endianness)
                        (sizeof uintptr_t))))

(define (unpacker color-type bit-depth)
  ;; 16 bit color channels need to be reduced to an 8 bit value.
  (define (u16->u8 n)
    (round (* (/ n 65536) 255)))
  (cond
   ((= color-type PNG_COLOR_TYPE_GRAY)
    (case bit-depth
      ;; For bit depths of 1, 2, and 4, multiple pixels are stored
      ;; within a single byte.  So, we need to use some bit shifting
      ;; magic to extract the grayscale values.  First, we find the
      ;; byte that contains the pixel.  Then we make a bitmask (using
      ;; a bit shift to the left) that selects the pixel bits (1, 2,
      ;; or 4 bits) from within the byte and apply a bitwise AND to
      ;; zero out the other, irrelevant bits.  To get a value within
      ;; the target grayscale range, we shift the bits back to the
      ;; right.  Finally, to arrive at a value for an 8-bit color
      ;; channel, we scale the value as needed (by 255 for 1 bit
      ;; grayscale, for example.)
      ((1)
       (lambda (row x)
         (let* ((byte (bytevector-u8-ref row (floor (/ x 8))))
                (bit (modulo x 8))
                (gray (* (ash (logand byte (ash #b00000001 bit)) (- bit)) 255)))
           (values gray gray gray 255))))
      ((2)
       (lambda (row x)
         (let* ((byte (bytevector-u8-ref row (floor (/ x 4))))
                (bit (* (modulo x 4) 2))
                (gray (* (ash (logand byte (ash #b00000011 bit)) (- bit)) 85)))
           (values gray gray gray 255))))
      ((4)
       (lambda (row x)
         (let* ((byte (bytevector-u8-ref row (floor (/ x 2))))
                (bit (* (modulo x 2) 4))
                (gray (* (ash (logand byte (ash #b00001111 bit)) (- bit)) 17)))
           (values gray gray gray 255))))
      ((8)
       (lambda (row x)
         (let ((gray (bytevector-u8-ref row x)))
           (values gray gray gray 255))))
      ((16)
       (lambda (row x)
         (let* ((offset (* x 2))
                (gray (u16->u8
                       (bytevector-u16-ref row offset
                                           (endianness big)))))
           (values gray gray gray 255))))))
   ((= color-type PNG_COLOR_TYPE_PALETTE)
    ;; TODO: Handle palettes.
    (raise-exception
     (make-exception-with-message
      "PNG images using color palettes are not supported")))
   ((= color-type PNG_COLOR_TYPE_RGB)
    (case bit-depth
      ((8)
       (lambda (row x)
         (let ((offset (* x 3)))
           (values (bytevector-u8-ref row offset)
                   (bytevector-u8-ref row (+ offset 1))
                   (bytevector-u8-ref row (+ offset 2))
                   255))))
      ((16)
       (lambda (row x)
         (let ((offset (* x 6)))
           ;; PNGs encode 16-bit color data in big endian format.
           (values (u16->u8
                    (bytevector-u16-ref row offset
                                        (endianness big)))
                   (u16->u8
                    (bytevector-u16-ref row (+ offset 2)
                                        (endianness big)))
                   (u16->u8
                    (bytevector-u16-ref row (+ offset 4)
                                        (endianness big)))
                   255))))))
   ((= color-type PNG_COLOR_TYPE_RGB_ALPHA)
    (case bit-depth
      ((8) #f) ; We can take a convenient shortcut in this case.
      ((16)
       (lambda (row x)
         (let* ((offset (* x 8))
                (r (u16->u8
                    (bytevector-u16-ref row offset
                                        (endianness big))))
                (g (u16->u8
                    (bytevector-u16-ref row (+ offset 2)
                                        (endianness big))))
                (b (u16->u8
                    (bytevector-u16-ref row (+ offset 4)
                                        (endianness big))))
                (a (u16->u8
                    (bytevector-u16-ref row (+ offset 6)
                                        (endianness big)))))
           (values r g b a))))))
   ((= color-type PNG_COLOR_TYPE_GRAY_ALPHA)
    (case bit-depth
      ((8)
       (lambda (row x)
         (let* ((offset (* x 2))
                (gray (bytevector-u8-ref row offset))
                (alpha (bytevector-u8-ref row (+ offset 1))))
           (values gray gray gray alpha))))
      ((16)
       (lambda (row x)
         (let* ((offset (* x 4))
                (gray (u16->u8
                       (bytevector-u16-ref row offset
                                           (endianness big))))
                (alpha (u16->u8
                        (bytevector-u16-ref row (+ offset 2)
                                            (endianness big)))))
           (values gray gray gray alpha))))))))

(define (make-png-reader port)
  (procedure->pointer void
                      (lambda (read-struct ptr length)
                        (let ((bv (pointer->bytevector ptr length)))
                          (get-bytevector-n! port bv 0 length)))
                      (list '* '* size_t)))

(define (load-png file-name)
  (if (file-exists? file-name)
      (call-with-input-file file-name
        (lambda (port)
          (let* ((read-struct (create-read-struct))
                 (info-struct (create-info-struct read-struct)))
            (dynamic-wind
              (lambda () #t)
              (lambda ()
                (unless (png-header? (get-bytevector-n port 8))
                  (raise-exception
                   (make-exception (make-external-error)
                                   (make-exception-with-message
                                    (string-append "not a PNG file: " file-name)))))
                ;; Tell libpng that we've already read the header.
                (set-sig-bytes! read-struct 8)
                ;; Need to set a custom read function because the
                ;; libpng default uses file pointers, but here we are
                ;; with a file port in Scheme land.
                (set-read-function! read-struct (make-png-reader port))
                (read-png read-struct info-struct)
                (let* ((width (image-width read-struct info-struct))
                       (height (image-height read-struct info-struct))
                       (color-type (color-type read-struct info-struct))
                       (bit-depth (bit-depth read-struct info-struct))
                       (pitch (bytes-per-row read-struct info-struct))
                       (rows (rows read-struct info-struct))
                       ;; Our target bitmap buffer, which uses RGBA
                       ;; format with 8 bits per color channel, so 4
                       ;; bytes per pixel.
                       (pixels (make-bytevector (* width height 4)))
                       (unpack (unpacker color-type bit-depth)))
                  (if unpack
                      ;; Unpack each row, converting from the PNG's
                      ;; color type and bit depth to RGBA colors of
                      ;; 8 bit depth.
                      (let y-loop ((y 0))
                        (when (< y height)
                          (let* ((row (dereference-pointer rows y))
                                 (bv (pointer->bytevector row pitch))
                                 (offset (* y width 4)))
                            (let x-loop ((x 0))
                              (when (< x width)
                                (let-values (((r g b a) (unpack bv x)))
                                  (let ((offset (+ (* y width 4) (* x 4))))
                                    (bytevector-u8-set! pixels offset r)
                                    (bytevector-u8-set! pixels (+ offset 1) g)
                                    (bytevector-u8-set! pixels (+ offset 2) b)
                                    (bytevector-u8-set! pixels (+ offset 3) a)))
                                (x-loop (+ x 1)))))
                          (y-loop (+ y 1))))
                      ;; Fastest case: The rows are already in RGBA
                      ;; format with 8 bit depth, so we can just
                      ;; copy each row verbatim.
                      (let y-loop ((y 0))
                        (when (< y height)
                          (let* ((row (dereference-pointer rows y))
                                 (bv (pointer->bytevector row pitch))
                                 (offset (* y width 4)))
                            (bytevector-copy! bv 0 pixels offset pitch))
                          (y-loop (+ y 1)))))
                  (values pixels
                          width
                          height)))
              (lambda ()
                (destroy-read-struct read-struct info-struct))))))
      (raise-exception
       (make-exception (make-external-error)
                       (make-exception-with-message
                        (string-append "file not found: " file-name))))))