;;;; wh-engine/render/shader.lisp
;;;; Lisp class for holding & handling shaders.
;; (in-package wh-engine/render)

(defun replace-xref-vars (forms vars replace-fun)
  (loop for elt on forms
        for form = (car elt)
        if (symbolp form)
          do (let ((var-entry (find form vars :key #'car)))
               (when var-entry
                 (rplaca elt (funcall replace-fun var-entry))))
        else if (consp form)
               do (replace-xref-vars (cdr form) vars replace-fun)))

(defun collect-shader-vars (forms)
  (loop for elt on forms
        for form = (car elt)
        if (symbolp form)
          collect form
        else if (consp form)
               nconc (collect-shader-vars (cdr form))))

(defvar *vert-builtins* '((*vertex-id* . "gl_VertexID")
                          (*instance-id* . "gl_InstanceID")
                          (*draw-id* . "gl_DrawID")
                          (*base-vertex* . "gl_BaseVertex")
                          (*base-instance* . "gl_BaseInstance")
                          ;; out
                          (*position* . "gl_Position")
                          (*point-size* . "gl_PointSize")
                          (*clip-distance* . "gl_ClipDistance")))

(defvar *frag-builtins* '((*frag-coord* . "gl_FragCoord")
                          (*front-facing* . "gl_FrontFacing")
                          (*point-coord* . "gl_PointCoord")
                          (*sample-id* . "gl_SampleID")
                          (*sample-position* . "gl_SamplePosition")
                          (*sample-mask-in* . "gl_SampleMaskIn")
                          (*clip-distance* . "gl_ClipDistance")
                          (*primitive-id* . "gl_PrimitiveID")
                          (*layer* . "gl_Layer")
                          (*viewport-index* . "gl_ViewportIndex")
                          ;; out
                          (*frag-depth* . "gl_FragDepth")
                               (*sample-mask* . "gl_SampleMask")))

(defvar *infix-ops* '((+ "+" t)
                           (- "-" t)
                           (* "*" t)
                           (/ "/" t)
                           (mod "%" t)
                           (= "==" t)
                           (eq "==" nil)
                           (eql "==" nil)
                           (equal "==" nil)
                           (/= "!=" nil)
                           (< "<" t)
                           (> ">" t)
                           (<= "<=" t)
                           (>= ">=" t)
                           (<< "<<" nil)
                           (>> ">>" nil)
                           (ash "<<" nil)
                           (logand "&" t)
                           (logior "|" t)
                           (logxor "^" t)
                           (and "&&" t)
                      (or "||" t)))

(defvar *types* '((:void . "void")
                  (:boolean . "bool")
                  (:integer . "int")
                  (:unsigned-integer . "uint")
                  (:float . "float")
                  (:double . "double")
                  (:vec2 . "vec2")
                  (:vec3 . "vec3")
                  (:vec4 . "vec4")
                  (:double-vec2 . "dvec2")
                  (:double-vec3 . "dvec3")
                  (:double-vec4 . "dvec4")
                  (:boolean-vec2 . "bvec2")
                  (:boolean-vec3 . "bvec3")
                  (:boolean-vec4 . "bvec4")
                  (:integer-vec2 . "ivec2")
                  (:integer-vec3 . "ivec3")
                  (:integer-vec4 . "ivec4")
                  (:unsigned-vec2 . "uvec2")
                  (:unsigned-vec3 . "uvec3")
                  (:unsigned-vec4 . "uvec4")
                  (:mat2 . "mat2")
                  (:mat3 . "mat3")
                  (:mat4 . "mat4")
                  (:sampler-1d . "sampler1D")
                  (:depth-sampler-1d . "sampler1DShadow")
                  (:array-sampler-1d . "sampler1DArray")
                  (:depth-array-sampler-1d . "sampler1DArrayShadow")
                  (:sampler-2d . "sampler2D")
                  (:depth-sampler-2d . "sampler2DShadow")
                  (:array-sampler-2d . "sampler2DArray")
                  (:depth-array-sampler-2d . "sampler2DArrayShadow")
                  (:multisampler-2d . "sampler2DMS")
                  (:array-multisampler-2d . "sampler2DMSArray")
                  (:rect-sampler-2d . "sampler2DRect")
                  (:depth-rect-sampler-2d . "sampler2DRectShadow")
                  (:sampler-3d . "sampler3D")
                  (:cube-sampler . "samplerCube")
                  (:depth-cube-sampler . "samplerCubeShadow")
                  (:array-cube-sampler . "samplerCubeArray")
                  (:depth-array-cube-sampler . "samplerCubeArrayShadow")
                  (:buffer-sampler . "samplerBuffer")
                  ))

