package com.suno.android.media

import android.Manifest
import android.content.Context
import android.content.pm.PackageManager
import android.media.AudioFormat
import android.media.AudioRecord
import android.media.MediaRecorder
import androidx.core.content.ContextCompat
import dagger.hilt.android.qualifiers.ApplicationContext
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.MutableSharedFlow
import kotlinx.coroutines.flow.asSharedFlow
import kotlinx.coroutines.isActive
import kotlinx.coroutines.withContext
import java.io.File
import java.io.FileOutputStream
import java.nio.ByteBuffer
import java.nio.ByteOrder
import javax.inject.Inject

class AudioRecorder @Inject constructor(
    @ApplicationContext private val context: Context,
) {
    private var audioRecord: AudioRecord? = null

    @Volatile private var isRecording = false
    private val bufferSize = AudioRecord.getMinBufferSize(SAMPLE_RATE, CHANNEL_CONFIG, AUDIO_FORMAT)

    private val _audioDataFlow = MutableSharedFlow<FloatArray>(replay = 1)
    val audioDataFlow = _audioDataFlow.asSharedFlow()

    suspend fun startRecording(
        outputFile: File,
    ): Result<Unit> =
        withContext(Dispatchers.IO) {
            if (isRecording) return@withContext Result.failure(IllegalStateException("Recording already in progress"))

            if (!checkPermission()) {
                return@withContext Result.failure(SecurityException("RECORD_AUDIO permission not granted"))
            }

            try {
                audioRecord = AudioRecord(
                    MediaRecorder.AudioSource.MIC,
                    SAMPLE_RATE,
                    CHANNEL_CONFIG,
                    AUDIO_FORMAT,
                    bufferSize,
                )

                if (audioRecord?.state != AudioRecord.STATE_INITIALIZED) {
                    return@withContext Result.failure(IllegalStateException("AudioRecord initialization failed"))
                }

                audioRecord?.startRecording()
                isRecording = true

                val pcmFile = File(context.cacheDir, "audio_record.pcm")
                FileOutputStream(pcmFile).use { output ->
                    val buffer = ShortArray(bufferSize / 2)
                    while (isActive && isRecording) {
                        val readSize = audioRecord?.read(buffer, 0, buffer.size) ?: 0
                        if (readSize > 0) {
                            output.write(buffer.toByteArray(), 0, readSize * 2)
                            val floatBuffer = buffer.take(readSize).map {
                                it / Short.MAX_VALUE.toFloat()
                            }.toFloatArray()
                            _audioDataFlow.emit(floatBuffer)
                        }
                    }
                }

                convertPcmToWav(pcmFile, outputFile)
                Result.success(Unit)
            } catch (e: Exception) {
                stopRecording()
                Result.failure(e)
            }
        }

    fun stopRecording() {
        isRecording = false
        audioRecord?.stop()
        audioRecord?.release()
        audioRecord = null
    }

    private fun convertPcmToWav(
        pcmFile: File,
        wavFile: File,
    ) {
        val pcmSize = pcmFile.length().toInt()
        val header = createWavHeader(pcmSize)

        FileOutputStream(wavFile).use { wavOutput ->
            wavOutput.write(header)
            pcmFile.inputStream().use { it.copyTo(wavOutput) }
        }
    }

    private fun createWavHeader(
        pcmSize: Int,
    ): ByteArray {
        val totalSize = 36 + pcmSize
        val byteRate = SAMPLE_RATE * CHANNELS * BITS_PER_SAMPLE / 8

        return ByteBuffer.allocate(44).order(ByteOrder.LITTLE_ENDIAN).apply {
            put("RIFF".toByteArray())
            putInt(totalSize)
            put("WAVE".toByteArray())
            put("fmt ".toByteArray())
            putInt(16) // Subchunk1Size for PCM
            putShort(1) // AudioFormat (1 = PCM)
            putShort(CHANNELS.toShort())
            putInt(SAMPLE_RATE)
            putInt(byteRate)
            putShort((CHANNELS * BITS_PER_SAMPLE / 8).toShort())
            putShort(BITS_PER_SAMPLE.toShort())
            put("data".toByteArray())
            putInt(pcmSize)
        }.array()
    }

    private fun checkPermission(): Boolean =
        ContextCompat.checkSelfPermission(
            context,
            Manifest.permission.RECORD_AUDIO,
        ) == PackageManager.PERMISSION_GRANTED

    companion object {
        const val SAMPLE_RATE = 44100
        const val CHANNELS = 1
        const val BITS_PER_SAMPLE = 16
        const val CHANNEL_CONFIG = AudioFormat.CHANNEL_IN_MONO
        const val AUDIO_FORMAT = AudioFormat.ENCODING_PCM_16BIT
    }
}

fun ShortArray.toByteArray(): ByteArray {
    val bytes = ByteArray(size * 2)
    for (i in indices) {
        bytes[i * 2] = (this[i].toInt() and 0xFF).toByte()
        bytes[i * 2 + 1] = (this[i].toInt() shr 8 and 0xFF).toByte()
    }
    return bytes
}