package com.turntablemedic.auralrest.wma

import java.io.IOException

/**
 * A .wma file as a stream of 16-bit interleaved PCM: the ASF reader feeding the WMA decoder.
 * Supports WMA 1 and 2 (what Windows Media Player made when ripping CDs).
 *
 * On open, every audio chunk's one-byte frame count is skimmed (nothing is decoded) to map the file,
 * which gives the exact length and lets seeks land on the exact sample.
 */
class WmaPcm(private val input: ByteInput) {
    private val asf = AsfReader(input)
    private val dec: WmaDecoder

    val sampleRate: Int get() = asf.sampleRate
    val channels: Int get() = asf.channels

    // Map: for each audio object, the packet it starts in and the decoder output count before it.
    private var objPacket = LongArray(0)
    private var objOut = LongArray(0)
    private var objCount = 0

    /** Exact number of sample frames this stream produces. */
    var totalFrames = 0L
        private set

    private var pending = Array(0) { FloatArray(0) }
    private var pendingAt = 0
    private var drop = 0L          // decoder output frames still to throw away
    private var ended = false

    init {
        if (!asf.isWma12) throw UnsupportedWma(when (asf.formatTag) {
            0x162 -> "it's WMA Pro, which this phone can't play"
            0x163 -> "it's WMA Lossless, which this phone can't play"
            0x0a -> "it's WMA Voice, which this phone can't play"
            else -> "its type of WMA isn't supported"
        })
        dec = WmaDecoder(if (asf.formatTag == 0x160) 1 else 2, asf.channels, asf.sampleRate, asf.bitRate, asf.blockAlign, asf.extradata)
        buildMap()
        asf.seekPacket(0)
        drop = dec.skipSamples.toLong()
    }

    private fun buildMap() {
        val ba = asf.blockAlign
        var out = 0L
        var hasPrev = false
        var n = 0
        var pk = LongArray(1024)
        var oo = LongArray(1024)
        while (true) {
            val obj = asf.next() ?: break
            if (n == pk.size) { pk = pk.copyOf(n * 2); oo = oo.copyOf(n * 2) }
            pk[n] = obj.packet; oo[n] = out; n++
            var p = 0
            while (p + ba <= obj.data.size) {
                if (dec.useBitReservoir) {
                    val frames = (obj.data[p].toInt() and 0x0f) - (if (hasPrev) 0 else 1)
                    if (frames < 0) hasPrev = false
                    else { out += frames.toLong() * dec.frameLen; hasPrev = true }
                } else out += dec.frameLen
                p += ba
            }
        }
        objPacket = pk; objOut = oo; objCount = n
        // Plus the final half-frame drained at the end, minus the decoder's start-up delay.
        totalFrames = (out + dec.frameLen - dec.skipSamples).coerceAtLeast(0)
    }

    /** Fills [out] with interleaved 16-bit samples, little-endian. Returns bytes written, -1 at the end. */
    fun read(out: ByteArray, off: Int, len: Int): Int {
        val frameBytes = channels * 2
        var n = 0
        while (n + frameBytes <= len) {
            if (pending.isEmpty() || pendingAt >= pending[0].size) {
                if (!refill()) break
                continue
            }
            val avail = pending[0].size - pendingAt
            val take = minOf(avail, (len - n) / frameBytes)
            for (i in 0 until take) {
                for (ch in 0 until channels) {
                    var v = pending[ch][pendingAt + i] * 32768f
                    v = if (v > 32767f) 32767f else if (v < -32768f) -32768f else v
                    val s = Math.round(v)
                    out[off + n] = s.toByte()
                    out[off + n + 1] = (s shr 8).toByte()
                    n += 2
                }
            }
            pendingAt += take
        }
        return if (n == 0 && ended) -1 else n
    }

    private fun decodeObject(obj: AsfReader.AudioObject): Array<FloatArray>? {
        val ba = asf.blockAlign
        val all = ArrayList<Array<FloatArray>>()
        var p = 0
        while (p + ba <= obj.data.size) {
            val s = dec.decodeSuperframe(obj.data, p)
            if (s != null && s[0].isNotEmpty()) all += s
            p += ba
        }
        if (all.size == 1) return all[0]
        val frames = all.sumOf { it[0].size }
        if (frames == 0) return null
        val merged = Array(channels) { FloatArray(frames) }
        var at = 0
        for (s in all) {
            for (ch in 0 until channels) System.arraycopy(s[ch], 0, merged[ch], at, s[ch].size)
            at += s[0].size
        }
        return merged
    }

    private fun use(samples: Array<FloatArray>): Boolean {
        val frames = samples[0].size
        var start = 0
        if (drop > 0) {
            val d = minOf(drop, frames.toLong()).toInt()
            drop -= d; start = d
        }
        if (start >= frames) return false
        pending = samples; pendingAt = start
        return true
    }

    /** Decodes more samples into [pending]; false at the end. */
    private fun refill(): Boolean {
        while (true) {
            if (ended) return false
            val obj = asf.next()
            if (obj == null) {
                ended = true
                // The last half-frame still waiting in the decoder's overlap buffer.
                return use(dec.drain())
            }
            val s = decodeObject(obj) ?: continue
            if (use(s)) return true
        }
    }

    /** Moves to sample frame [frame] exactly. */
    fun seek(frame: Long) {
        val target = frame.coerceIn(0, totalFrames) + dec.skipSamples   // in decoder output terms
        pending = Array(0) { FloatArray(0) }; pendingAt = 0; ended = false
        dec.flush()
        if (objCount == 0) { asf.seekPacket(0); drop = target; return }
        // The last object starting at or before the target...
        var lo = 0; var hi = objCount - 1
        while (lo < hi) { val mid = (lo + hi + 1) / 2; if (objOut[mid] <= target) lo = mid else hi = mid - 1 }
        val j = lo
        // ...and one before it, decoded and thrown away so the decoder's carry-over is right.
        val w = if (j > 0) j - 1 else 0
        asf.seekPacket(objPacket[w])
        var first = w
        while (first > 0 && objPacket[first - 1] == objPacket[w]) first--
        repeat(w - first) { asf.next() }
        if (w < j) asf.next()?.let { decodeObject(it) }
        drop = target - objOut[j]
    }
}