/*
 * Reader for ASF files (.wma), following the layout FFmpeg's libavformat/asfdec_f.c reads
 * (LGPL-2.1-or-later). It finds the first WMA 1/2 audio stream and hands back its media objects
 * (whole compressed audio chunks) in order, with their presentation times.
 */
package com.turntablemedic.auralrest.wma

import java.io.IOException

/** Random-access bytes: a file on disk, or a phone's content:// file via its descriptor. */
interface ByteInput {
    val size: Long
    /** Reads up to [len] bytes at [pos]; returns the count read, or -1 at the end. */
    fun read(pos: Long, buf: ByteArray, off: Int, len: Int): Int
}

/** A file this decoder can't play, with the reason in words a listener understands. */
class UnsupportedWma(reason: String) : IOException(reason)

class AsfReader(private val input: ByteInput) {

    class AudioObject(val data: ByteArray, val timeMs: Long, val packet: Long)

    // Stream description
    var formatTag = 0; private set
    var channels = 0; private set
    var sampleRate = 0; private set
    var bitRate = 0; private set
    var blockAlign = 0; private set
    var extradata = ByteArray(0); private set
    var prerollMs = 0L; private set
    /** Duration in milliseconds, preroll removed. */
    var durationMs = 0L; private set
    /** True when the file is locked with Windows Media DRM ("copy protect music"). */
    var isProtected = false; private set

    private var streamNumber = -1
    private var packetSize = 0
    private var dataStart = 0L      // first packet
    private var dataEnd = 0L
    private var packetCount = 0L
    private var dsSpan = 0
    private var dsPacketSize = 0
    private var dsChunkSize = 0

    // Reading state
    private var nextPacket = 0L
    private var pkt = ByteArray(0)
    private val ready = ArrayDeque<AudioObject>()
    private var cur: ByteArray? = null
    private var curNum = -1
    private var curFilled = 0
    private var curTime = 0L
    private var curPacket = 0L
    private var reading = 0L

    init { parseHeader() }

    val isWma12: Boolean get() = formatTag == 0x160 || formatTag == 0x161

    private fun readFully(pos: Long, n: Int): ByteArray {
        val b = ByteArray(n)
        var got = 0
        while (got < n) {
            val r = input.read(pos + got, b, got, n - got)
            if (r <= 0) throw IOException("File ends early")
            got += r
        }
        return b
    }

    private fun parseHeader() {
        val top = readFully(0, 30)
        if (!guidIs(top, 0, HEADER_GUID)) throw IOException("Not an ASF/WMA file")
        val headerSize = le64(top, 16)
        if (headerSize < 30 || headerSize > 64L * 1024 * 1024) throw IOException("Bad ASF header")
        val h = readFully(0, headerSize.toInt())
        var p = 30
        while (p + 24 <= h.size) {
            val size = le64(h, p + 16)
            if (size < 24 || p + size > h.size) break
            val body = p + 24
            when {
                guidIs(h, p, FILE_PROPS_GUID) -> {
                    packetCount = le64(h, body + 16 + 8 + 8)
                    val playDuration = le64(h, body + 16 + 8 + 8 + 8)   // 100 ns units
                    prerollMs = le64(h, body + 16 + 8 + 8 + 8 + 8 + 8)
                    packetSize = le32(h, body + 16 + 8 + 8 + 8 + 8 + 8 + 8 + 4)
                    durationMs = (playDuration / 10_000L - prerollMs).coerceAtLeast(0)
                }
                guidIs(h, p, STREAM_PROPS_GUID) && streamNumber < 0 -> parseStream(h, body)
                guidIs(h, p, CONTENT_ENCRYPTION_GUID) || guidIs(h, p, EXT_CONTENT_ENCRYPTION_GUID) -> isProtected = true
            }
            p += size.toInt()
        }
        if (isProtected) throw UnsupportedWma("it's copy-protected (Windows Media DRM), so only Windows can play it")
        if (streamNumber < 0) throw IOException("No audio in this file")
        if (packetSize <= 0) throw IOException("Unsupported ASF packet layout")
        val d = readFully(headerSize, 50)
        if (!guidIs(d, 0, DATA_GUID)) throw IOException("ASF data not found")
        dataStart = headerSize + 50
        val dataSize = le64(d, 16)
        dataEnd = if (dataSize >= 50) minOf(headerSize + dataSize, input.size) else input.size
        val fit = (dataEnd - dataStart) / packetSize
        if (packetCount <= 0 || packetCount > fit) packetCount = fit
        pkt = ByteArray(packetSize)
    }

