(module audio-encoder racket/base

  (require racket/path
           racket/string
           racket/contract
           racket/runtime-path
           "flac-encoder.rkt"
           "opus-encoder.rkt"
           "taglib.rkt"
           "private/pcm-converter.rkt"
           "private/utils.rkt")

  (provide audio-encode
           audio-supported-encoder-extensions
           audio-register-encoder!
           make-audio-encoder
           audio-encoder?)

  (define-struct audio-encoder (exts open write finish settings))

  (define-runtime-module-path-index audio-decoder-module "audio-decoder.rkt")

  (define audio-encoders (make-hash))

  (define (audio-register-encoder! type encoder)
    (hash-set! audio-encoders type encoder))

  (audio-register-encoder!
   'flac
   (make-audio-encoder '("flac")
                       flac-encoder-open
                       flac-encoder-write
                       flac-encoder-finish
                       flac-encoder-prepare-settings))

  (audio-register-encoder!
   'opus
   (make-audio-encoder '("opus" "oga")
                       opus-encoder-open
                       opus-encoder-write
                       opus-encoder-finish
                       opus-encoder-prepare-settings))

  (define (audio-supported-encoder-extensions)
    (apply append (map audio-encoder-exts (hash-values audio-encoders))))

  (define (path-extension-symbol file)
    (let ((ext (path-get-extension (build-path file))))
      (and ext (string->symbol (string-downcase (substring (bytes->string/utf-8 ext) 1))))))

  (define (encoder-for-output output-file explicit-kind)
    (let ((kind (or explicit-kind (path-extension-symbol output-file))))
      (cond [(and kind (hash-ref audio-encoders kind #f)) (values kind (hash-ref audio-encoders kind))]
            [else (error 'audio-encode "cannot infer encoder from output file ~a" output-file)])))

  (define (tag-value-copy! src dst getter setter empty?)
    (let ((v (getter src)))
      (unless (empty? v) (setter dst v))))

  (define (empty-string? v) (or (eq? v #f) (and (string? v) (string=? v ""))))
  (define (empty-number? v) (or (eq? v #f) (and (number? v) (< v 0))))

  (define (merge-hash a b)
    (let ((out (make-hash)))
      (when (hash? a)
        (for-each (lambda (k) (hash-set! out k (hash-ref a k))) (hash-keys a)))
      (when (hash? b)
        (for-each (lambda (k) (hash-set! out k (hash-ref b k))) (hash-keys b)))
      out))

  (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 (maybe-string s) (and (string? s) (not (string=? s "")) s))
  (define (maybe-number n) (and (number? n) (>= n 0) (number->string n)))

  (define (source-tags->opus-settings input-file settings)
    ;; For Opus, embedded pictures must be written into the OpusTags packet
    ;; before the encoder starts.  TagLib post-processing is not reliable for
    ;; this path, so transfer the regular comments and cover art through
    ;; libopusenc comments instead.
    (with-handlers ([exn:fail? (lambda (e)
                                 (warn-sound "Could not read source tags from ~a for Opus comments: ~a"
                                             input-file (exn-message e))
                                 settings)])
      (call-with-id3-tags
       input-file
       (lambda (src)
         (if (not (tags-valid? src))
             settings
             (let ((out (copy-hash settings))
                   (comments (make-hash)))
               (let ((title (maybe-string (tags-title src)))) (when title (hash-set! comments 'title title)))
               (let ((album (maybe-string (tags-album src)))) (when album (hash-set! comments 'album album)))
               (let ((artist (maybe-string (tags-artist src)))) (when artist (hash-set! comments 'artist artist)))
               (let ((comment (maybe-string (tags-comment src)))) (when comment (hash-set! comments 'comment comment)))
               (let ((genre (maybe-string (tags-genre src)))) (when genre (hash-set! comments 'genre genre)))
               (let ((composer (maybe-string (tags-composer src)))) (when composer (hash-set! comments 'composer composer)))
               (let ((album-artist (maybe-string (tags-album-artist src)))) (when album-artist (hash-set! comments 'albumartist album-artist)))
               (let ((year (maybe-number (tags-year src)))) (when year (hash-set! comments 'date year)))
               (let ((track (maybe-number (tags-track src)))) (when track (hash-set! comments 'tracknumber track)))
               (let ((disc (tags-disc-number src)))
                 (cond [(string? disc) (unless (string=? disc "") (hash-set! comments 'discnumber disc))]
                       [(and (number? disc) (>= disc 0)) (hash-set! comments 'discnumber (number->string disc))]
                       [else (void)]))
               (unless (null? (hash-keys comments)) (hash-set! out 'comments comments))
               (let ((picture (tags-picture src)))
                 (unless (eq? picture #f) (hash-set! out 'picture picture)))
               out)))
       #:mode 'read)))

  (define (make-tag-result method success? picture note)
    (let ((h (make-hash)))
      (hash-set! h 'method method)
      (hash-set! h 'success? success?)
      (hash-set! h 'picture? (not (eq? picture #f)))
      (when (id3-picture? picture)
        (hash-set! h 'picture-size (id3-picture-size picture))
        (hash-set! h 'picture-mimetype (id3-picture-mimetype picture)))
      (when note (hash-set! h 'note note))
      h))

  (define (copy-tags! input-file output-file)
    (with-handlers ([exn:fail? (lambda (e)
                                 (warn-sound "Could not copy tags from ~a to ~a: ~a"
                                             input-file output-file (exn-message e))
                                 (make-tag-result 'taglib-post-copy #f #f (exn-message e)))])
      (call-with-id3-tags
       input-file
       (lambda (src)
         (call-with-id3-tags
          output-file
          (lambda (dst)
            (if (and (tags-valid? src) (tags-valid? dst))
                (begin
                  (tag-value-copy! src dst tags-title tags-title! empty-string?)
                  (tag-value-copy! src dst tags-album tags-album! empty-string?)
                  (tag-value-copy! src dst tags-artist tags-artist! empty-string?)
                  (tag-value-copy! src dst tags-comment tags-comment! empty-string?)
                  (tag-value-copy! src dst tags-genre tags-genre! empty-string?)
                  (tag-value-copy! src dst tags-composer tags-composer! empty-string?)
                  (tag-value-copy! src dst tags-album-artist tags-album-artist! empty-string?)
                  (tag-value-copy! src dst tags-year tags-year! empty-number?)
                  (tag-value-copy! src dst tags-track tags-track! empty-number?)
                  (tag-value-copy! src dst tags-disc-number tags-disc-number! empty-number?)
                  (let ((picture (tags-picture src)))
                    (unless (eq? picture #f) (tags-picture! dst picture))
                    (tags-save! dst)
                    (make-tag-result 'taglib-post-copy #t picture #f)))
                (make-tag-result 'taglib-post-copy #f #f "source or destination tags invalid")))
          #:mode 'read-write))
       #:mode 'read)))

  (define (input-frames-in-buffer fmt buf-len)
    (let* ((channels (hash-ref fmt 'channels 1))
           (bits (hash-ref fmt 'bits-per-sample (hash-ref fmt 'pcm-bits-per-sample 16)))
           (bytes-per-sample (max 1 (quotient bits 8)))
           (frame-bytes (* channels bytes-per-sample)))
      (if (> frame-bytes 0) (quotient buf-len frame-bytes) 0)))

  (define (total-input-frames fmt)
    (and (hash? fmt)
         (or (hash-ref fmt 'total-samples #f)
             (hash-ref fmt 'total-frames #f)
             (hash-ref fmt 'frames #f))))

  (define (audio-encode input-file output-file settings
                        #:encoder [explicit-kind #f]
                        #:copy-tags? [copy-tags? #t]
                        #:progress-callback [progress-callback #f])
    (define-values (kind encoder) (encoder-for-output output-file explicit-kind))
    (define effective-settings (if (and copy-tags? (eq? kind 'opus))
                                   (source-tags->opus-settings input-file settings)
                                   settings))
    (define backend-handle #f)
    (define format #f)
    (define output-format #f)
    (define converter #f)
    (define frames-written 0)
    (define frames-read 0)
    (define last-progress -1.0)
    (define tags-result #f)

    (define (progress! phase input-format)
      (when progress-callback
        (let* ((total (total-input-frames input-format))
               (progress (and (integer? total) (> total 0)
                              (min 1.0 (/ frames-read total))))
               (h (make-hash)))
          (hash-set! h 'phase phase)
          (hash-set! h 'encoder kind)
          (hash-set! h 'input input-file)
          (hash-set! h 'output output-file)
          (hash-set! h 'frames-read frames-read)
          (hash-set! h 'frames-written frames-written)
          (hash-set! h 'total-frames total)
          (hash-set! h 'progress progress)
          (hash-set! h 'input-format input-format)
          (when output-format (hash-set! h 'output-format output-format))
          (progress-callback h)
          (when (number? progress) (set! last-progress progress)))))

    (define (ensure-open! fmt)
      (when (eq? backend-handle #f)
        ;; Record the resolved output format, not merely the incoming PCM format.
        ;; This matters when only FLAC bit depth changes, because no swresample
        ;; converter is needed but the resulting FLAC stream metadata still differs.
        (set! output-format ((audio-encoder-settings encoder) effective-settings fmt))
        (set! backend-handle ((audio-encoder-open encoder) output-file effective-settings fmt))))

    (define (write-backend! fmt buffer buf-len)
      (ensure-open! fmt)
      (set! frames-written (+ frames-written ((audio-encoder-write encoder) backend-handle fmt buffer buf-len))))

    (define (ensure-converter! input-format)
      ;; FLAC may need conversion because the caller requested a target sample
      ;; rate or bit depth.  Opus is deliberately not routed through this
      ;; converter by default: libopusenc accepts the source input rate and has
      ;; its own resampler, and opus-encoder.rkt feeds it float PCM directly.
      (when (and (eq? kind 'flac) (eq? converter #f))
        (when (pcm-conversion-needed? input-format effective-settings)
          (set! converter (make-pcm-converter input-format effective-settings)))))

    (define (write-converted! input-format buffer buf-len)
      (ensure-converter! input-format)
      (cond [converter
             (let-values (((out out-samples) (pcm-converter-convert converter buffer buf-len input-format)))
               (when (> out-samples 0)
                 (write-backend! (pcm-converter-output-format converter) out (bytes-length out))))]
            [else (write-backend! input-format buffer buf-len)]))

    (define (drain-converter!)
      (when converter
        (let loop ()
          (let-values (((out out-samples) (pcm-converter-drain converter)))
            (when (> out-samples 0)
              (write-backend! (pcm-converter-output-format converter) out (bytes-length out))
              (loop))))))

    (define (on-format audio-kind ao-kind handle fmt)
      ;; Keep stream metadata, but delay encoder creation until the first audio
      ;; buffer.  Some decoders report an output-oriented stream format first
      ;; and then the exact PCM frame format in buf-info.
      (set! format fmt)
      (progress! 'format fmt))

    (define (on-audio audio-kind ao-kind handle buf-info buffer buf-len)
      (let ((effective-format (merge-hash format buf-info)))
        (set! format effective-format)
        (set! frames-read (+ frames-read (input-frames-in-buffer effective-format buf-len)))
        (write-converted! effective-format buffer buf-len)
        (progress! 'audio effective-format)))

    (let* ((audio-open-proc (dynamic-require audio-decoder-module 'audio-open))
           (audio-read-proc (dynamic-require audio-decoder-module 'audio-read))
           (decoder (audio-open-proc input-file on-format on-audio)))
      (dynamic-wind
       void
       (lambda () (audio-read-proc decoder))
       (lambda ()
         (dynamic-wind
          drain-converter!
          (lambda () (when backend-handle ((audio-encoder-finish encoder) backend-handle)))
          (lambda () (when converter (pcm-converter-close! converter)))))))

    (progress! 'finished-encoding format)
    (set! tags-result
          (cond [(not copy-tags?) (make-tag-result 'none #t #f "tag copy disabled")]
                [(eq? kind 'opus)
                 (make-tag-result 'libopusenc-comments #t (hash-ref effective-settings 'picture #f) #f)]
                [else (copy-tags! input-file output-file)]))
    (progress! 'finished format)
    (let ((r (make-hash)))
      (hash-set! r 'encoder kind)
      (hash-set! r 'input input-file)
      (hash-set! r 'output output-file)
      (hash-set! r 'input-format format)
      (hash-set! r 'output-format output-format)
      (hash-set! r 'frames-read frames-read)
      (hash-set! r 'frames-written frames-written)
      (hash-set! r 'tag-copy tags-result)
      r))

  ) ; end of module