(module pcm-converter racket/base

  (require ffi/unsafe
           "utils.rkt"
           "../ffmpeg-definitions.rkt")

  (provide pcm-conversion-needed?
           make-pcm-converter
           pcm-converter?
           pcm-converter-input-format
           pcm-converter-output-format
           pcm-converter-convert
           pcm-converter-drain
           pcm-converter-close!)

  (define S32-BYTES 4)

  (define-struct pcm-converter (swr-ctx in-layout out-layout input-format output-format channels in-rate out-rate closed?)
    #:mutable
    #:constructor-name make-raw-pcm-converter)

  (define (hash-ref/default h k default)
    (if (and (hash? h) (hash-has-key? h k)) (hash-ref h k) default))

  (define (copy-hash h)
    (let ((out (make-hash)))
      (when (hash? h)
        (for-each (lambda (k) (hash-set! out k (hash-ref h k))) (hash-keys h)))
      out))

  (define (native-signed-ref bs start bytes)
    (int-bytes->integer bs #t (system-big-endian?) start (+ start bytes)))

  (define (native-signed-set! bs start bytes value)
    (integer->integer-bytes value bytes #t (system-big-endian?) bs start))

  (define (clamp-s32 v)
    (cond [(< v -2147483648) -2147483648]
          [(> v 2147483647) 2147483647]
          [else v]))

  (define (expand-sample-to-s32 sample in-bits)
    (clamp-s32 (if (< in-bits 32) (arithmetic-shift sample (- 32 in-bits)) sample)))

  (define (pcm-bytes->s32-bytes buffer size in-bits)
    (cond [(= in-bits 32) (if (= size (bytes-length buffer)) buffer (subbytes buffer 0 size))]
          [else
           (let* ((in-bytes (quotient in-bits 8))
                  (sample-count (quotient size in-bytes))
                  (out (make-bytes (* sample-count S32-BYTES))))
             (for ([i (in-range sample-count)])
               (let* ((in-off (* i in-bytes))
                      (out-off (* i S32-BYTES))
                      (sample (native-signed-ref buffer in-off in-bytes)))
                 (native-signed-set! out out-off S32-BYTES (expand-sample-to-s32 sample in-bits))))
             out)]))

  (define (bytes->native bs size)
    (let ((p (malloc size 'atomic-interior)))
      (memcpy p 0 bs 0 size)
      p))

  (define (native->bytes p size)
    (let ((bs (make-bytes size)))
      (memcpy bs 0 p 0 size)
      bs))

  (define (make-plane ptr)
    (let ((planes (malloc _pointer 1 'atomic-interior)))
      (ptr-set! planes _pointer 0 ptr)
      planes))

  (define (source-value v source)
    (if (eq? v 'source) source v))

  (define (target-sample-rate settings input-format)
    (source-value
     (hash-ref/default settings 'target-sample-rate
                       (hash-ref/default settings 'sample-rate
                                         (hash-ref input-format 'sample-rate)))
     (hash-ref input-format 'sample-rate)))

  (define (target-channels settings input-format)
    (source-value
     (hash-ref/default settings 'target-channels
                       (hash-ref/default settings 'channels
                                         (hash-ref input-format 'channels)))
     (hash-ref input-format 'channels)))

  (define (target-bits settings input-format)
    (let ((source-bits (let ((bits (hash-ref/default input-format 'bits-per-sample 24)))
                         (if (and (integer? bits) (<= bits 24)) bits 24))))
      (source-value
       (hash-ref/default settings 'target-bits-per-sample
                         (hash-ref/default settings 'bits-per-sample source-bits))
       source-bits)))

  (define (make-output-format input-format settings)
    (let* ((out (copy-hash input-format))
           (in-rate (hash-ref input-format 'sample-rate))
           (out-rate (target-sample-rate settings input-format))
           (total (hash-ref/default input-format 'total-samples #f)))
      (hash-set! out 'sample-rate out-rate)
      (hash-set! out 'channels (target-channels settings input-format))
      (hash-set! out 'bits-per-sample (target-bits settings input-format))
      (hash-set! out 'pcm-bits-per-sample 32)
      (hash-set! out 'type 'interleaved)
      (hash-set! out 'endianness 'native-endian)
      (when (and (integer? total) (>= total 0) (integer? in-rate) (> in-rate 0))
        (hash-set! out 'total-samples (inexact->exact (round (* total (/ out-rate in-rate))))))
      out))

  (define (pcm-conversion-needed? input-format settings)
    (let ((in-rate (hash-ref input-format 'sample-rate))
          (in-channels (hash-ref input-format 'channels))
          (out-rate (target-sample-rate settings input-format))
          (out-channels (target-channels settings input-format)))
      (or (not (= in-rate out-rate))
          (not (= in-channels out-channels)))))

  (define (make-pcm-converter input-format settings)
    (let* ((channels-in (hash-ref input-format 'channels))
           (channels-out (target-channels settings input-format))
           (rate-in (hash-ref input-format 'sample-rate))
           (rate-out (target-sample-rate settings input-format))
           (in-layout (ffmpeg-make-default-channel-layout channels-in))
           (out-layout (ffmpeg-make-default-channel-layout channels-out)))
      (let-values (((ret ctx) (swr_alloc_set_opts2 #f
                                                   out-layout AV_SAMPLE_FMT_S32 rate-out
                                                   in-layout AV_SAMPLE_FMT_S32 rate-in
                                                   0 #f)))
        (when (< ret 0) (error 'make-pcm-converter "swr_alloc_set_opts2 failed: ~a" ret))
        (let ((ret-init (swr_init ctx)))
          (when (< ret-init 0) (error 'make-pcm-converter "swr_init failed: ~a" ret-init))
          (make-raw-pcm-converter ctx in-layout out-layout input-format (make-output-format input-format settings)
                                  channels-out rate-in rate-out #f)))))

  (define (ensure-open! c who)
    (when (or (not (pcm-converter? c)) (pcm-converter-closed? c))
      (error who "PCM converter is closed")))

  (define (convert* c in-bytes in-samples)
    (let* ((channels (pcm-converter-channels c))
           (max-out-samples (swr_get_out_samples (pcm-converter-swr-ctx c) in-samples)))
      (cond [(<= max-out-samples 0) (values #"" 0)]
            [else
             (let* ((out-size (* max-out-samples channels S32-BYTES))
                    (out-ptr (malloc out-size 'atomic-interior))
                    (out-planes (make-plane out-ptr))
                    (in-ptr (and in-bytes (bytes->native in-bytes (bytes-length in-bytes))))
                    (in-planes (and in-ptr (make-plane in-ptr)))
                    (out-samples (swr_convert (pcm-converter-swr-ctx c) out-planes max-out-samples in-planes in-samples)))
               (when (< out-samples 0) (error 'pcm-converter-convert "swr_convert failed: ~a" out-samples))
               (values (native->bytes out-ptr (* out-samples channels S32-BYTES)) out-samples))])))

  (define (pcm-converter-convert c buffer size buf-info)
    (ensure-open! c 'pcm-converter-convert)
    (let* ((in-bits (hash-ref/default buf-info 'bits-per-sample (hash-ref (pcm-converter-input-format c) 'bits-per-sample)))
           (in-bytes (quotient in-bits 8))
           (in-channels (hash-ref (pcm-converter-input-format c) 'channels))
           (in-samples (quotient (quotient size in-bytes) in-channels))
           (s32 (pcm-bytes->s32-bytes buffer size in-bits)))
      (convert* c s32 in-samples)))

  (define (pcm-converter-drain c)
    (ensure-open! c 'pcm-converter-drain)
    (let* ((ctx (pcm-converter-swr-ctx c))
           (delay (swr_get_delay ctx (pcm-converter-out-rate c)))
           (channels (pcm-converter-channels c)))
      (cond [(<= delay 0) (values #"" 0)]
            [else
             (let* ((out-size (* delay channels S32-BYTES))
                    (out-ptr (malloc out-size 'atomic-interior))
                    (out-planes (make-plane out-ptr))
                    (out-samples (swr_convert ctx out-planes delay #f 0)))
               (when (< out-samples 0) (error 'pcm-converter-drain "swr_convert drain failed: ~a" out-samples))
               (values (native->bytes out-ptr (* out-samples channels S32-BYTES)) out-samples))])))

  (define (pcm-converter-close! c)
    (when (and (pcm-converter? c) (not (pcm-converter-closed? c)))
      (set-pcm-converter-swr-ctx! c (swr_free (pcm-converter-swr-ctx c)))
      (ffmpeg-channel-layout-uninit! (pcm-converter-in-layout c))
      (ffmpeg-channel-layout-uninit! (pcm-converter-out-layout c))
      (set-pcm-converter-closed?! c #t))
    #t)

  ) ; end of module