;;; Guile DRMAA --- Guile bindings for DRMAA
;;; Copyright © 2021 Ricardo Wurmus <rekado@elephly.net>
;;;
;;; This file is part of Guile DRMAA.
;;;
;;; Guile DRMAA 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.
;;;
;;; Guile DRMAA 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 Guile DRMAA.  If not, see <http://www.gnu.org/licenses/>.

(define-module (drmaa v1 high)
  #:use-module (drmaa errors)
  #:use-module ((drmaa v1 low) #:prefix low:)
  #:use-module (oop goops)
  #:use-module (srfi srfi-1)
  #:use-module (srfi srfi-19)
  #:use-module (srfi srfi-34)
  #:use-module (srfi srfi-35)
  #:use-module (ice-9 match)
  #:use-module (ice-9 format)
  #:export (job-template

            with-drmaa-session

            run-job
            run-bulk-jobs)
  #:re-export ((low:control . control)
               (low:job-ps . job-ps)
               (low:synchronize . synchronize)
               (low:wait . wait)

               (low:w-if-exited? . w-if-exited?)
               (low:w-exit-status . w-exit-status)
               (low:w-if-signaled? . w-if-signaled?)
               (low:w-termsig . w-termsig)
               (low:w-coredump? . w-coredump?)
               (low:w-if-aborted? . w-if-aborted?)

               (low:get-contact . get-contact)
               (low:drmaa-version . drmaa-version)
               (low:get-drm-system . get-drm-system)
               (low:get-drmaa-implementation . get-drmaa-implementation)))

(define-class <drmaa-class> (<object>))