    private fun parseStream(h: ByteArray, body: Int) {
        if (!guidIs(h, body, AUDIO_MEDIA_GUID)) return
        val ecGuidAt = body + 16
        val typeLen = le32(h, body + 40)
        val ecLen = le32(h, body + 44)
        val flags = le16(h, body + 48)
        val ts = body + 54
        if (typeLen < 18) return
        val tag = le16(h, ts)
        formatTag = tag
        channels = le16(h, ts + 2)
        sampleRate = le32(h, ts + 4)
        bitRate = le32(h, ts + 8) * 8
        blockAlign = le16(h, ts + 12)
        val cb = le16(h, ts + 16)
        extradata = h.copyOfRange(ts + 18, ts + 18 + minOf(cb, typeLen - 18))
        streamNumber = flags and 0x7f
        if (flags and 0x8000 != 0) isProtected = true
        if (guidIs(h, ecGuidAt, AUDIO_SPREAD_GUID) && ecLen >= 5) {
            val e = ts + typeLen
            dsSpan = h[e].toInt() and 255
            dsPacketSize = le16(h, e + 1)
            dsChunkSize = le16(h, e + 3)
            if (dsChunkSize == 0 || dsPacketSize / dsChunkSize <= 1 || dsPacketSize % dsChunkSize != 0) dsSpan = 0
        }
    }

    /** Next complete audio object, or null at the end of the file. */
    fun next(): AudioObject? {
        while (ready.isEmpty()) {
            if (nextPacket >= packetCount) return null
            readPacket(nextPacket++)
        }
        return ready.removeFirst()
    }

    /** Continue reading from packet [index]; objects that start in it or later come next. */
    fun seekPacket(index: Long) {
        ready.clear(); cur = null; curNum = -1
        nextPacket = index.coerceIn(0, packetCount)
    }

    private fun readPacket(index: Long) {
        reading = index
        val b = pkt
        try {
            val got = readFully(dataStart + index * packetSize, packetSize)
            System.arraycopy(got, 0, b, 0, packetSize)
        } catch (_: IOException) { nextPacket = packetCount; return }
        try { parsePacket(b) } catch (_: IndexOutOfBoundsException) { /* damaged packet: skip */ }
    }

    private fun parsePacket(b: ByteArray) {
        var p = 0
        var flags = b[p++].toInt() and 255
        if (flags and 0x80 != 0) {
            p += flags and 0x0f
            flags = b[p++].toInt() and 255
        }
        val prop = b[p++].toInt() and 255
        val explicitLen = readVar(b, p, flags ushr 5); p += lenOf(flags ushr 5)
        p += lenOf(flags ushr 1)   // sequence
        var padding = readVar(b, p, flags ushr 3); p += lenOf(flags ushr 3)
        p += 4 + 2                 // send time, duration
        if (explicitLen in 1 until packetSize) padding += packetSize - explicitLen
        val end = packetSize - padding
        val multiple = flags and 1 != 0
        var count = 1
        var payloadLenType = 0
        if (multiple) {
            val pf = b[p++].toInt() and 255
            count = pf and 0x3f
            payloadLenType = pf ushr 6
        }
        repeat(count) {
            if (p >= end) return
            val stream = b[p++].toInt() and 0x7f
            val objNum = readVar(b, p, prop ushr 4); p += lenOf(prop ushr 4)
            val offset = readVar(b, p, prop ushr 2); p += lenOf(prop ushr 2)
            val repLen = readVar(b, p, prop); p += lenOf(prop)
            var objSize = -1
            var objTime = 0L
            var compressed = false
            if (repLen == 1) {
                compressed = true
                objTime = offset.toLong()
                p += 1     // presentation time delta
            } else {
                if (repLen >= 8) {
                    objSize = le32(b, p)
                    objTime = le32(b, p + 4).toLong() and 0xffffffffL
                }
                p += repLen
            }
            val len = if (multiple) { val l = readVar(b, p, payloadLenType); p += lenOf(payloadLenType); l } else end - p
            if (len < 0 || p + len > packetSize) return
            if (stream == streamNumber) {
                if (compressed) {
                    var q = p
                    while (q < p + len) {
                        val sub = b[q++].toInt() and 255
                        if (q + sub > p + len) break
                        emit(b.copyOfRange(q, q + sub), objTime, reading)
                        q += sub
                    }
                } else addFragment(objNum, offset, objSize, objTime, b, p, len)
            }
            p += len
        }
    }

