package net.siteed.audiostream

import android.media.MediaCodec
import android.media.MediaCodecInfo
import android.media.MediaFormat
import android.os.Handler
import android.os.HandlerThread
import java.io.File
import java.io.FileOutputStream
import java.io.IOException
import java.nio.ByteBuffer
import java.util.concurrent.CountDownLatch
import java.util.concurrent.TimeUnit
import java.util.concurrent.atomic.AtomicBoolean

/**
 * Direct-to-file AAC-LC encoder that consumes the same PCM bytes produced
 * by the capture [android.media.AudioRecord] loop and writes an ADTS-framed
 * `.aac` file on disk. AAC bytes never cross the JS bridge.
 *
 * Crashproof by construction: every ADTS frame starts with the sync word
 * `0xFFF9`, and no global header or footer is required. The OS flushes the
 * underlying [FileOutputStream] in the background, so a `SIGKILL` only
 * corrupts at most the last (partial) frame — the file remains decodable
 * up to the last complete frame.
 *
 * All encoder and file I/O work runs on a dedicated [HandlerThread]
 * (`aac-encoder`). [encode] hands off copies of the incoming PCM buffer to
 * that thread so the AudioRecord read loop is never blocked.
 */
class AACFileWriter {

    companion object {
        private const val TAG = "AACFileWriter"
        private const val MIME = "audio/mp4a-latm"

        /** ADTS sample-rate table (MPEG-4 Part 3, Table 1.18). */
        val ADTS_SAMPLE_RATE_TABLE: IntArray = intArrayOf(
            96_000, 88_200, 64_000, 48_000, 44_100, 32_000,
            24_000, 22_050, 16_000, 12_000, 11_025, 8_000, 7_350
        )

        fun adtsSampleRateIndex(sampleRate: Int): Int {
            val idx = ADTS_SAMPLE_RATE_TABLE.indexOf(sampleRate)
            return if (idx >= 0) idx else -1
        }

        /**
         * Builds an ADTS (Audio Data Transport Stream) frame header — 7
         * bytes, no CRC, MPEG-2. Mirrors the iOS [AACFileWriter] exactly.
         *
         * @param packetLength Size of the AAC raw data block that follows.
         * @param sampleRateIndex Index in [ADTS_SAMPLE_RATE_TABLE].
         * @param channelConfig 1 for mono, 2 for stereo.
         */
        fun makeADTSHeader(
            packetLength: Int,
            sampleRateIndex: Int,
            channelConfig: Int
        ): ByteArray {
            val profile = 2 // AAC-LC
            val frameLength = packetLength + 7
            val header = ByteArray(7)
            header[0] = 0xFF.toByte()
            header[1] = 0xF9.toByte() // MPEG-2, no CRC
            header[2] = (((profile - 1) shl 6) or (sampleRateIndex shl 2) or (channelConfig shr 2)).toByte()
            header[3] = (((channelConfig and 0x3) shl 6) or ((frameLength shr 11) and 0x3)).toByte()
            header[4] = ((frameLength shr 3) and 0xFF).toByte()
            header[5] = (((frameLength and 0x7) shl 5) or 0x1F).toByte()
            header[6] = 0xFC.toByte()
            return header
        }
    }

    private val isRunning = AtomicBoolean(false)
    private var handlerThread: HandlerThread? = null
    private var handler: Handler? = null
    private var encoder: MediaCodec? = null
    private var fos: FileOutputStream? = null
    private var file: File? = null

    // Encoder parameters (set at start())
    private var sampleRate: Int = 0
    private var channels: Int = 1
    private var bitrate: Int = 0
    private var sampleRateIndex: Int = 0
    private var channelConfig: Int = 1
    private var bytesPerSample: Int = 2 // PCM 16-bit
    private var presentationTimeUs: Long = 0L

