package com.sherpaonnxofflinestt.utils

import java.io.File
import java.io.FileOutputStream
import java.io.RandomAccessFile
import java.nio.ByteBuffer
import java.nio.ByteOrder

/**
 * Utility class to write PCM audio data to a WAV file.
 * Supports streaming writes - call start(), write() multiple times, then close().
 */
class WavFileWriter(
    private val filePath: String,
    private val sampleRate: Int = 16000,
    private val channels: Int = 1,
    private val bitsPerSample: Int = 16
) {
    private var outputStream: FileOutputStream? = null
    private var totalSamplesWritten: Long = 0
    private var isOpen = false

    /**
     * Start writing to the WAV file. Writes a placeholder header.
     */
    fun start(): Boolean {
        return try {
            val file = File(filePath)
            file.parentFile?.mkdirs()
            outputStream = FileOutputStream(file)

            // Write placeholder header (will be updated on close)
            writeHeader(0)
            isOpen = true
            totalSamplesWritten = 0
            android.util.Log.i("WavFileWriter", "Started writing to $filePath")
            true
        } catch (e: Exception) {
            android.util.Log.e("WavFileWriter", "Failed to start: ${e.message}")
            false
        }
    }

    /**
     * Write audio samples (float array, values between -1.0 and 1.0)
     */
    fun write(samples: FloatArray) {
        if (!isOpen || outputStream == null) return

        try {
            val buffer = ByteBuffer.allocate(samples.size * 2)
            buffer.order(ByteOrder.LITTLE_ENDIAN)

            for (sample in samples) {
                // Convert float [-1.0, 1.0] to 16-bit PCM
                val pcmValue = (sample * 32767).toInt().coerceIn(-32768, 32767).toShort()
                buffer.putShort(pcmValue)
            }

            outputStream?.write(buffer.array())
            totalSamplesWritten += samples.size
        } catch (e: Exception) {
            android.util.Log.e("WavFileWriter", "Failed to write samples: ${e.message}")
        }
    }

    /**
     * Write audio samples from a list of float arrays
     */
    fun write(samplesList: List<FloatArray>) {
        for (samples in samplesList) {
            write(samples)
        }
    }

    /**
     * Close the file and update the header with correct sizes
     */
    fun close(): Boolean {
        if (!isOpen) return false

        return try {
            outputStream?.flush()
            outputStream?.close()
            outputStream = null

            android.util.Log.i("WavFileWriter", "Closed stream, updating header with $totalSamplesWritten samples")

            // Update header with actual data size
            val dataSize = (totalSamplesWritten * bitsPerSample / 8).toInt()
            android.util.Log.i("WavFileWriter", "Data size to write: $dataSize bytes")

            updateHeader()

            isOpen = false
            android.util.Log.i("WavFileWriter", "WAV file complete: $filePath")
            true
        } catch (e: Exception) {
            android.util.Log.e("WavFileWriter", "Failed to close: ${e.message}", e)
            false
        }
    }

    private fun writeHeader(dataSize: Int) {
        val byteRate = sampleRate * channels * bitsPerSample / 8
        val blockAlign = channels * bitsPerSample / 8

        val header = ByteBuffer.allocate(44)
        header.order(ByteOrder.LITTLE_ENDIAN)

        // RIFF header
        header.put("RIFF".toByteArray())
        header.putInt(36 + dataSize) // File size - 8
        header.put("WAVE".toByteArray())

        // fmt subchunk
        header.put("fmt ".toByteArray())
        header.putInt(16) // Subchunk1Size (16 for PCM)
        header.putShort(1) // AudioFormat (1 = PCM)
        header.putShort(channels.toShort())
        header.putInt(sampleRate)
        header.putInt(byteRate)
        header.putShort(blockAlign.toShort())
        header.putShort(bitsPerSample.toShort())

        // data subchunk
        header.put("data".toByteArray())
        header.putInt(dataSize)

        outputStream?.write(header.array())
    }

    private fun updateHeader() {
        val dataSize = (totalSamplesWritten * bitsPerSample / 8).toInt()
        android.util.Log.i("WavFileWriter", "updateHeader: dataSize=$dataSize, filePath=$filePath")

        try {
            RandomAccessFile(filePath, "rw").use { raf ->
                // Update RIFF chunk size at offset 4
                raf.seek(4)
                raf.write(intToLittleEndian(36 + dataSize))
                android.util.Log.i("WavFileWriter", "updateHeader: wrote RIFF size ${36 + dataSize}")

                // Update data chunk size at offset 40
                raf.seek(40)
                raf.write(intToLittleEndian(dataSize))
                android.util.Log.i("WavFileWriter", "updateHeader: wrote data size $dataSize")
            }
        } catch (e: Exception) {
            android.util.Log.e("WavFileWriter", "updateHeader failed: ${e.message}", e)
            throw e
        }
    }

    private fun intToLittleEndian(value: Int): ByteArray {
        return ByteBuffer.allocate(4).order(ByteOrder.LITTLE_ENDIAN).putInt(value).array()
    }

    val isRecording: Boolean
        get() = isOpen
}