WaveformView.kt

package com.keyboardr.bluejay.ui.playlist

import android.content.Context
import android.graphics.Canvas
import android.graphics.Color
import android.graphics.Paint
import android.media.MediaCodec
import android.media.MediaExtractor
import android.media.MediaFormat
import android.util.AttributeSet
import android.util.Log
import android.view.View
import com.keyboardr.bluejay.model.MediaItem
import com.keyboardr.bluejay.util.*
import kotlinx.coroutines.*
import kotlinx.coroutines.flow.collect
import kotlinx.coroutines.flow.flow
import kotlin.math.absoluteValue
import kotlin.math.log10
import kotlin.math.sqrt

@FlowPreview
class WaveformView @JvmOverloads constructor(context: Context,
                                             attrs: AttributeSet? = null,
                                             defStyle: Int = 0) : View(context, attrs, defStyle), LogScope {

  private var scope = if (isInEditMode) GlobalScope else CoroutineScopeRegistry()
  private var activeJob: Job? = null

  // each item in range [0, 1]
  private var formArray: FloatArray? = null

  var mediaItem: MediaItem? = null
    set(value) {
      if (field == value) return
      field = value
      invalidateFormArray()
    }

  var position: Long? = null
    set(value) {
      field = value
      invalidate()
    }

  private val slicePaint = Paint().apply {
    color = Color.WHITE
  }
  private val progressPaint = Paint().apply {
    color = Color.BLACK
    alpha = 0x80
  }

  override fun onDraw(canvas: Canvas) {
    super.onDraw(canvas)

    val duration = mediaItem?.duration?.takeUnless { it == 0L }
    val progressFraction = duration?.let { position?.toFloat()?.div(it) } ?: 0f
    canvas.drawRect(0f, 0f, progressFraction * width, height.toFloat(), progressPaint)

    formArray?.forEachIndexed { x, size ->
      val y = (height * size / 2).coerceAtLeast(1f)
      canvas.drawLine(x.toFloat(), height / 2 - y, x.toFloat(), height / 2 + y, slicePaint)
    }
        ?: canvas.drawLine(0F, (height / 2).toFloat(), width.toFloat(), (height / 2).toFloat(), slicePaint)
  }

  override fun onAttachedToWindow() {
    super.onAttachedToWindow()
    (scope as? CoroutineScopeRegistry)?.start()
  }

  override fun onDetachedFromWindow() {
    super.onDetachedFromWindow()
    (scope as? CoroutineScopeRegistry)?.stop()
  }

  override fun onMeasure(widthMeasureSpec: Int, heightMeasureSpec: Int) {
    val width = widthMeasureSpec.specSize()
    val height = heightMeasureSpec.specSize().coerceAtMost(48.dpToPx(context))
    setMeasuredDimension(width, height)
    if (width != formArray?.size) {
      invalidateFormArray()
    }
  }

  private fun invalidateFormArray() {
    formArray = if (measuredWidth == 0) null else FloatArray(measuredWidth)
    activeJob?.cancel()
    activeJob = null
    position = null

    val localFormArray = formArray ?: return
    val localMediaItem = mediaItem ?: return
    val path = localMediaItem.path ?: return
    val duration = localMediaItem.duration
    val usPerSlice = duration * 1000 / (localFormArray.size - 1)

    activeJob = scope.launch(Dispatchers.IO) {
      var slice = 0

      val extractor = MediaExtractor()
      try {
        extractor.setDataSource(path)
        val codec = extractor.selectTrack() ?: run {
          Log.e(LOG_TAG, "Can't find audio info")
          return@launch
        }
        codec.start()
        flow {
          val info = MediaCodec.BufferInfo()
          var outputFormat = codec.outputFormat
          var isEOS = false

          val sliceBuffer = SliceBuffer()

          do {
            if (!isEOS) {
              val inBufferId = codec.dequeueInputBuffer(10000)
              if (inBufferId >= 0) {
                val buffer = codec.getInputBuffer(inBufferId)!!
                val sampleSize = extractor.readSampleData(buffer, 0)
                if (sampleSize < 0) {
                  // We shouldn't stop the playback at this point, just pass the EOS
                  // flag to decoder, we will get it again from the
                  // dequeueOutputBuffer
                  Log.d(LOG_TAG, "InputBuffer BUFFER_FLAG_END_OF_STREAM")
                  codec.queueInputBuffer(inBufferId, 0, 0, 0, MediaCodec.BUFFER_FLAG_END_OF_STREAM)
                  isEOS = true
                } else {
                  codec.queueInputBuffer(inBufferId, 0, sampleSize, extractor.sampleTime, 0)
                  extractor.advance()
                }
              }
            }

            when (val outBufferId = codec.dequeueOutputBuffer(info, 10000)) {
              MediaCodec.INFO_OUTPUT_FORMAT_CHANGED -> {
                outputFormat = codec.outputFormat
                Log.d(LOG_TAG, "New format $outputFormat")
              }
              MediaCodec.INFO_TRY_AGAIN_LATER -> Log.w(LOG_TAG, "dequeueOutputBuffer timed out!")
              else -> {
                val numChannels = outputFormat.getInteger(MediaFormat.KEY_CHANNEL_COUNT)
                val outputBuffer = codec.getOutputBuffer(outBufferId)!!
                while (outputBuffer.hasRemaining()) {
                  sliceBuffer.add(outputBuffer.getShort())
                  if (numChannels > 1) {
                    for (i in 2 until numChannels) {
                      if (outputBuffer.hasRemaining()) {
                        outputBuffer.getShort()
                      }
                    }
                  }
                }

                codec.releaseOutputBuffer(outBufferId, true)
                if (info.presentationTimeUs > (slice * usPerSlice)) {
                  emit(Slice(slice, sliceBuffer.toDbRms()))
                  slice++
                  sliceBuffer.clear()
                }
              }
            }

          } while (!info.flags.hasFlag(MediaCodec.BUFFER_FLAG_END_OF_STREAM))
        }.collect { (slice, value) ->
          localFormArray[slice] = value
          launch(Dispatchers.Main) { invalidate() }
        }

        codec.stop()
        codec.release()
      } finally {
        extractor.release()
      }
      activeJob = null
    }

  }

  private fun SliceBuffer.toDbRms(): Float {
    var sum = 0f
    for (index in 0 until size) {
      val element = get(index)
      val db = log10(element.toFraction().absoluteValue*9 + 1)
      sum += db * db
    }

    return sqrt(sum / size)
  }

  private fun MediaExtractor.selectTrack(): MediaCodec? {
    for (track in 0 until trackCount) {
      val format = getTrackFormat(track)
      val mime = format.mimeType
      if (mime.startsWith("audio/")) {
        selectTrack(track)
        return MediaCodec.createDecoderByType(mime).apply {
          configure(format, null, null, 0)
        }
      }
    }
    return null
  }

  private val MediaFormat.mimeType
    get() = getString(MediaFormat.KEY_MIME)
}

private data class Slice(val index: Int, val value: Float)

private class SliceBuffer(initialSize: Int = 1000) {

  private var array = ShortArray(initialSize)

  var size = 0
    private set

  operator fun get(index: Int): Short {
    return array[index]
  }

  fun add(value: Short) {
    if (size == array.size) {
      array = array.copyOf(array.size * 2)
    }
    array[size++] = value
  }

  fun clear() {
    array.fill(0, 0, (size - 1).coerceAtLeast(0))
    size = 0
  }

}

private fun Int.specSize() = View.MeasureSpec.getSize(this)