ProgressResponseBody.kt

class ProgressResponseBody(
    private val uniqueId: String,
    private val delegate: ResponseBody
) : ResponseBody() {

    private val progressPublisher: MutableStateFlow<Progress>?

    init {
        addProgressPublisher(uniqueId)
        progressPublisher = progressPublishers[uniqueId]
    }

    private val bufferedSource: BufferedSource by lazy {
        source(delegate.source()).buffer()
    }

    override fun contentType(): MediaType? = delegate.contentType()

    override fun contentLength() = delegate.contentLength()

    override fun source() = bufferedSource

    private fun source(source: Source): Source {
        return object : ForwardingSource(source) {
            private var totalBytesRead = 0L

            @Throws(IOException::class)
            override fun read(sink: Buffer, byteCount: Long): Long {
                val bytesRead = super.read(sink, byteCount)

                if (bytesRead == -1L) {
                    //exhausted
                    this@ProgressResponseBody.log(SHOW_LOG) { "exhausted" }
                    progressPublisher?.value = Progress(
                        contentLength = contentLength(),
                        readLength = contentLength(),
                        progress = 100f
                    )
                    closeAndRemovePublisher()
                    return bytesRead
                }
                // read() returns the number of bytes read, or -1 if this source is exhausted.
                totalBytesRead += bytesRead

                val totalBytesToRead = contentLength()
                if (totalBytesToRead == -1L) {
                    this@ProgressResponseBody.log(SHOW_LOG) { "finished" }
                    progressPublisher?.value = Progress(
                        contentLength = contentLength(),
                        readLength = contentLength(),
                        progress = 100f
                    )
                    closeAndRemovePublisher()
                } else {
                    val readProgress = (totalBytesRead.toFloat() / totalBytesToRead) * 100f
                    progressPublisher?.value = Progress(
                        contentLength = contentLength(),
                        readLength = totalBytesRead,
                        progress = readProgress
                    )
                    if (readProgress.roundToInt() == 100) {
                        closeAndRemovePublisher()
                    }
                    this@ProgressResponseBody.log(SHOW_LOG) { "progress: $readProgress" }
                }
                return bytesRead
            }
        }
    }

    private fun closeAndRemovePublisher() {
        removeProgressPublisher(uniqueId)
    }

    companion object {
        private const val SHOW_LOG = false
        private val progressPublishers =
            mutableMapOf<String, MutableStateFlow<Progress>>()

        /*
        * Use this to get Progress updates as Flow from anywhere (probably viewModel).
        * takeWhile is needed to complete the flow, otherwise the flow never completes.
        * Ref: https://github.com/Kotlin/kotlinx.coroutines/issues/1973#issuecomment-633898538
        * */
        fun getProgressPublisher(uniqueId: String) =
            progressPublishers[uniqueId]?.takeWhile { it.progress != 100f && it.readLength != it.contentLength }

        //call this before sending the request
        private fun addProgressPublisher(uniqueId: String) {
            progressPublishers[uniqueId] = MutableStateFlow(
                Progress(
                    contentLength = Long.MAX_VALUE,
                    readLength = 0L,
                    progress = 0f
                )
            )
        }

        //call this after sending the request
        private fun removeProgressPublisher(uniqueId: String) {
            progressPublishers.remove(uniqueId)
        }
    }

    data class Progress(val contentLength: Long, val readLength: Long, val progress: Float)
}