;; TODO: maybe move validation to setters instead of only running them
;; at initialization time.
(define-method (initialize (instance <drmaa-class>) initargs)
  (define klass (class-of instance))
  ;; Validate slots
  (for-each
   (lambda (slot)
     (let* ((options (slot-definition-options slot))
            (keyword (slot-definition-init-keyword slot)))
       ;; Ensure required slots are filled.
       (and (memq #:required? options)
            (or (member keyword initargs)
                (raise (condition
                        (&formatted-message
                         (format "~a: required field `~a' missing.~%")
                         (arguments (list (class-name klass)
                                          (slot-definition-name slot))))))))

       ;; Only perform these checks if a value is provided.
       (and=> (member keyword initargs)
              (lambda (tail)
                ;; Run validators on slot values
                (match (memq #:valid? options)
                  ((_ validate . rest)
                   (match tail
                     ((_ value . rest)
                      ;; TODO: allow for better error messages
                      (or (validate value)
                          (raise (condition
                                  (&formatted-message
                                   (format "~a: field `~a' has the wrong type.~%")
                                   (arguments (list (class-name klass)
                                                    (slot-definition-name slot))))))))))
                  (_ #t))))))
   (class-slots klass))

  ;; Reject extraneous fields
  (let* ((allowed (map slot-definition-init-keyword (class-slots klass)))
         (provided (filter keyword? initargs)))
    (match (lset-difference eq? provided allowed)
      (() #t)
      (extraneous
       (raise (condition
               (&formatted-message
                (format "~a: extraneous fields: ~{~a ~}~%")
                (arguments (list (class-name klass)
                                 (map keyword->symbol extraneous)))))))))
  (next-method))

(define-class <job-template> (<drmaa-class>)
  ;; Slots
  (job-name
   #:accessor job-template-job-name
   #:init-keyword #:job-name
   #:init-form (format #false "guile-drmaa-job-~a"
                       (time-second (current-time time-monotonic)))
   #:drmaa-name (low:DRMAA 'JOB_NAME))

  (remote-command
   #:accessor job-template-remote-command
   #:init-keyword #:remote-command
   #:init-value "/bin/true"
   #:drmaa-name (low:DRMAA 'REMOTE_COMMAND))
  (arguments
   #:accessor job-template-arguments
   #:init-keyword #:arguments
   #:init-value (list)
   #:drmaa-name (low:DRMAA 'V_ARGV))

  (working-directory
   #:accessor job-template-working-directory
   #:init-keyword #:working-directory
   #:init-value ""
   #:drmaa-name (low:DRMAA 'WD))

  (hold?
   #:accessor job-template-hold?
   #:init-keyword #:hold?
   #:init-value #true
   #:drmaa-name (low:DRMAA 'JS_STATE)
   #:valid? boolean?
   #:transform (lambda (value)
                 (if value
                     (low:DRMAA 'SUBMISSION_STATE_HOLD)
                     (low:DRMAA 'SUBMISSION_STATE_ACTIVE))))
  (join-files?
   #:accessor job-template-join-files?
   #:init-keyword #:join-files?
   #:valid? boolean?
   #:transform (lambda (value)
                 (if value "y" "n"))
   #:drmaa-name (low:DRMAA 'JOIN_FILES))
  (block-email?
   #:accessor job-template-block-email?
   #:init-keyword #:block-email?
   #:valid? boolean?
   #:transform (lambda (value)
                 (if value "0" "1"))
   #:drmaa-name (low:DRMAA 'BLOCK_EMAIL))

  (input-path
   #:accessor job-template-input-path
   #:init-keyword #:input-path
   #:drmaa-name (low:DRMAA 'INPUT_PATH))
  (output-path
   #:accessor job-template-output-path
   #:init-keyword #:output-path
   #:drmaa-name (low:DRMAA 'OUTPUT_PATH))
  (error-path
   #:accessor job-template-error-path
   #:init-keyword #:error-path
   #:drmaa-name (low:DRMAA 'ERROR_PATH))
  (transfer-files
   #:accessor job-template-transfer-files
   #:init-keyword #:transfer-files
   #:valid? (lambda (value)
              (and (list? value)
                   (for-each
                    (lambda (item)
                      (or (member item '(error input output))
                          (raise (condition
                                  (&message
                                   (message "transfer-files: must be a list containing the symbols error, input, or output."))))))
                    value)))
   #:transform (lambda (value)
                 (string-join (map (lambda (item)
                                     (match item
                                       ('error "e")
                                       ('input "i")
                                       ('output "o")))
                                   value)))
   #:drmaa-name (low:DRMAA 'TRANSFER_FILES))

  (job-category
   #:accessor job-template-job-category
   #:init-keyword #:job-category
   #:drmaa-name (low:DRMAA 'JOB_CATEGORY))
  (native-specification
   #:accessor job-template-native-specification
   #:init-keyword #:native-specification
   #:drmaa-name (low:DRMAA 'NATIVE_SPECIFICATION))

  (environment
   #:accessor job-template-environment
   #:init-keyword #:environment
   #:valid? (lambda (value)
              (or (and (list? value)
                       (for-each
                        (match-lambda
                          ((key . value) #true)
                          (_ #false))
                        value))
                  (raise (condition
                          (&message
                           (message "environment: must be a list of pairs."))))))
   #:transform (lambda (value)
                 (map (match-lambda
                        ((key . value)
                         (format #false "~a=~a" key value)))
                      value))
   #:drmaa-name (low:DRMAA 'V_ENV))
  (emails
   #:accessor job-template-emails
   #:init-keyword #:emails
   #:valid? (lambda (value)
              (or (and (list? value)
                       (every string? value))
                  (raise (condition
                          (&message
                           (message "emails: must be a list of strings."))))))
   #:drmaa-name (low:DRMAA 'V_EMAIL))

  (start-time
   #:accessor job-template-start-time
   #:init-keyword #:start-time
   #:valid? date?
   #:transform (lambda (value)
                 (date->string value "~Y/~m/~d ~H:~M:~S"))
   #:drmaa-name (low:DRMAA 'START_TIME))
  (deadline-time
   #:accessor job-template-deadline-time
   #:init-keyword #:deadline-time
   #:valid? date?
   #:transform (lambda (value)
                 (date->string value "~Y/~m/~d ~H:~M:~S"))
   #:drmaa-name (low:DRMAA 'DEADLINE_TIME))
  (duration-hlimit
   #:accessor job-template-duration-hlimit
   #:init-keyword #:duration-hlimit
   #:valid? (lambda (value)
              (and (time? value)
                   (eq? time-duration (time-type value))))
   #:transform (lambda (value)
                 (format #false "00:00:~2,'0d" (time-second value)))
   #:drmaa-name (low:DRMAA 'DURATION_HLIMIT))
  (duration-slimit
   #:accessor job-template-duration-slimit
   #:init-keyword #:duration-slimit
   #:valid? (lambda (value)
              (and (time? value)
                   (eq? time-duration (time-type value))))
   #:transform (lambda (value)
                 (format #false "00:00:~2,'0d" (time-second value)))
   #:drmaa-name (low:DRMAA 'DURATION_SLIMIT))
  (wct-hlimit
   #:accessor job-template-wct-hlimit
   #:init-keyword #:wct-hlimit
   #:valid? (lambda (value)
              (and (time? value)
                   (eq? time-duration (time-type value))))
   #:transform (lambda (value)
                 (format #false "00:00:~2,'0d" (time-second value)))
   #:drmaa-name (low:DRMAA 'WCT_HLIMIT))
  (wct-slimit
   #:accessor job-template-wct-slimit
   #:init-keyword #:wct-slimit
   #:valid? (lambda (value)
              (and (time? value)
                   (eq? time-duration (time-type value))))
   #:transform (lambda (value)
                 (format #false "00:00:~2,'0d" (time-second value)))
   #:drmaa-name (low:DRMAA 'WCT_SLIMIT))

  ;; Class options
  #:name "job-template")

(define (job-template . args)
  (apply make <job-template> args))

(define implementation
  (let ((result #false))
    (lambda ()
      "Return the DRMAA implementation.  This is a memoized procedure."
      (or result
          (begin
            (set! result (low:get-drmaa-implementation))
            result)))))

(define (grid-engine?)
  (or (string-prefix? "UGE " (implementation))   ; Univa Grid Engine
      (string-prefix? "SGE " (implementation))))

(define (add-to-native-specification! t spec)
  (let ((specs (low:get-attribute t (low:DRMAA 'NATIVE_SPECIFICATION))))
    (low:set-attribute! t (low:DRMAA 'NATIVE_SPECIFICATION)
                        (string-append specs " " spec))))

(define-method (lower-job-template (instance <job-template>))
  (define klass (class-of instance))
  (define implemented-attributes
    (let ((result #false))
      (lambda ()
        (or result
            (begin
              (set! result (low:get-attribute-names))
              result)))))
  (define t #false)
  (catch #true
    (lambda ()
      (set! t (low:allocate-job-template!))
      (for-each
       (lambda (slot)
         (let* ((options   (slot-definition-options slot))
                (keyword   (slot-definition-init-keyword slot))
                (drmaa     (and=> (memq #:drmaa-name options) cadr))
                (transform (or (and=> (memq #:transform options) cadr)
                               identity))
                (slot-name (slot-definition-name slot)))
           (when (slot-bound? instance slot-name)
             (let* ((value (slot-ref instance slot-name))
                    (transformed (transform value)))
               (cond
                ;; Empty lists and #false are not valid values to set.
                ((or (null? transformed) (not transformed))
                 #true)
                ((pair? transformed)
                 (apply low:set-vector-attribute! t drmaa transformed))
                (else
                 (if (member drmaa (implemented-attributes))
                     (low:set-attribute! t drmaa transformed)
                     ;; Use native specification if possible 
                     (cond
                      ((and (grid-engine?)
                            (eq? drmaa (low:DRMAA 'DURATION_HLIMIT)))
                       (add-to-native-specification! t
                                                     (format #false "-h_rt ~a" transformed)))
                      ((and (grid-engine?)
                            (eq? drmaa (low:DRMAA 'DURATION_SLIMIT)))
                       (add-to-native-specification! t
                                                     (format #false "-s_rt ~a" transformed)))
                      ((and (grid-engine?)
                            (eq? drmaa (low:DRMAA 'WCT_HLIMIT)))
                       (add-to-native-specification! t
                                                     (format #false "-h_rt ~a" transformed)))
                      ((and (grid-engine?)
                            (eq? drmaa (low:DRMAA 'WCT_SLIMIT)))
                       (add-to-native-specification! t
                                                     (format #false "-s_rt ~a" transformed)))
                      (else
                       (raise (condition
                               (&formatted-message
                                (format "This DRMAA implementation does not support the attribute `~a'.~%")
                                (arguments (list drmaa))))))))))))))
       (class-slots klass))
      t)
    (lambda (key . args)
      (low:delete-job-template! t)
      (apply throw key args))))

(define-syntax-rule (with-drmaa-session body ...)
  (dynamic-wind
    (lambda () (low:init-session!))
    (lambda () body ...)
    (lambda () (low:exit-session!))))

(define-syntax-rule (with-job-template instance proc)
  (let ((template #false))
    (dynamic-wind
      (lambda ()
        (set! template (lower-job-template instance)))
      (lambda ()
        (proc template))
      (lambda ()
        (when template
          (low:delete-job-template! template))))))

(define-method (run-job (template <job-template>))
  "Submit a single job with the attributes defined in the job
template.  Upon success return the job identifier."
  (with-job-template template low:run-job))

(define* (run-bulk-job template end
                       #:key
                       (start 1) (increment 1))
  "Submit a set of parametric jobs, which can be run concurrently.
The attributes defined in TEMPLATE are used for every job in the set.
Each job is identical except for its index.  The first job has an
index equal to START; the next has an index equal to the sum of START
and INCREMENT.  END is the maximum value of the index, but the value
of the last job's index may not be equal to END if the difference
between START and END is not evenly divisible by INCREMENT.

Return a list of job identifiers."
  (with-job-template template
                     (lambda (t)
                       (low:run-bulk-jobs t end #:start start #:increment increment))))