    /**
     * Opens the output file and starts the encoder. Must be called before
     * [encode].
     *
     * @param sampleRate Target sample rate. Must exist in
     *                   [ADTS_SAMPLE_RATE_TABLE].
     * @param channels 1 (mono) or 2 (stereo).
     * @param bitrate AAC-LC target bitrate in bps.
     * @param outFile Destination file. Truncated if it exists.
     * @throws IllegalArgumentException For invalid parameters.
     * @throws IOException For file / codec initialisation failures.
     */
    @Throws(IOException::class)
    fun start(sampleRate: Int, channels: Int, bitrate: Int, outFile: File) {
        val idx = adtsSampleRateIndex(sampleRate)
        require(idx >= 0) { "Sample rate $sampleRate is not in the ADTS sample-rate table" }
        require(channels in 1..2) { "Unsupported channel count: $channels (only 1 or 2)" }

        // Prepare file (truncate if exists, ensure parent dir).
        outFile.parentFile?.mkdirs()
        if (outFile.exists() && !outFile.delete()) {
            LogUtils.w(TAG, "Could not remove existing file at ${outFile.absolutePath}")
        }
        if (!outFile.createNewFile()) {
            throw IOException("Could not create file at ${outFile.absolutePath}")
        }
        val stream = FileOutputStream(outFile, false)

        // Dedicated HandlerThread so the AudioRecord read loop is never blocked.
        val thread = HandlerThread("aac-encoder").apply { start() }
        val h = Handler(thread.looper)

        // Configure MediaCodec synchronously (config errors should fail fast).
        val format = MediaFormat.createAudioFormat(MIME, sampleRate, channels).apply {
            setInteger(MediaFormat.KEY_AAC_PROFILE, MediaCodecInfo.CodecProfileLevel.AACObjectLC)
            setInteger(MediaFormat.KEY_BIT_RATE, bitrate)
            setInteger(MediaFormat.KEY_MAX_INPUT_SIZE, 16_384)
        }
        val codec = try {
            MediaCodec.createEncoderByType(MIME).apply {
                configure(format, null, null, MediaCodec.CONFIGURE_FLAG_ENCODE)
                start()
            }
        } catch (e: Exception) {
            try { stream.close() } catch (_: Exception) {}
            thread.quitSafely()
            throw IOException("Failed to configure MediaCodec AAC encoder", e)
        }

        this.sampleRate = sampleRate
        this.channels = channels
        this.bitrate = bitrate
        this.sampleRateIndex = idx
        this.channelConfig = channels
        this.bytesPerSample = 2
        this.presentationTimeUs = 0L
        this.file = outFile
        this.fos = stream
        this.encoder = codec
        this.handlerThread = thread
        this.handler = h
        this.isRunning.set(true)

        LogUtils.i(TAG, "Started ${outFile.name} @ ${sampleRate}Hz / ${channels}ch / ${bitrate}bps")
    }

    /**
     * Feeds a copy of the given PCM slice to the encoder asynchronously.
     * Returns immediately; the actual encoding + file I/O runs on the
     * internal `aac-encoder` HandlerThread.
     */
    fun encode(pcm: ByteArray, length: Int) {
        if (!isRunning.get()) return
        if (length <= 0) return
        val copy = pcm.copyOfRange(0, length)
        handler?.post { encodeChunk(copy) }
    }

    /**
     * Drains any output packets the encoder can produce right now and
     * flushes the [FileOutputStream] to disk. Does not signal EOS nor
     * close anything; encoding can continue afterwards.
     *
     * Used on pause / audio-focus loss so the file on disk reflects all
     * packets the encoder has already produced.
     */
    fun flush() {
        val h = handler ?: return
        if (!isRunning.get()) return
        val latch = CountDownLatch(1)
        h.post {
            try {
                drainOutput(endOfStream = false)
                fos?.flush()
            } catch (e: Exception) {
                LogUtils.w(TAG, "flush failed: ${e.message}", e)
            } finally {
                latch.countDown()
            }
        }
        latch.await(2, TimeUnit.SECONDS)
    }

    /**
     * Signals end-of-stream, drains all remaining packets, then closes
     * both the encoder and the output file. Returns the final file.
     */
    fun stop(): File? {
        val h = handler
        val result = file
        if (h == null || !isRunning.compareAndSet(true, false)) {
            return result
        }
        val latch = CountDownLatch(1)
        h.post {
            try {
                val codec = encoder
                if (codec != null) {
                    try {
                        val index = codec.dequeueInputBuffer(200_000)
                        if (index >= 0) {
                            codec.queueInputBuffer(
                                index, 0, 0, presentationTimeUs,
                                MediaCodec.BUFFER_FLAG_END_OF_STREAM
                            )
                        }
                        drainOutput(endOfStream = true)
                    } catch (e: Exception) {
                        LogUtils.w(TAG, "stop: drain failed: ${e.message}", e)
                    }
                    try { codec.stop() } catch (_: Exception) {}
                    try { codec.release() } catch (_: Exception) {}
                }
                try {
                    fos?.flush()
                    // Best-effort fsync so the OS durably writes the trailing
                    // bytes; not strictly needed for ADTS correctness.
                    fos?.fd?.sync()
                } catch (_: Exception) { }
                try { fos?.close() } catch (_: Exception) {}
            } finally {
                encoder = null
                fos = null
                latch.countDown()
            }
        }
        latch.await(5, TimeUnit.SECONDS)
        handlerThread?.quitSafely()
        handlerThread = null
        handler = null
        LogUtils.i(TAG, "Stopped ${result?.name ?: "<null>"}")
        return result
    }