    private fun addFragment(num: Int, offset: Int, size: Int, time: Long, b: ByteArray, p: Int, len: Int) {
        if (offset == 0 || cur == null || num != curNum) {
            if (offset != 0 || size <= 0 || size > 4 * 1024 * 1024) { cur = null; return }
            cur = ByteArray(size); curNum = num; curFilled = 0; curTime = time; curPacket = reading
        }
        val c = cur!!
        if (offset != curFilled || offset + len > c.size) { cur = null; return }
        System.arraycopy(b, p, c, offset, len)
        curFilled += len
        if (curFilled == c.size) {
            cur = null
            emit(c, curTime, curPacket)
        }
    }

    private fun emit(data: ByteArray, time: Long, packet: Long) {
        val out = if (dsSpan > 1 && data.size == dsPacketSize * dsSpan) descramble(data) else data
        ready.addLast(AudioObject(out, time - prerollMs, packet))
    }

    private fun descramble(d: ByteArray): ByteArray {
        val out = ByteArray(d.size)
        var offset = 0
        while (offset < d.size) {
            val off = offset / dsChunkSize
            val row = off / dsSpan
            val col = off % dsSpan
            val idx = row + col * dsPacketSize / dsChunkSize
            System.arraycopy(d, idx * dsChunkSize, out, offset, dsChunkSize)
            offset += dsChunkSize
        }
        return out
    }

    private companion object {
        val HEADER_GUID = guid("75B22630-668E-11CF-A6D9-00AA0062CE6C")
        val FILE_PROPS_GUID = guid("8CABDCA1-A947-11CF-8EE4-00C00C205365")
        val STREAM_PROPS_GUID = guid("B7DC0791-A9B7-11CF-8EE6-00C00C205365")
        val AUDIO_MEDIA_GUID = guid("F8699E40-5B4D-11CF-A8FD-00805F5C442B")
        val AUDIO_SPREAD_GUID = guid("BFC3CD50-618F-11CF-8BB2-00AA00B4E220")
        val DATA_GUID = guid("75B22636-668E-11CF-A6D9-00AA0062CE6C")
        val CONTENT_ENCRYPTION_GUID = guid("2211B3FB-BD23-11D2-B4B7-00C04FB0E39A")
        val EXT_CONTENT_ENCRYPTION_GUID = guid("298AE614-2622-4C17-B935-DAE07EE9289C")

        /** GUID text to its on-disk bytes (first three groups little-endian). */
        fun guid(s: String): ByteArray {
            val h = s.replace("-", "")
            val raw = ByteArray(16) { h.substring(it * 2, it * 2 + 2).toInt(16).toByte() }
            val o = raw.copyOf()
            for (i in 0 until 4) o[i] = raw[3 - i]
            o[4] = raw[5]; o[5] = raw[4]; o[6] = raw[7]; o[7] = raw[6]
            return o
        }

        fun guidIs(b: ByteArray, at: Int, g: ByteArray): Boolean {
            if (at + 16 > b.size) return false
            for (i in 0 until 16) if (b[at + i] != g[i]) return false
            return true
        }

        fun lenOf(type: Int) = when (type and 3) { 0 -> 0; 1 -> 1; 2 -> 2; else -> 4 }
        fun readVar(b: ByteArray, p: Int, type: Int) = when (type and 3) {
            0 -> 0
            1 -> b[p].toInt() and 255
            2 -> le16(b, p)
            else -> le32(b, p)
        }
        fun le16(b: ByteArray, o: Int) = (b[o].toInt() and 255) or ((b[o + 1].toInt() and 255) shl 8)
        fun le32(b: ByteArray, o: Int) = le16(b, o) or (le16(b, o + 2) shl 16)
        fun le64(b: ByteArray, o: Int) = (le32(b, o).toLong() and 0xffffffffL) or (le32(b, o + 4).toLong() shl 32)
    }
}