package com.sherpaonnxofflinestt.managers

import com.k2fsa.sherpa.onnx.*
import com.sherpaonnxofflinestt.STTConstants
import com.sherpaonnxofflinestt.utils.STTLogger
import java.io.File

/**
 * Result from STT processing with metrics
 */
data class STTProcessingResult(
    val text: String,
    val processingTimeMs: Long,
    val confidence: Float = 0f
)

/**
 * Result of model validation check
 */
enum class ModelValidationResult {
    Valid,
    NotFound,
    NotReadable,
    TooSmall,
    InvalidFormat
}

/**
 * Listener interface for STT recognition events
 */
interface STTRecognitionListener {
    fun onPartialResult(text: String, processingTimeMs: Long, confidence: Float)
    fun onFinalResult(text: String, processingTimeMs: Long, confidence: Float)
    fun onEndpoint()
    fun onError(code: String, message: String)
}

/**
 * Manages sherpa-onnx speech recognition (both streaming and offline modes).
 *
 * Features:
 * - Model validation before loading
 * - Streaming (online) and batch (offline) recognition
 * - Configurable ONNX provider (CPU, NNAPI, GPU)
 * - Thread-safe operations
 */
class STTRecognitionManager(
    private val listener: STTRecognitionListener
) {
    private var onlineRecognizer: OnlineRecognizer? = null
    private var offlineRecognizer: OfflineRecognizer? = null
    private var onlineStream: OnlineStream? = null
    private var modelType: String = "streaming"

    var sampleRate: Int = STTConstants.DEFAULT_SAMPLE_RATE

    /**
     * Validate an ONNX model file.
     * Checks existence, readability, minimum size, and ONNX format.
     * @param path Path to the model file
     * @return Validation result
     */
    fun validateModel(path: String): ModelValidationResult {
        val file = File(path)
        if (!file.exists()) return ModelValidationResult.NotFound
        if (!file.canRead()) return ModelValidationResult.NotReadable
        if (file.length() < 1024) return ModelValidationResult.TooSmall

        // Check ONNX magic bytes: The ONNX format starts with a protobuf header
        // For ONNX files, we check for reasonable file content
        return try {
            file.inputStream().use { stream ->
                val header = ByteArray(8)
                val bytesRead = stream.read(header)
                if (bytesRead < 8) {
                    return ModelValidationResult.TooSmall
                }
                // ONNX files typically start with certain protobuf patterns
                // A simple heuristic: check if it's not a text file or obviously wrong format
                if (header.all { it in 32..126 } && String(header).startsWith("<!DOCTYPE")) {
                    return ModelValidationResult.InvalidFormat
                }
                ModelValidationResult.Valid
            }
        } catch (e: Exception) {
            STTLogger.Recognition.e("Model validation failed for $path: ${e.message}")
            ModelValidationResult.NotReadable
        }
    }

    /**
     * Validate all required model files for initialization.
     * @param architecture Model architecture: "transducer" or "whisper"
     * @return Map of file type to validation result (empty map means all valid)
     */
    fun validateModels(
        encoderPath: String,
        decoderPath: String,
        joinerPath: String,
        tokensPath: String,
        architecture: String = "transducer"
    ): Map<String, ModelValidationResult> {
        val results = mutableMapOf<String, ModelValidationResult>()

        val encoderResult = validateModel(encoderPath)
        if (encoderResult != ModelValidationResult.Valid) {
            results["encoder"] = encoderResult
        }

        val decoderResult = validateModel(decoderPath)
        if (decoderResult != ModelValidationResult.Valid) {
            results["decoder"] = decoderResult
        }

        // Joiner is only required for transducer architecture
        if (architecture == "transducer") {
            val joinerResult = validateModel(joinerPath)
            if (joinerResult != ModelValidationResult.Valid) {
                results["joiner"] = joinerResult
            }
        }

        // Tokens file is a text file, just check existence and readability
        val tokensFile = File(tokensPath)
        if (!tokensFile.exists()) {
            results["tokens"] = ModelValidationResult.NotFound
        } else if (!tokensFile.canRead()) {
            results["tokens"] = ModelValidationResult.NotReadable
        }

        return results
    }

    /**
     * Initialize STT recognizer
     * @param encoderPath Full path to encoder ONNX model file
     * @param decoderPath Full path to decoder ONNX model file
     * @param joinerPath Full path to joiner ONNX model file (required for transducer, empty for whisper)
     * @param tokensPath Full path to tokens file
     * @param modelType "streaming" or "offline"
     * @param modelArchitecture "transducer" or "whisper"
     * @param whisperLanguage Target language for whisper ("" for auto-detect, "en", "fr", etc.)
     * @param whisperTask "transcribe" or "translate" for whisper models
     * @param provider ONNX Runtime provider ("cpu", "nnapi", "gpu")
     * @param numThreads Number of threads for inference
     * @return true if initialization successful
     */
    fun initialize(
        encoderPath: String,
        decoderPath: String,
        joinerPath: String,
        tokensPath: String,
        modelType: String = "streaming",
        modelArchitecture: String = "transducer",
        whisperLanguage: String = "",
        whisperTask: String = "transcribe",
        provider: String = STTConstants.DEFAULT_PROVIDER,
        numThreads: Int = STTConstants.DEFAULT_STT_NUM_THREADS
    ): Boolean {
        this.modelType = modelType

        // Validate all model files
        val validationErrors = validateModels(encoderPath, decoderPath, joinerPath, tokensPath, modelArchitecture)
        if (validationErrors.isNotEmpty()) {
            for ((fileType, error) in validationErrors) {
                STTLogger.Recognition.e("Model validation failed for $fileType: $error")
            }
            listener.onError("MODEL_VALIDATION_ERROR", "Model validation failed: $validationErrors")
            return false
        }

        STTLogger.Recognition.i("Model files validated successfully")
        STTLogger.Recognition.i("  architecture: $modelArchitecture")
        STTLogger.Recognition.i("  encoder: $encoderPath")
        STTLogger.Recognition.i("  decoder: $decoderPath")
        if (modelArchitecture == "transducer") {
            STTLogger.Recognition.i("  joiner: $joinerPath")
        }
        STTLogger.Recognition.i("  tokens: $tokensPath")

        return try {
            val startTime = System.currentTimeMillis()

            if (modelType == "offline") {
                STTLogger.Recognition.i("Loading OFFLINE STT model (architecture: $modelArchitecture)...")

                val offlineConfig = when (modelArchitecture) {
                    "whisper" -> {
                        // Whisper architecture: encoder-decoder only
                        STTLogger.Recognition.i("Whisper config: language='$whisperLanguage', task='$whisperTask'")
                        OfflineRecognizerConfig(
                            modelConfig = OfflineModelConfig(
                                whisper = OfflineWhisperModelConfig(
                                    encoder = encoderPath,
                                    decoder = decoderPath,
                                    language = whisperLanguage, // Empty = auto-detect
                                    task = whisperTask,
                                    tailPaddings = 1000
                                ),
                                tokens = tokensPath,
                                numThreads = numThreads,
                                debug = false,
                                provider = provider,
                                modelType = "whisper"
                            ),
                            decodingMethod = "greedy_search"
                        )
                    }
                    else -> {
                        // Transducer architecture: encoder-decoder-joiner
                        val offlineModelType = if (encoderPath.contains("parakeet")) "nemo_transducer" else "transducer"
                        OfflineRecognizerConfig(
                            modelConfig = OfflineModelConfig(
                                transducer = OfflineTransducerModelConfig(
                                    encoder = encoderPath,
                                    decoder = decoderPath,
                                    joiner = joinerPath
                                ),
                                tokens = tokensPath,
                                numThreads = numThreads,
                                debug = false,
                                provider = provider,
                                modelType = offlineModelType
                            ),
                            decodingMethod = "greedy_search"
                        )
                    }
                }
                offlineRecognizer = OfflineRecognizer(assetManager = null, config = offlineConfig)
            } else {
                // Streaming mode - only transducer supported
                if (modelArchitecture == "whisper") {
                    listener.onError("INVALID_CONFIG", "Whisper models only support offline mode")
                    return false
                }
                android.util.Log.i("STTRecognitionManager", "Loading STREAMING STT model...")
                val recognizerConfig = OnlineRecognizerConfig(
                    modelConfig = OnlineModelConfig(
                        transducer = OnlineTransducerModelConfig(
                            encoder = encoderPath,
                            decoder = decoderPath,
                            joiner = joinerPath
                        ),
                        tokens = tokensPath,
                        numThreads = numThreads,
                        debug = false,
                        provider = provider
                    ),
                    endpointConfig = EndpointConfig(
                        rule1 = EndpointRule(false, 2.4f, 0.0f),
                        rule2 = EndpointRule(true, 1.2f, 0.0f),
                        rule3 = EndpointRule(false, 0.0f, 20.0f)
                    ),
                    enableEndpoint = true
                )
                onlineRecognizer = OnlineRecognizer(assetManager = null, config = recognizerConfig)
            }

            android.util.Log.i("STTRecognitionManager", "STT model loaded in ${System.currentTimeMillis() - startTime}ms")
            true
        } catch (e: Exception) {
            android.util.Log.e("STTRecognitionManager", "Failed to initialize: ${e.message}", e)
            listener.onError("STT_INIT_ERROR", "Failed to initialize STT: ${e.message}")
            false
        }
    }

    /**
     * Create a new streaming session (for online mode)
     */
    fun createStream() {
        if (modelType == "streaming") {
            onlineStream?.release()
            onlineStream = onlineRecognizer?.createStream()
        }
    }

    /**
     * Process audio samples through streaming recognizer
     * @param samples Float audio samples
     * @param lastTranscript The last sent transcript to detect changes
     * @return Processing result with text and timing, or null if unchanged
     */
    fun processStreamingAudio(samples: FloatArray, lastTranscript: String): STTProcessingResult? {
        val recognizer = onlineRecognizer ?: return null
        val stream = onlineStream ?: return null

        val startTime = System.currentTimeMillis()

        stream.acceptWaveform(samples, sampleRate)

        while (recognizer.isReady(stream)) {
            recognizer.decode(stream)
        }

        val processingTime = System.currentTimeMillis() - startTime

        val result = recognizer.getResult(stream)
        if (!result.text.isNullOrEmpty() && result.text != lastTranscript) {
            listener.onPartialResult(result.text, processingTime, 0f)
            return STTProcessingResult(result.text, processingTime, 0f)
        }

        // Check for endpoint
        if (recognizer.isEndpoint(stream)) {
            val finalResult = recognizer.getResult(stream)
            if (!finalResult.text.isNullOrEmpty()) {
                listener.onFinalResult(finalResult.text, processingTime, 0f)
                listener.onEndpoint()
            }
            recognizer.reset(stream)
            return STTProcessingResult("", processingTime, 0f)
        }

        return null
    }

    /**
     * Process accumulated speech buffer with offline recognizer
     * @param speechBuffer List of audio chunks from VAD
     * @return Processing result with text and timing, or null if failed
     */
    fun processOfflineBuffer(speechBuffer: List<FloatArray>): STTProcessingResult? {
        val recognizer = offlineRecognizer ?: return null

        val totalSamples = speechBuffer.sumOf { it.size }
        if (totalSamples == 0) {
            android.util.Log.w("STTRecognitionManager", "No audio to process")
            return null
        }

        val allSamples = FloatArray(totalSamples)
        var offset = 0
        speechBuffer.forEach { buffer ->
            buffer.copyInto(allSamples, offset)
            offset += buffer.size
        }

        android.util.Log.i("STTRecognitionManager", "Processing offline: ${allSamples.size} samples (${String.format("%.2f", allSamples.size.toFloat() / sampleRate)}s)")

        return processOfflineSamples(allSamples)
    }

    /**
     * Process samples through offline recognizer
     * @param samples Float audio samples
     * @return Processing result with text and timing, or null if failed
     */
    fun processOfflineSamples(samples: FloatArray): STTProcessingResult? {
        val recognizer = offlineRecognizer
        if (recognizer == null) {
            android.util.Log.e("STTRecognitionManager", "[STT DEBUG] processOfflineSamples: recognizer is NULL!")
            return null
        }

        // Calculate audio statistics
        val durationSec = samples.size.toFloat() / sampleRate
        var minSample = Float.MAX_VALUE
        var maxSample = Float.MIN_VALUE
        var sumAbsSamples = 0.0
        var zeroCount = 0
        for (sample in samples) {
            if (sample < minSample) minSample = sample
            if (sample > maxSample) maxSample = sample
            sumAbsSamples += kotlin.math.abs(sample.toDouble())
            if (sample == 0f) zeroCount++
        }
        val meanAbsSample = sumAbsSamples / samples.size
        val zeroPercent = (zeroCount.toFloat() / samples.size * 100)

        android.util.Log.i("STTRecognitionManager", "[STT DEBUG] processOfflineSamples called")
        android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Input: ${samples.size} samples = ${"%.2f".format(durationSec)}s @ ${sampleRate}Hz")
        android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Audio stats: min=${"%.5f".format(minSample)}, max=${"%.5f".format(maxSample)}, meanAbs=${"%.5f".format(meanAbsSample)}")
        android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Zero samples: $zeroCount (${"%.1f".format(zeroPercent)}%)")

        // Check for suspicious audio
        if (zeroPercent > 90) {
            android.util.Log.w("STTRecognitionManager", "[STT DEBUG] WARNING: >90% of samples are zero - audio may be silent/corrupt!")
        }
        if (meanAbsSample < 0.001) {
            android.util.Log.w("STTRecognitionManager", "[STT DEBUG] WARNING: Very low amplitude audio (meanAbs < 0.001) - may be too quiet!")
        }
        if (durationSec < 0.3) {
            android.util.Log.w("STTRecognitionManager", "[STT DEBUG] WARNING: Very short audio (<300ms) - may be too short for STT!")
        }

        return try {
            val startTime = System.currentTimeMillis()

            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Creating stream...")
            val stream = recognizer.createStream()

            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Accepting waveform...")
            stream.acceptWaveform(samples, sampleRate)

            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Decoding...")
            recognizer.decode(stream)

            val processingTime = System.currentTimeMillis() - startTime
            val rtf = processingTime.toFloat() / (durationSec * 1000)

            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Getting result...")
            val result = recognizer.getResult(stream)

            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Raw result text: \"${result.text}\"")
            android.util.Log.i("STTRecognitionManager", "[STT DEBUG] Processing time: ${processingTime}ms (RTF: ${"%.2f".format(rtf)}x)")

            stream.release()

            if (result.text.isNotEmpty()) {
                android.util.Log.i("STTRecognitionManager", "[STT DEBUG] SUCCESS: \"${result.text}\"")
                listener.onFinalResult(result.text, processingTime, 0f)
                STTProcessingResult(result.text, processingTime, 0f)
            } else {
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG] EMPTY RESULT - model returned no transcription!")
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG] Possible causes:")
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG]   - Audio too short or quiet")
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG]   - Speech not recognized by model")
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG]   - Audio language doesn't match model")
                android.util.Log.w("STTRecognitionManager", "[STT DEBUG]   - Heavy background noise")
                null
            }
        } catch (e: Exception) {
            android.util.Log.e("STTRecognitionManager", "[STT DEBUG] EXCEPTION during STT: ${e.message}", e)
            listener.onError("OFFLINE_STT_ERROR", "Failed to process speech: ${e.message}")
            null
        }
    }

    /**
     * Reset the streaming session
     */
    fun resetStream() {
        onlineStream?.let { stream ->
            onlineRecognizer?.reset(stream)
        }
    }

    /**
     * Flush any pending streaming result (for when stopping mid-speech)
     * @return The pending text, or null if none
     */
    fun flushStreamingResult(): String? {
        val recognizer = onlineRecognizer ?: return null
        val stream = onlineStream ?: return null

        try {
            // Decode any remaining audio in the stream
            while (recognizer.isReady(stream)) {
                recognizer.decode(stream)
            }

            val result = recognizer.getResult(stream)
            if (!result.text.isNullOrEmpty()) {
                android.util.Log.i("STTRecognitionManager", "[FLUSH] Flushing pending text: \"${result.text}\"")
                return result.text
            }
        } catch (e: Exception) {
            android.util.Log.e("STTRecognitionManager", "[FLUSH] Error flushing: ${e.message}")
        }

        return null
    }

    /**
     * Release streaming session
     */
    fun releaseStream() {
        onlineStream?.release()
        onlineStream = null
    }

    /**
     * Get current model type
     */
    fun getModelType(): String = modelType

    /**
     * Check if streaming mode
     */
    fun isStreaming(): Boolean = modelType == "streaming"

    /**
     * Check if offline mode
     */
    fun isOffline(): Boolean = modelType == "offline"

    /**
     * Check if recognizer is initialized
     */
    fun isInitialized(): Boolean = onlineRecognizer != null || offlineRecognizer != null

    /**
     * Release all resources
     */
    fun release() {
        releaseStream()
        onlineRecognizer?.release()
        onlineRecognizer = null
        offlineRecognizer?.release()
        offlineRecognizer = null
        android.util.Log.i("STTRecognitionManager", "STT released")
    }
}