package dev.kaizen.app.live

import android.Manifest
import android.content.Context
import android.content.pm.PackageManager
import android.media.AudioAttributes
import android.media.AudioFormat
import android.media.AudioRecord
import android.media.AudioTrack
import android.media.MediaRecorder
import android.util.Base64
import android.util.Log
import androidx.core.content.ContextCompat
import kotlinx.coroutines.*
import java.nio.ByteBuffer
import java.nio.ByteOrder
import kotlin.math.abs
import kotlin.math.sqrt

/**
 * Bidirectional audio pipeline for Live Call.
 *
 * Capture: AudioRecord (16kHz PCM16 mono, VOICE_COMMUNICATION for hardware AEC)
 *          → 100ms chunks → Base64 → [onChunk] callback.
 *
 * Playback: AudioTrack (24kHz PCM16 mono, LOW_LATENCY, MODE_STREAM)
 *           ← PCM16 chunks from server.
 *
 * Mic-ducking: gain reduced to 10% while the KI speaks (prevents echo feedback
 * loop on devices with weak AEC). Not a digital gain — we skip sending chunks
 * whose RMS is below the ducking threshold.
 */
class AudioPipeline(private val context: Context) {

    companion object {
        private const val TAG = "AudioPipeline"

        private const val CAPTURE_SAMPLE_RATE = 16_000
        private const val PLAYBACK_SAMPLE_RATE = 24_000
        private const val CHANNEL_IN = AudioFormat.CHANNEL_IN_MONO
        private const val CHANNEL_OUT = AudioFormat.CHANNEL_OUT_MONO
        private const val ENCODING = AudioFormat.ENCODING_PCM_16BIT

        private const val CHUNK_MS = 100
        private const val CAPTURE_CHUNK_SAMPLES = CAPTURE_SAMPLE_RATE * CHUNK_MS / 1000 // 1600
        private const val CAPTURE_CHUNK_BYTES = CAPTURE_CHUNK_SAMPLES * 2               // 3200

        private const val DUCKING_GAIN = 0.10f
    }

    var onChunk: ((base64Pcm: String) -> Unit)? = null
    var onAmplitude: ((Float) -> Unit)? = null

    private var recorder: AudioRecord? = null
    private var track: AudioTrack? = null
    private var captureJob: Job? = null

    @Volatile private var ducking = false
    @Volatile private var capturing = false

    // ─── Capture ──────────────────────────────────────────────────────────

    fun startCapture(scope: CoroutineScope): Boolean {
        if (ContextCompat.checkSelfPermission(context, Manifest.permission.RECORD_AUDIO) != PackageManager.PERMISSION_GRANTED) {
            Log.w(TAG, "RECORD_AUDIO permission not granted")
            return false
        }

        val minBuf = AudioRecord.getMinBufferSize(CAPTURE_SAMPLE_RATE, CHANNEL_IN, ENCODING)
        if (minBuf == AudioRecord.ERROR || minBuf == AudioRecord.ERROR_BAD_VALUE) {
            Log.e(TAG, "Invalid min buffer size: $minBuf")
            return false
        }

        val rec = AudioRecord(
            MediaRecorder.AudioSource.VOICE_COMMUNICATION,
            CAPTURE_SAMPLE_RATE,
            CHANNEL_IN,
            ENCODING,
            maxOf(minBuf * 4, CAPTURE_CHUNK_BYTES * 4),
        )

        if (rec.state != AudioRecord.STATE_INITIALIZED) {
            Log.e(TAG, "AudioRecord init failed")
            rec.release()
            return false
        }

        recorder = rec
        rec.startRecording()
        capturing = true

        captureJob = scope.launch(Dispatchers.IO) {
            val buffer = ShortArray(CAPTURE_CHUNK_SAMPLES)
            val byteBuffer = ByteBuffer.allocate(CAPTURE_CHUNK_BYTES).order(ByteOrder.LITTLE_ENDIAN)

            while (isActive && capturing) {
                val read = rec.read(buffer, 0, CAPTURE_CHUNK_SAMPLES)
                if (read <= 0) continue

                val rms = computeRms(buffer, read)
                onAmplitude?.invoke(rms)

                if (ducking && rms < DUCKING_GAIN) continue

                byteBuffer.clear()
                for (i in 0 until read) byteBuffer.putShort(buffer[i])

                val base64 = Base64.encodeToString(byteBuffer.array(), 0, read * 2, Base64.NO_WRAP)
                onChunk?.invoke(base64)
            }
        }

        return true
    }

    fun stopCapture() {
        capturing = false
        captureJob?.cancel()
        captureJob = null
        try {
            recorder?.stop()
        } catch (_: IllegalStateException) {}
        recorder?.release()
        recorder = null
    }

    // ─── Playback ─────────────────────────────────────────────────────────

    fun initPlayback(): Boolean {
        val minBuf = AudioTrack.getMinBufferSize(PLAYBACK_SAMPLE_RATE, CHANNEL_OUT, ENCODING)
        if (minBuf == AudioTrack.ERROR || minBuf == AudioTrack.ERROR_BAD_VALUE) {
            Log.e(TAG, "Invalid playback buffer size: $minBuf")
            return false
        }

        val t = AudioTrack.Builder()
            .setAudioAttributes(
                AudioAttributes.Builder()
                    .setUsage(AudioAttributes.USAGE_ASSISTANT)
                    .setContentType(AudioAttributes.CONTENT_TYPE_SPEECH)
                    .build()
            )
            .setAudioFormat(
                AudioFormat.Builder()
                    .setSampleRate(PLAYBACK_SAMPLE_RATE)
                    .setEncoding(ENCODING)
                    .setChannelMask(CHANNEL_OUT)
                    .build()
            )
            .setBufferSizeInBytes(maxOf(minBuf * 4, 24_000 * 2))
            .setPerformanceMode(AudioTrack.PERFORMANCE_MODE_LOW_LATENCY)
            .setTransferMode(AudioTrack.MODE_STREAM)
            .build()

        if (t.state != AudioTrack.STATE_INITIALIZED) {
            Log.e(TAG, "AudioTrack init failed")
            t.release()
            return false
        }

        track = t
        t.play()
        return true
    }

    fun playChunk(base64Pcm: String) {
        val bytes = Base64.decode(base64Pcm, Base64.NO_WRAP)
        track?.write(bytes, 0, bytes.size)
    }

    fun flushPlayback() {
        track?.pause()
        track?.flush()
        track?.play()
    }

    fun stopPlayback() {
        try {
            track?.pause()
            track?.flush()
        } catch (_: IllegalStateException) {}
        track?.release()
        track = null
    }

    // ─── Ducking ──────────────────────────────────────────────────────────

    fun setDucking(enabled: Boolean) {
        ducking = enabled
    }

    // ─── Cleanup ──────────────────────────────────────────────────────────

    fun release() {
        stopCapture()
        stopPlayback()
    }

    // ─── Helpers ──────────────────────────────────────────────────────────

    private fun computeRms(buffer: ShortArray, count: Int): Float {
        var sum = 0.0
        for (i in 0 until count) {
            val sample = buffer[i].toDouble() / Short.MAX_VALUE
            sum += sample * sample
        }
        return sqrt(sum / count).toFloat()
    }
}