(defun swizzlep (sym)
  (let ((sym-str (string-downcase (symbol-name sym))))
    (and (char= (char sym-str 0) #\v)
         (loop for c across (subseq sym-str 1)
               always (member c '(#\x #\y #\z #\w) :test #'char=))
         (subseq sym-str 1))))

(defun conv-shader-form (form funs vars)
  (flet ((conv-form (x) (conv-shader-form x funs vars)))
    (etypecase form
      (cons
       (let* ((sym (car form))
              (args (cdr form))
              (infix-op (assoc sym *infix-ops*))
              (swizzle (swizzlep sym)))
         (cond
           (infix-op
            ;; special case for -, which is also a unary op
            (if (and (eq sym '-) (= (length args) 1))
                (format nil "(-~A)"
                        (conv-form (first args)))
                (progn
                  (unless (or (= (length args) 2) (and (nth 2 infix-op) (> (length args) 2)))
                    (error "Invalid # of args for shader infix-op: ~S" infix-op))
                  ;; don't wanna repeat this a bajillion times, so just splice the
                  ;; GLSL version of the operator into the format control string
                  (format nil (concatenate 'string "(~{~A~^ " (nth 1 infix-op) " ~})")
                          (loop for x in args collect (conv-form x))))))
           ;; the unary operators ! and ~
           ((member sym '(not lognot))
            (unless (= (length args) 1)
              (error "Invalid # of args for shader unary-op: ~S" sym))
            (format nil (if (eq sym 'not) "(!~A)" "(~~~A)")
                    (conv-form (first args))))
           ;; hijack the (exprA, exprB, ...) to allow multi-setf
           ((eq sym 'setf)
            (unless (= (mod (length args) 2) 0)
              (error "Invalid # of args for shader setf"))
            (format nil "(~{~A = ~A~^, ~})"
                    (loop for x in args collect (conv-form x))))
           ((eq sym 'nth)
            (unless (= (length args) 2)
              (error "Invalid # of args for shader nth"))
            (format nil "~A[~A]"
                    (conv-form (second args))
                    (conv-form (first args))))
           ((eq sym 'aref)
            (unless (>= (length args) 2)
              (error "Invalid # of args for shader aref"))
            (format nil "~A~{[~A]~}"
                    (conv-form (first args))
                    (loop for x in (rest args) collect (conv-form x))))
           ;; non-statement (if) is the ternary operator
           ((eq sym 'if)
            (unless (= (length args) 3)
              (error "Invalid # of args for shader if"))
            (format nil "(~A ? ~A : ~A)"
                    (conv-form (first args))
                    (conv-form (second args))
                    (conv-form (third args))))
           ;; WHY IS THIS A FUNCTION HISS
           ((eq sym 'length)
            (unless (= (length args) 1)
              (error "Invalid # of args for shader length"))
            (format nil "~A.length()"
                    (conv-form (first args))))
           ;; apparently you can do this, (exprA, exprB) is non-statement progn
           ;; thanks C
           ((eq sym 'progn)
            (unless (>= (length args) 1)
              (error "Invalid # of args for shader progn"))
            (format nil "(~{~A~^, ~})"
                    (loop for x in args collect (conv-form x))))
           (swizzle
            (unless (= (length args) 1)
              (error "Invalid # of args for shader swizzle"))
            (format nil "~A.~A"
                    (conv-form (first args))
                    swizzle))
           ;; need this for (for)
           ((eq sym 'declare-var)
            (unless (and (<= 1 (length args) 2) (= (length (first args)) 2))
              (error "Invalid # of args for shader declare-var"))
            (format nil "~A ~A~:[~; = ~A~]"
                    (or (cdr (assoc (cadar args) *types*)) (error "Invalid type specifier: ~S" (second args)))
                    (conv-form (caar args))
                    (= (length args) 2)
                    (conv-form (second args))))
           (t
            (format nil "~A(~{~A~^, ~})"
                    (or (cdr (assoc sym funs)) (error "Invalid shader fun: ~S" sym))
                    (loop for x in args collect (conv-form x)))))))
      (boolean (if form "true" "false"))
      (symbol (or (cdr (assoc form vars)) (error "Invalid shader var: ~S" form)))
      (integer (format nil "~D" form))
      (float (format nil "~F" form)))))

(defun conv-shader-stmt (stmt funs vars returnp)
  (labels ((conv-form (x) (conv-shader-form x funs vars))
           (conv-stmt (x r) (conv-shader-stmt x funs vars r))
           (loop-conv-stmts (list r)
             (loop for x on list
                   collect (conv-stmt (car x) (and r (not (cdr x)))))))
    (if (consp stmt)
        (let* ((sym (car stmt))
               (args (cdr stmt)))
          (case sym
            ((progn)
             (format nil "{~%~{~A~}}~%"
                     (loop-conv-stmts args returnp)))
            ((if)
             (unless (<= 2 (length args) 3)
               (error "Invalid # of args for shader if"))
             (format nil "if (~A) ~{~%~A~^else~}"
                     (conv-form (first args))
                     (loop-conv-stmts (rest args) returnp)))
            ;; handle both forms in one case clause
            ((when unless)
             (unless (>= (length args) 2)
               (error "Invalid # of args for shader when/unless"))
             (format nil "if (~:[~;!(~]~A~2:*~:[~;)~]~1*) {~%~{~A~}}~%"
                     (eq sym 'unless)
                     (conv-form (first args))
                     (loop-conv-stmts (rest args) returnp)))
            ((cond)
             (unless (>= (length args) 1)
               (error "Invalid # of args for shader cond"))
             (format nil "~{~:[if (~A) ~;~1*~]{~%~{~A~}}~^ else ~}~%"
                     (loop for clause in args
                           nconc (list (eq (first clause) t)
                                       (conv-form (first clause))
                                       (loop-conv-stmts (rest clause) returnp)))))
            ((case)
             (unless (>= (length args) 2)
               (error "Invalid # of args for shader case"))
             (format nil "switch (~A) {~%~{~:[~{case ~A:~%~}~;default:~%~1*~]~{~A~}break;~%~}}~%"
                     (conv-form (first args))
                     (loop for clause in (rest args)
                           nconc (list (eq (car clause) t)
                                       (if (listp (car clause))
                                           (loop for x in (car clause) collect (conv-form x))
                                           (list (conv-form (car clause))))
                                       (loop-conv-stmts (cdr clause) returnp)))))
            ((return)
             (unless (<= 0 (length args) 1)
               (error "Invalid # of args for shader return"))
             (format nil "return~:[~; ~A~];~%"
                     args
                     (and args (conv-form (first args)))))
            ((break continue discard)
             (when args
               (error "Invalid # of args for shader break"))
             (format nil "~(~A~);~%" sym))
            ((while)
             (unless (>= (length args) 1)
               (error "Invalid # of args for shader while"))
             (format nil "while (~A) {~%~{~A~}}~%"
                     (conv-form (first args))
                     (loop-conv-stmts (rest args))))
            ((for)
             (unless (and (>= (length args) 1) (= (length (first args)) 3))
               (error "Invalid # of args for shader for"))
             (format nil "for (~A;~A;~A) {~%~{~A~}}~%"
                     (conv-form (caar args))
                     (conv-form (cadar args))
                     (conv-form (caddar args))
                     (loop-conv-stmts (cdr args) returnp)))
            ((let*)
             (conv-stmt `(progn
                           ,@(loop for x in (first args)
                                   if (consp (car x))
                                     collect `(declare-var ,@x)
                                   else
                                     collect `(declare-var ,x))
                           ,@(rest args))
                        returnp))
            ((setf)
             (unless (= (mod (length args) 2) 0)
               (error "Invalid # of args for shader setf"))
             (format nil "~{~A = ~A;~%~}"
                     (loop for x in args collect (conv-form x))))
            (t
             (format nil "~:[~;return ~]~A;~%" returnp (conv-form stmt)))))
        ;; something else
        (format nil "~:[~;return ~]~A;~%" returnp (conv-form stmt)))
    ))

(defun transpile-shader (name &key
                                version
                                ((:in vert-inputs))
                                ((:inter inter-vars))
                                ((:out frag-outputs))
                                ((:uniform uniform-vars))
                                ((:vert vert-body))
                                ((:frag frag-body)))
  "Translate Lisp shader code to GLSL."
  ;; create implicit inter vars
  (replace-xref-vars frag-body vert-inputs
                     (lambda (var-entry)
                       (let* ((orig-var (car var-entry))
                              (inter-var-entry (find orig-var inter-vars
                                                     :key (lambda (e) (getf (cdr e) :orig-var)))))
                         (if inter-var-entry
                             (car inter-var-entry)
                             (let ((new-var (gensym (symbol-name orig-var))))
                               (push `(,new-var :orig-var ,orig-var :type ,(getf (cdr var-entry) :type))
                                     inter-vars)
                               (push `(setf ,new-var ,orig-var) vert-body)
                               new-var)))))
  `(:in ,vert-inputs :inter ,inter-vars :out ,frag-outputs :uniform ,uniform-vars
    :vert ,vert-body :frag ,frag-body))

(defmacro define-shader (name (&key
                                 ((:in (&rest vert-inputs)) ())
                                 ((:inter (&rest inter-vars)) ())
                                 ((:out (&rest frag-outputs)) ())
                                 ((:uniform (&rest uniform-vars)) ()))
                         &key
                           ((:vert (&body vert-body)))
                           ((:frag (&body frag-body)))
                           version)
  "Define a GLSL shader with Lisp code."
  `(defparameter ,name (transpile-shader ',name
                                         :version ,version
                                         :in ',vert-inputs
                                         :inter ',inter-vars
                                         :out ',frag-outputs
                                         :uniform ',uniform-vars
                                         :vert ',vert-body :frag ',frag-body)))