package com.sherpaonnxofflinestt.managers

import com.sherpaonnxofflinestt.STTConstants
import com.sherpaonnxofflinestt.TenVad
import com.sherpaonnxofflinestt.models.VADState
import java.util.concurrent.atomic.AtomicReference
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

/**
 * Listener interface for VAD events
 */
interface VADListener {
    fun onSpeechStart()
    fun onSpeechContinue(floatSamples: FloatArray)
    fun onSpeechEnd(speechBuffer: List<FloatArray>)
    fun onStateChange(state: VADState, probability: Float, speechMs: Int, silenceMs: Int)
}

/**
 * Extended VAD listener with forced segment break support
 */
interface VADListenerExtended : VADListener {
    fun onForcedSegmentBreak(speechBuffer: List<FloatArray>)
}

/**
 * Manages TEN-VAD voice activity detection with state machine.
 * Thread-safe implementation using atomic state and synchronization.
 *
 * State machine: SILENCE -> SPEECH_START -> SPEECH -> SPEECH_END -> SILENCE
 */
class VADManager(
    private val eventManager: EventEmissionManager,
    private val listener: VADListener
) {
    // Thread-safe VAD instance access
    @Volatile
    private var tenVad: TenVad? = null

    // Atomic state for thread-safe state transitions
    private val vadStateRef = AtomicReference<VADState>(VADState.Silence)

    // Lock for state machine transitions
    private val stateLock = ReentrantLock()

    @Volatile
    private var speechStartTime: Long = 0
    @Volatile
    private var silenceStartTime: Long = 0

    // Thread-safe speech buffer
    private val speechBuffer: MutableList<FloatArray> = mutableListOf()
    private val bufferLock = ReentrantLock()

    private var vadFrameBuffer: ShortArray = ShortArray(0)
    private var vadFrameBufferPos: Int = 0
    private var vadHopSize: Int = STTConstants.VAD_HOP_SIZE

    // VAD configuration (using constants as defaults)
    var vadMinSpeechMs: Int = STTConstants.DEFAULT_MIN_SPEECH_MS
    var vadMinSilenceMs: Int = STTConstants.DEFAULT_MIN_SILENCE_MS
    var maxSpeechDurationMs: Int = STTConstants.DEFAULT_MAX_SPEECH_MS
    var speechPaddingMs: Int = STTConstants.DEFAULT_SPEECH_PADDING_MS
    var mode: String = "normal"            // aggressive, normal, sensitive

    // Continuous speech tracking for forced breaks
    @Volatile
    private var continuousSpeechMs: Long = 0
    @Volatile
    private var lastFrameTime: Long = 0

    // Helper property to get/set vadState through atomic reference
    private var vadState: VADState
        get() = vadStateRef.get()
        set(value) = vadStateRef.set(value)

    /**
     * Get effective threshold based on mode.
     * Uses STTConstants for threshold adjustments.
     */
    private fun getEffectiveThreshold(baseThreshold: Float): Float {
        return when (mode) {
            "aggressive" -> (baseThreshold + STTConstants.VAD_MODE_AGGRESSIVE_ADJUSTMENT)
                .coerceAtMost(STTConstants.VAD_THRESHOLD_MAX)  // Less sensitive
            "sensitive" -> (baseThreshold + STTConstants.VAD_MODE_SENSITIVE_ADJUSTMENT)
                .coerceAtLeast(STTConstants.VAD_THRESHOLD_MIN)  // More sensitive
            else -> baseThreshold  // normal
        }
    }

    /**
     * Initialize TEN-VAD.
     * Thread-safe: synchronized initialization.
     * @param hopSize Frame size for VAD (default from STTConstants)
     * @param threshold Speech probability threshold (default from STTConstants)
     */
    @Synchronized
    fun initialize(
        hopSize: Int = STTConstants.VAD_HOP_SIZE,
        threshold: Float = STTConstants.DEFAULT_VAD_THRESHOLD
    ) {
        vadHopSize = hopSize
        vadFrameBuffer = ShortArray(vadHopSize)
        vadFrameBufferPos = 0
        continuousSpeechMs = 0
        lastFrameTime = 0

        val effectiveThreshold = getEffectiveThreshold(threshold)

        tenVad = TenVad(
            hopSize = vadHopSize,
            threshold = effectiveThreshold
        )

        android.util.Log.i("VADManager", "VAD initialized with hopSize=$hopSize, threshold=$effectiveThreshold (mode=$mode)")
    }

    /**
     * Process audio samples through VAD
     * @param samples Raw audio samples (Short)
     * @param floatSamples Floating point samples for buffering
     * @param isInSpeech Whether currently in speech mode (for offline buffering) - DEPRECATED, now uses internal state
     */
    fun processAudio(samples: ShortArray, floatSamples: FloatArray?, isInSpeech: Boolean = false) {
        val vad = tenVad ?: return

        var offset = 0
        while (offset < samples.size) {
            // Fill the frame buffer
            val remaining = vadHopSize - vadFrameBufferPos
            val toCopy = minOf(remaining, samples.size - offset)

            System.arraycopy(samples, offset, vadFrameBuffer, vadFrameBufferPos, toCopy)
            vadFrameBufferPos += toCopy
            offset += toCopy

            // Process complete frames
            if (vadFrameBufferPos >= vadHopSize) {
                val result = vad.process(vadFrameBuffer)
                vadFrameBufferPos = 0

                if (result != null) {
                    processVADResult(result.isSpeech, result.probability, floatSamples)
                }
            }
        }

        // After VAD processing, if we're now in speech, buffer the audio chunk
        // This ensures the first chunk that triggers speech detection is also buffered
        val nowInSpeech = isInSpeech()
        if (floatSamples != null && nowInSpeech) {
            listener.onSpeechContinue(floatSamples.copyOf())
        }
    }

    /**
     * Get current VAD state
     */
    fun getCurrentState(): VADState = vadState

    /**
     * Check if currently in speech
     */
    fun isInSpeech(): Boolean = vadState == VADState.Speech || vadState == VADState.SpeechStart

    /**
     * Reset VAD state.
     * Thread-safe: synchronized method with buffer lock.
     */
    @Synchronized
    fun reset() {
        vadState = VADState.Silence
        speechStartTime = 0
        silenceStartTime = 0
        bufferLock.withLock {
            speechBuffer.clear()
        }
        vadFrameBufferPos = 0
        continuousSpeechMs = 0
        lastFrameTime = 0
    }

    /**
     * Release VAD resources.
     * Thread-safe: synchronized method.
     */
    @Synchronized
    fun destroy() {
        tenVad?.destroy()
        tenVad = null
        bufferLock.withLock {
            speechBuffer.clear()
        }
        android.util.Log.i("VADManager", "VAD destroyed")
    }

    /**
     * VAD state machine processing.
     * Thread-safe: uses stateLock for state transitions and bufferLock for buffer access.
     */
    private fun processVADResult(isSpeech: Boolean, probability: Float, currentChunk: FloatArray?) = stateLock.withLock {
        val currentTime = System.currentTimeMillis()
        val previousState = vadState

        // Convert VAD frame to float for speech buffer
        val floatFrame = FloatArray(vadFrameBuffer.size) { vadFrameBuffer[it] / 32768.0f }

        // Track frame timing for continuous speech duration
        val frameDurationMs = if (lastFrameTime > 0) currentTime - lastFrameTime else 10L
        lastFrameTime = currentTime

        if (isSpeech) {
            when (vadState) {
                VADState.Silence -> {
                    vadState = VADState.SpeechStart
                    speechStartTime = currentTime
                    speechBuffer.clear()
                    speechBuffer.add(floatFrame)
                    continuousSpeechMs = 0
                    listener.onSpeechStart()
                }
                VADState.SpeechStart -> {
                    speechBuffer.add(floatFrame)
                    continuousSpeechMs += frameDurationMs
                    val duration = currentTime - speechStartTime
                    if (duration >= vadMinSpeechMs) {
                        vadState = VADState.Speech
                    }
                }
                VADState.Speech -> {
                    speechBuffer.add(floatFrame)
                    silenceStartTime = 0
                    continuousSpeechMs += frameDurationMs

                    // Check for forced segment break due to max duration
                    if (continuousSpeechMs >= maxSpeechDurationMs) {
                        android.util.Log.i("VADManager", "Forced segment break at ${continuousSpeechMs}ms")

                        // Send forced break event
                        eventManager.sendVADUpdate(
                            "speech_end",
                            probability,
                            continuousSpeechMs.toInt(),
                            0
                        )

                        // Notify listener with accumulated speech buffer
                        val bufferCopy = speechBuffer.toList()
                        if (listener is VADListenerExtended) {
                            (listener as VADListenerExtended).onForcedSegmentBreak(bufferCopy)
                        } else {
                            listener.onSpeechEnd(bufferCopy)
                        }

                        // Reset for new segment but stay in speech state
                        speechBuffer.clear()
                        speechStartTime = currentTime
                        continuousSpeechMs = 0
                    }
                }
                VADState.SpeechEnd -> {
                    // Should not happen, but handle gracefully
                    vadState = VADState.Speech
                    continuousSpeechMs = 0
                }
            }
        } else {
            when (vadState) {
                VADState.SpeechStart -> {
                    // Speech was too short, cancel
                    vadState = VADState.Silence
                    speechBuffer.clear()
                    continuousSpeechMs = 0
                }
                VADState.Speech -> {
                    if (silenceStartTime == 0L) {
                        silenceStartTime = currentTime
                    }
                    val silenceDuration = currentTime - silenceStartTime
                    if (silenceDuration >= vadMinSilenceMs) {
                        vadState = VADState.SpeechEnd

                        // Send speech_end event BEFORE triggering processing
                        eventManager.sendVADUpdate(
                            "speech_end",
                            probability,
                            (currentTime - speechStartTime).toInt(),
                            silenceDuration.toInt()
                        )

                        // Notify listener with accumulated speech buffer
                        val bufferCopy = speechBuffer.toList()
                        listener.onSpeechEnd(bufferCopy)

                        // Transition back to silence
                        vadState = VADState.Silence
                        speechBuffer.clear()
                        continuousSpeechMs = 0
                    }
                }
                VADState.Silence, VADState.SpeechEnd -> {
                    // Already in silence, nothing to do
                    continuousSpeechMs = 0
                }
            }
        }

        // Notify state change if changed
        if (vadState != previousState && vadState != VADState.SpeechEnd) {
            listener.onStateChange(
                vadState,
                probability,
                if (vadState == VADState.Speech) (currentTime - speechStartTime).toInt() else 0,
                if (silenceStartTime > 0) (currentTime - silenceStartTime).toInt() else 0
            )

            eventManager.sendVADUpdate(
                vadState.name,
                probability,
                if (vadState == VADState.Speech) (currentTime - speechStartTime).toInt() else 0,
                if (silenceStartTime > 0) (currentTime - silenceStartTime).toInt() else 0
            )
        }
    }
}