;;; Sly
;;; Copyright (C) 2014 David Thompson <dthompson2@worcester.edu>
;;;
;;; This program 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.
;;;
;;; This program 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/>.

(define-module (sly render shader)
  #:use-module (ice-9 rdelim)
  #:use-module (ice-9 match)
  #:use-module (system foreign)
  #:use-module (rnrs bytevectors)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-4)
  #:use-module (srfi srfi-9)
  #:use-module (gl)
  #:use-module (gl low-level)
  #:use-module (sly utils)
  #:use-module (sly guardian)
  #:use-module (sly math transform)
  #:use-module (sly math vector)
  #:use-module (sly render color)
  #:use-module (sly config)
  #:use-module (sly wrappers gl)
  #:export (make-shader
            shader?
            load-shader
            load-default-shader
            vertex-position-location
            vertex-texture-location
            shader-uniform-location
            shader-id
            shader-uniforms
            shader-linked?
            null-shader
            apply-shader
            %uniform-setters
            register-uniform-setter!
            uniform?
            uniform-name
            uniform-gl-name
            uniform-default
            uniform-set!
            uniforms))

(define-syntax-rule (define-logger name length-proc log-proc)
  (define (name obj)
    (let ((log-length (u32vector 0)))
      (length-proc obj (version-2-0 info-log-length)
                   (bytevector->pointer log-length))
      (let ((log (make-u8vector (1+ (u32vector-ref log-length 0)))))
        (log-proc obj (u32vector-ref log-length 0) %null-pointer
                  (bytevector->pointer log))
        (format #t "~a\n" (utf8->string log))))))

(define-syntax-rule (define-status name status-proc status-name)
  (define (name obj)
    (let ((status (u32vector 0)))
      (status-proc obj (version-2-0 status-name)
                   (bytevector->pointer status))
      (= (u32vector-ref status 0) 1))))

;;;
;;; Shaders
;;;

(define-record-type <shader-stage>
  (%make-shader-stage type id)
  shader-stage?
  (type shader-stage-type)
  (id shader-stage-id))

(define (vertex-shader-stage? stage)
  "Return #t if STAGE is a vertex shader, #f otherwise."
  (eq? (shader-stage-type stage) 'vertex))

(define (fragment-shader-stage? stage)
  "Return #t if STAGE is a fragment shader, #f otherwise."
  (eq? (shader-stage-type stage) 'fragment))

;; Reap GL shaders when their wrapper objects are GC'd.
(define (free-shader-stage stage)
  (false-if-exception
   (glDeleteShader (shader-stage-id stage))))

(register-finalizer shader-stage? free-shader-stage)

(define-status %shader-stage-compiled? glGetShaderiv compile-status)

(define (shader-stage-compiled? stage)
  "Return #t if STAGE has been compiled already."
  (%shader-stage-compiled? (shader-stage-id stage)))

(define-logger %display-compilation-error glGetShaderiv glGetShaderInfoLog)

(define (display-compilation-error stage)
  (%display-compilation-error (shader-stage-id stage)))

(define (compile-shader-stage stage)
  "Attempt to compile STAGE.  Compilation errors are written to
stdout."
  (glCompileShader (shader-stage-id stage))
  (unless (shader-stage-compiled? stage)
    (display "Failed to compile shader:\n")
    (display-compilation-error stage)))

(define (set-shader-stage-source stage source)
  "Use the GLSL source code in the string SOURCE for the shader
STAGE."
  (let ((length (u32vector (string-length source)))
        (str (u64vector (pointer-address (string->pointer source)))))
    (glShaderSource (shader-stage-id stage) 1 (bytevector->pointer str)
                    (bytevector->pointer length))))

(define (gl-shader-type type)
  "Convert the symbol TYPE to the appropriate OpenGL shader constant."
  (cond ((eq? type 'vertex)
         (version-2-0 vertex-shader))
        ((eq? type 'fragment)
         (version-2-0 fragment-shader))
        (else
         (error "Invalid shader type: " type))))

(define (make-shader-stage type source)
  "Create a new GLSL shader stage of the given TYPE (either 'vertex'
or 'fragment') and compile the GLSL program contained in the string
SOURCE."
  (let* ((id (glCreateShader (gl-shader-type type)))
         (stage (%make-shader-stage type id)))
    (set-shader-stage-source stage source)
    (compile-shader-stage stage)
    (guard stage)))

(define (make-vertex-shader-stage source)
  "Create a new GLSL vertex shader stage and compile the GLSL program
contained in the string SOURCE."
  (make-shader-stage 'vertex source))

(define (make-fragment-shader-stage source)
  "Create a new GLSL fragment shader stage and compile the GLSL
program contained in the string SOURCE."
  (make-shader-stage 'fragment source))

(define (load-shader-stage type file-name)
  "Create a new GLSL shader stage of the given TYPE (either 'vertex'
or 'fragment') and compile the GLSL program stored in the file
FILE-NAME."
  (if (file-exists? file-name)
      (make-shader-stage type (call-with-input-file file-name read-string))
      (error "File not found!" file-name)))

(define (load-vertex-shader file-name)
  "Create a new GLSL vertex shader stage and compile the GLSL program
stored in the file FILE-NAME."
  (load-shader-stage 'vertex file-name))

(define (load-fragment-shader file-name)
  "Create a new GLSL vertex shader stage and compile the GLSL program
stored in the file FILE-NAME."
  (load-shader-stage 'fragment file-name))

;;;
;;; Shader Programs
;;;

(define-record-type <uniform>
  (make-uniform name gl-name location default)
  uniform?
  (name uniform-name)
  (gl-name uniform-gl-name)
  (location uniform-location)
  (default uniform-default))

(define-record-type <attribute>
  (make-attribute name location)
  attribute?
  (name attribute-name)
  (location attribute-location))

(define-record-type <shader>
  (%make-shader id uniforms)
  shader?
  (id shader-id)
  (uniforms shader-uniforms))

;; Hard-coded vertex attribute locations.  These are fixed so that all
;; Sly shaders abide by the same interface.
(define vertex-position-location 0)
(define vertex-texture-location 1)

(define (shader-uniform-location shader name)
  (let ((uniform (find (lambda (uniform)
                         (eq? (uniform-name uniform) name))
                       (shader-uniforms shader))))
    (if uniform
        (uniform-location uniform)
        (error "Uniform not found: " name))))

(define (free-shader shader)
  (false-if-exception
   (glDeleteProgram (shader-id shader))))

(register-finalizer shader? free-shader)

(define-status shader-linked? glGetProgramiv link-status)
(define-logger display-linking-error glGetProgramiv glGetProgramInfoLog)

(define (make-shader vertex-stage fragment-stage uniforms)
  "Create a new shader program that has been linked with the given
VERTEX-STAGE and FRAGMENT-STAGE."
  (unless (and (vertex-shader-stage? vertex-stage)
               (fragment-shader-stage? fragment-stage))
    (error "Expected vertex and fragment shader stages"
           vertex-stage fragment-stage))
  (let ((id (glCreateProgram))
        (stages (list vertex-stage fragment-stage)))
    (define build-uniform
      (match-lambda
        ((name gl-name default)
         (let ((location (glGetUniformLocation id gl-name)))
           (if (= location -1)
               (error "Uniform not found: " gl-name)
               (make-uniform name gl-name location default))))))

    (define (string->attribute attribute-name)
      (let ((location (glGetAttribLocation id attribute-name)))
        (if (= location -1)
            (error "Attribute not found: " attribute-name)
            (make-attribute attribute-name location))))

    (catch #t
      (lambda ()
        (for-each (lambda (stage)
                    (glAttachShader id (shader-stage-id stage)))
                  stages)
        ;; Bind attribute locations
        (glBindAttribLocation id vertex-position-location "position")
        (glBindAttribLocation id vertex-texture-location "tex")
        (glLinkProgram id)
        (unless (shader-linked? id)
          (display "Failed to link shader program:\n")
          (display-linking-error id))
        ;; Once the program has been linked, the shaders can be detached.
        (for-each (lambda (stage)
                    (glDetachShader id (shader-stage-id stage)))
                  stages)
        (let* ((uniforms (map build-uniform uniforms)))
          (guard (%make-shader id uniforms))))
      throw
      (lambda _
        ;; Make sure to delete program in case linking fails.
        (glDeleteProgram id)))))

(define* (load-shader #:key vertex-source fragment-source
                      uniforms)
  (make-shader (load-vertex-shader vertex-source)
               (load-fragment-shader fragment-source)
               uniforms))

(define null-shader
  (%make-shader 0 '()))

(define (apply-shader shader)
  (glUseProgram (shader-id shader)))

;;;
;;; Uniforms
;;;

(define-record-type <uniform-setter>
  (make-uniform-setter predicate proc)
  uniform-setter?
  (predicate uniform-setter-predicate)
  (proc uniform-setter-proc))

(define %uniform-setters '())

(define (register-uniform-setter! predicate setter)
  "Add a new type of uniform setter for shader programs where
PREDICATE tests the type of a given value and SETTER performs the
necessary OpenGL calls to set the uniform value in the proper
location."
  (set! %uniform-setters
        (cons (make-uniform-setter predicate setter)
              %uniform-setters)))

;; Built-in uniform setters for booleans, numbers, vectors, and
;; colors.
(register-uniform-setter! boolean?
                          (lambda (location b)
                            (glUniform1i location (if b 1 0))))

(register-uniform-setter! number?
                          (lambda (location n)
                            (glUniform1f location n)))

(register-uniform-setter! vector2?
                          (lambda (location v)
                            (glUniform2f location (vx v) (vy v))))

(register-uniform-setter! vector3?
                          (lambda (location v)
                            (glUniform3f location (vx v) (vy v) (vz v))))

(register-uniform-setter! vector4?
                          (lambda (location v)
                            (glUniform4f location (vx v) (vy v) (vz v) (vw v))))

(register-uniform-setter! transform?
                          (lambda (location t)
                            (let ((pointer (transform->pointer t)))
                              (glUniformMatrix4fv location 1 #f
                                                  pointer))))

(register-uniform-setter! color?
                          (lambda (location c)
                            (glUniform4f location
                                         (color-r c)
                                         (color-g c)
                                         (color-b c)
                                         (color-a c))))

(define (uniform-set! shader name value)
  "Use the appropriate setter procedure to translate VALUE into OpenGL
compatible data and assign it to the location of the uniform NAME
within SHADER."
  (let ((setter (find (lambda (setter)
                        ((uniform-setter-predicate setter) value))
                      %uniform-setters))
        (location (shader-uniform-location shader name)))
    (if setter
        ((uniform-setter-proc setter) location value)
        (error "Not a valid uniform data type" value))))

;;;
;;; Built-in Shaders
;;;

;; Use lazy evaluation because shader loading will fail until there is
;; an OpenGL context available.
(define %default-shader
  (delay
    (load-shader
     #:vertex-source (scope-datadir "/shaders/default-vertex.glsl")
     #:fragment-source (scope-datadir "/shaders/default-fragment.glsl")
     #:uniforms `((mvp "mvp" ,identity-transform)
                  (color "color" ,white)
                  (texture? "use_texture" #f)))))

(define (load-default-shader)
  "Load and return the default shader program."
  (force %default-shader))