    // ---- Internal (all invoked on the encoder thread) ---------------------

    private fun encodeChunk(pcm: ByteArray) {
        val codec = encoder ?: return
        if (!isRunning.get()) return
        var offset = 0
        while (offset < pcm.size && isRunning.get()) {
            val index = try {
                codec.dequeueInputBuffer(10_000)
            } catch (e: IllegalStateException) {
                LogUtils.w(TAG, "dequeueInputBuffer failed: ${e.message}")
                return
            }
            if (index < 0) {
                // No input buffer right now — opportunistically drain output
                // so the encoder can free some input slots, then retry.
                drainOutput(endOfStream = false)
                continue
            }
            val inputBuffer: ByteBuffer = try {
                codec.getInputBuffer(index) ?: run {
                    LogUtils.w(TAG, "Null input buffer at index $index")
                    return
                }
            } catch (e: IllegalStateException) {
                LogUtils.w(TAG, "getInputBuffer failed: ${e.message}")
                return
            }
            inputBuffer.clear()
            val remaining = inputBuffer.remaining()
            val take = minOf(remaining, pcm.size - offset)
            inputBuffer.put(pcm, offset, take)
            try {
                codec.queueInputBuffer(index, 0, take, presentationTimeUs, 0)
            } catch (e: IllegalStateException) {
                LogUtils.w(TAG, "queueInputBuffer failed: ${e.message}")
                return
            }
            val samples = take.toLong() / (bytesPerSample * channels)
            if (sampleRate > 0) {
                presentationTimeUs += (samples * 1_000_000L) / sampleRate.toLong()
            }
            offset += take
            drainOutput(endOfStream = false)
        }
    }

    private fun drainOutput(endOfStream: Boolean) {
        val codec = encoder ?: return
        val info = MediaCodec.BufferInfo()
        // Short timeout on the first call, then non-blocking. When draining
        // at EOS we loop until we actually see the EOS flag.
        val firstTimeoutUs = if (endOfStream) 200_000L else 0L
        var firstIteration = true
        while (true) {
            val index = try {
                codec.dequeueOutputBuffer(info, if (firstIteration) firstTimeoutUs else 0L)
            } catch (e: IllegalStateException) {
                LogUtils.w(TAG, "dequeueOutputBuffer failed: ${e.message}")
                return
            }
            firstIteration = false
            when {
                index == MediaCodec.INFO_TRY_AGAIN_LATER -> {
                    if (endOfStream) {
                        // Keep polling a bit; the encoder will eventually emit EOS.
                        try {
                            val retryIndex = codec.dequeueOutputBuffer(info, 100_000L)
                            if (retryIndex < 0) return
                            handleOutputIndex(codec, retryIndex, info)
                            if ((info.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) return
                        } catch (e: IllegalStateException) {
                            return
                        }
                    } else {
                        return
                    }
                }
                index == MediaCodec.INFO_OUTPUT_FORMAT_CHANGED -> continue
                index < 0 -> continue
                else -> {
                    handleOutputIndex(codec, index, info)
                    if ((info.flags and MediaCodec.BUFFER_FLAG_END_OF_STREAM) != 0) return
                }
            }
        }
    }

    private fun handleOutputIndex(codec: MediaCodec, index: Int, info: MediaCodec.BufferInfo) {
        try {
            val outputBuffer = codec.getOutputBuffer(index)
            if (outputBuffer != null && info.size > 0 &&
                (info.flags and MediaCodec.BUFFER_FLAG_CODEC_CONFIG) == 0
            ) {
                val header = makeADTSHeader(info.size, sampleRateIndex, channelConfig)
                val aacData = ByteArray(info.size)
                outputBuffer.position(info.offset)
                outputBuffer.limit(info.offset + info.size)
                outputBuffer.get(aacData)
                fos?.let { stream ->
                    stream.write(header)
                    stream.write(aacData)
                }
            }
        } catch (e: IOException) {
            LogUtils.e(TAG, "File write failed: ${e.message}", e)
        } catch (e: IllegalStateException) {
            LogUtils.w(TAG, "handleOutputIndex failed: ${e.message}")
        } finally {
            try { codec.releaseOutputBuffer(index, false) } catch (_: Exception) {}
        }
    }
}