(define-module (webid-oidc jws)
  #:use-module (webid-oidc jwk)
  #:use-module (webid-oidc errors)
  #:use-module ((webid-oidc stubs) #:prefix stubs:)
  #:use-module (rnrs bytevectors)
  #:use-module (ice-9 receive))

(define-public (the-jws-header x)
  (with-exception-handler
      (lambda (cause)
        (raise-not-a-jws-header x cause))
    (lambda ()
      (let ((alg (assq-ref x 'alg)))
        (unless alg
          (raise-missing-alist-key x 'alg))
        (unless (string? alg)
          (raise-unsupported-alg alg))
        (case (string->symbol alg)
          ((HS256 HS384 HS512 RS256 RS384 RS512 ES256 ES384 ES512 PS256 PS384 PS512)
           x)
          (else
           (raise-unsupported-alg (string->symbol alg))))))))

(define-public (the-jws-payload x)
  (with-exception-handler
      (lambda (cause)
        (raise-not-a-jws-payload x cause))
    (lambda ()
      (unless (list? x)
        (scm-error 'wrong-type-arg "the-jws-payload" "expected a list" '() (list x)))
      x)))

(define-public (the-jws x)
  (with-exception-handler
      (lambda (cause)
        (raise-not-a-jws x cause))
    (lambda ()
      (unless (pair? x)
        (scm-error 'wrong-type-arg "the-jws" "expected a pair" '() (list x)))
      (cons (the-jws-header (car x))
            (the-jws-payload (cdr x))))))

(define-public (jws-header? x)
  (false-if-exception
   (and (the-jws-header x) #t)))

(define-public (jws-payload? x)
  (false-if-exception
   (and (the-jws-payload x) #t)))

(define-public (jws? x)
  (false-if-exception
   (and (the-jws x) #t)))

(define-public (make-jws header payload)
  (the-jws (cons (the-jws-header header)
                 (the-jws-payload payload))))

(define-public (jws-header jws)
  (car (the-jws jws)))

(define-public (jws-payload jws)
  (cdr (the-jws jws)))

(define-public (jws-alg jws)
  (if (jws? jws)
      (jws-alg (jws-header jws))
      (string->symbol (assq-ref (the-jws-header jws) 'alg))))

(define (split-in-3-parts string separator)
  (let ((parts (list->vector (string-split string separator))))
    (unless (eqv? (vector-length parts) 3)
      (raise-not-in-3-parts string separator))
    (values (vector-ref parts 0) (vector-ref parts 1) (vector-ref parts 2))))

(define (base64-decode-json str)
  (with-exception-handler
      (lambda (error)
        (cond
         (((record-predicate &not-base64) error)
          (raise-exception error))
         (((record-predicate &not-json) error)
          (raise-exception error))
         (else
          ;; From utf8->string
          (raise-not-base64 str error))))
    (lambda ()
      (stubs:json-string->scm (utf8->string (stubs:base64-decode str))))))

(define (parse str verify)
  (receive (header payload signature)
      (split-in-3-parts str #\.)
    (let ((base (string-append header "." payload))
          (header (base64-decode-json header))
          (payload (base64-decode-json payload)))
      (let ((ret (make-jws header payload)))
        (verify ret base signature)
        ret))))

(define (verify-any alg keys payload signature)
  (define (aux candidates)
    (if (null? keys)
        (raise-no-matching-key keys alg payload signature)
        (let ((next-ok
               (with-exception-handler
                   (lambda (error)
                     #f)
                 (lambda ()
                   (stubs:verify alg (car candidates) payload signature)
                   #t)
                 #:unwind? #t
                 #:unwind-for-type &invalid-signature)))
          (or next-ok
              (aux (cdr candidates))))))
  (aux keys))

(define-public (jws-decode str lookup-keys)
  (with-exception-handler
      (lambda (error)
        (raise-cannot-decode-jws str error))
    (lambda ()
      (parse str
             (lambda (jws payload signature)
               (let ((keys (lookup-keys jws)))
                 (let ((keys (cond ((jwk? keys) (list keys))
                                   ((jwks? keys) (jwks-keys keys))
                                   (else keys))))
                   (verify-any (jws-alg jws) keys payload signature))))))))

(define-public (jws-encode jws key)
  (with-exception-handler
      (lambda (error)
        (raise-cannot-encode-jws jws key error))
    (lambda ()
      (let ((header (jws-header jws))
            (payload (jws-payload jws)))
        (let ((header (stubs:scm->json-string header))
              (payload (stubs:scm->json-string payload)))
          (let ((header (stubs:base64-encode header))
                (payload (stubs:base64-encode payload)))
            (let ((payload (string-append header "." payload)))
              (let ((signature (stubs:sign (jws-alg jws) key payload)))
                (string-append payload "." signature)))))))))