package xyz.cottongin.radio247.audio

import java.io.ByteArrayOutputStream

data class Mp3FrameInfo(
    val sampleRate: Int,
    val channelCount: Int,
    val samplesPerFrame: Int
)

private data class ParsedHeader(
    val frameSize: Int,
    val sampleRate: Int,
    val channelCount: Int,
    val samplesPerFrame: Int
)

// MPEG1 Layer 3 bitrates (kbps), index 0 and 15 invalid
private val MPEG1_L3_BITRATE = intArrayOf(0, 32, 40, 48, 56, 64, 80, 96, 112, 128, 160, 192, 224, 256, 320, 0)

// MPEG2/2.5 Layer 3 bitrates (kbps)
private val MPEG2_L3_BITRATE = intArrayOf(0, 8, 16, 24, 32, 40, 48, 56, 64, 80, 96, 112, 128, 144, 160, 0)

// Sample rates (Hz)
private val MPEG1_SAMPLE_RATE = intArrayOf(44100, 48000, 32000, 0)
private val MPEG2_SAMPLE_RATE = intArrayOf(22050, 24000, 16000, 0)
private val MPEG25_SAMPLE_RATE = intArrayOf(11025, 12000, 8000, 0)

class Mp3FrameSync(
    private val onFrame: (ByteArray) -> Unit
) {
    private val buffer = ByteArrayOutputStream()

    var detectedFormat: Mp3FrameInfo? = null
        private set

    fun feed(data: ByteArray, offset: Int = 0, length: Int = data.size) {
        buffer.write(data, offset, length)
        processBuffer()
    }

    fun flush() {
        processBuffer()
    }

    private fun processBuffer() {
        val bytes = buffer.toByteArray()
        var pos = 0

        while (pos + 4 <= bytes.size) {
            if (!isSyncWord(bytes, pos)) {
                pos++
                continue
            }

            val header = parseHeader(bytes, pos)
            if (header == null) {
                pos++
                continue
            }
            val frameSize = header.frameSize

            if (pos + frameSize > bytes.size) {
                break
            }

            val nextFramePos = pos + frameSize
            if (nextFramePos + 4 <= bytes.size) {
                if (!isSyncWord(bytes, nextFramePos) || parseHeader(bytes, nextFramePos) == null) {
                    pos++
                    continue
                }
            }

            if (detectedFormat == null) {
                detectedFormat = Mp3FrameInfo(
                    sampleRate = header.sampleRate,
                    channelCount = header.channelCount,
                    samplesPerFrame = header.samplesPerFrame
                )
            }

            val frame = bytes.copyOfRange(pos, pos + frameSize)
            onFrame(frame)
            pos += frameSize
        }

        buffer.reset()
        if (pos < bytes.size) {
            buffer.write(bytes, pos, bytes.size - pos)
        }
    }

    private fun isSyncWord(bytes: ByteArray, pos: Int): Boolean {
        val b0 = bytes[pos].toInt() and 0xFF
        val b1 = bytes[pos + 1].toInt() and 0xFF
        return b0 == 0xFF && (b1 and 0xE0) == 0xE0
    }

    companion object {
        fun parseFrameInfo(frame: ByteArray): Mp3FrameInfo? {
            if (frame.size < 4) return null
            val b0 = frame[0].toInt() and 0xFF
            val b1 = frame[1].toInt() and 0xFF
            if (b0 != 0xFF || (b1 and 0xE0) != 0xE0) return null

            val b2 = frame[2].toInt() and 0xFF
            val b3 = frame[3].toInt() and 0xFF

            val header = (b0 shl 24) or (b1 shl 16) or (b2 shl 8) or b3
            val mpegVersion = (header ushr 19) and 0x03
            val layer = (header ushr 17) and 0x03
            if (mpegVersion == 1 || layer == 0) return null

            val sampleRateIndex = (b2 shr 2) and 0x03
            if (sampleRateIndex == 3) return null

            val sampleRate = when (mpegVersion) {
                3 -> MPEG1_SAMPLE_RATE[sampleRateIndex]
                2 -> MPEG2_SAMPLE_RATE[sampleRateIndex]
                0 -> MPEG25_SAMPLE_RATE[sampleRateIndex]
                else -> return null
            }
            if (sampleRate == 0) return null

            val channelMode = (b3 shr 6) and 0x03
            val channelCount = if (channelMode == 3) 1 else 2

            val samplesPerFrame = when (layer) {
                3 -> 384                                       // Layer I
                2 -> 1152                                      // Layer II
                1 -> if (mpegVersion == 3) 1152 else 576       // Layer III
                else -> return null
            }

            return Mp3FrameInfo(sampleRate, channelCount, samplesPerFrame)
        }
    }

    private fun parseHeader(bytes: ByteArray, pos: Int): ParsedHeader? {
        if (pos + 4 > bytes.size) return null
        val header = ((bytes[pos].toInt() and 0xFF) shl 24) or
            ((bytes[pos + 1].toInt() and 0xFF) shl 16) or
            ((bytes[pos + 2].toInt() and 0xFF) shl 8) or
            (bytes[pos + 3].toInt() and 0xFF)

        val b2 = bytes[pos + 2].toInt() and 0xFF
        val b3 = bytes[pos + 3].toInt() and 0xFF

        val mpegVersion = (header ushr 19) and 0x03
        val layer = (header ushr 17) and 0x03
        val bitrateIndex = (b2 shr 4) and 0x0F
        val sampleRateIndex = (b2 shr 2) and 0x03
        val padding = ((b2 shr 1) and 0x01) == 1

        if (mpegVersion == 1) return null
        if (layer == 0) return null
        if (bitrateIndex == 0 || bitrateIndex == 15) return null
        if (sampleRateIndex == 3) return null

        val bitrate = when (mpegVersion) {
            3 -> MPEG1_L3_BITRATE[bitrateIndex]
            2 -> MPEG2_L3_BITRATE[bitrateIndex]
            0 -> MPEG2_L3_BITRATE[bitrateIndex]
            else -> return null
        }
        if (bitrate == 0) return null

        val sampleRate = when (mpegVersion) {
            3 -> MPEG1_SAMPLE_RATE[sampleRateIndex]
            2 -> MPEG2_SAMPLE_RATE[sampleRateIndex]
            0 -> MPEG25_SAMPLE_RATE[sampleRateIndex]
            else -> return null
        }
        if (sampleRate == 0) return null

        val channelMode = (b3 shr 6) and 0x03
        val channelCount = if (channelMode == 3) 1 else 2

        val samplesPerFrame = when (layer) {
            3 -> 384                                       // Layer I
            2 -> 1152                                      // Layer II
            1 -> if (mpegVersion == 3) 1152 else 576       // Layer III
            else -> return null
        }

        val frameSize = when (layer) {
            1 -> { // Layer III
                when (mpegVersion) {
                    3 -> (144 * bitrate * 1000 / sampleRate) + if (padding) 1 else 0
                    else -> (72 * bitrate * 1000 / sampleRate) + if (padding) 1 else 0
                }
            }
            2 -> { // Layer II
                when (mpegVersion) {
                    3 -> (144 * bitrate * 1000 / sampleRate) + if (padding) 1 else 0
                    else -> (72 * bitrate * 1000 / sampleRate) + if (padding) 1 else 0
                }
            }
            3 -> { // Layer I
                ((12 * bitrate * 1000 / sampleRate) + if (padding) 1 else 0) * 4
            }
            else -> return null
        }

        if (frameSize <= 0) return null
        return ParsedHeader(frameSize, sampleRate, channelCount, samplesPerFrame)
    }
}
