package com.openclaw.alfred.voice

import android.content.Context
import android.util.Log
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.withContext
import org.json.JSONObject
import org.vosk.Model
import org.vosk.Recognizer
import org.vosk.android.RecognitionListener
import org.vosk.android.SpeechService
import java.io.File
import java.io.IOException

/**
 * Unified Vosk recognition manager for both wake word detection and full transcription.
 * 
 * Modes:
 * - WAKE_WORD: Listens for "hey alfred" or "alfred" (lightweight grammar)
 * - FULL_RECOGNITION: Transcribes full sentences (full vocabulary)
 */
class VoskRecognitionManager(
    private val context: Context,
    private val onWakeWordDetected: () -> Unit,
    private val onTranscriptionResult: (text: String, confidence: Float) -> Unit,
    private val onError: (String) -> Unit,
    private val onInitialized: () -> Unit = {}
) {
    
    private val TAG = "VoskRecognitionManager"
    private var model: Model? = null
    private var speechService: SpeechService? = null
    private var currentMode: RecognitionMode = RecognitionMode.STOPPED
    private var audioBuffer: MutableList<ByteArray> = mutableListOf()
    private var isRecordingAudio = false
    
    /**
     * Get wake words from preferences.
     */
    private fun getWakeWords(): Set<String> {
        val prefs = context.getSharedPreferences("alfred_settings", android.content.Context.MODE_PRIVATE)
        val customWord = prefs.getString("wake_word", "alfred") ?: "alfred"
        
        return setOf(
            customWord.lowercase().trim(),
            "hey ${customWord.lowercase().trim()}",
            "ok ${customWord.lowercase().trim()}"
        )
    }
    
    enum class RecognitionMode {
        STOPPED,
        WAKE_WORD,
        FULL_RECOGNITION
    }
    
    /**
     * Initialize the Vosk model (must be called before start).
     * Copies model from assets and loads it.
     */
    suspend fun initialize() {
        withContext(Dispatchers.IO) {
            try {
                Log.d(TAG, "Initializing Vosk model...")
                
                // Target directory in app's internal storage
                val modelDir = File(context.filesDir, "vosk-model")
                
                // Copy model from assets if not already there
                if (!modelDir.exists() || !File(modelDir, "am").exists()) {
                    Log.d(TAG, "Copying model from assets...")
                    copyModelFromAssets(modelDir)
                    Log.d(TAG, "Model copied successfully")
                }
                
                // Load the model
                Log.d(TAG, "Loading model from ${modelDir.absolutePath}")
                model = Model(modelDir.absolutePath)
                Log.d(TAG, "Model loaded successfully")
                
                // Notify success on main thread
                withContext(Dispatchers.Main) {
                    onInitialized()
                }
                
            } catch (e: Exception) {
                Log.e(TAG, "Failed to initialize model", e)
                withContext(Dispatchers.Main) {
                    onError("Recognition setup failed: ${e.message}")
                }
            }
        }
    }
    
    /**
     * Start wake word detection mode.
     */
    fun startWakeWordMode() {
        if (currentMode != RecognitionMode.STOPPED) {
            Log.w(TAG, "Already running in mode: $currentMode")
            return
        }
        
        val currentModel = model
        if (currentModel == null) {
            onError("Model not initialized. Call initialize() first.")
            return
        }
        
        try {
            Log.d(TAG, "=== Starting wake word mode ===")
            
            // Create recognizer for wake word detection (partial results only)
            val recognizer = Recognizer(currentModel, 16000.0f)
            recognizer.setMaxAlternatives(0)
            recognizer.setWords(false)
            
            Log.d(TAG, "Recognizer created for wake word, starting listener...")
            startListening(recognizer, RecognitionMode.WAKE_WORD)
            Log.d(TAG, "Wake word listener started")
            
        } catch (e: Exception) {
            Log.e(TAG, "Failed to start wake word mode", e)
            onError("Failed to start wake word detection: ${e.message}")
        }
    }
    
    /**
     * Start full recognition mode (transcribe full sentences).
     */
    fun startFullRecognitionMode() {
        if (currentMode != RecognitionMode.STOPPED) {
            stop()
        }
        
        val currentModel = model
        if (currentModel == null) {
            onError("Model not initialized. Call initialize() first.")
            return
        }
        
        try {
            Log.d(TAG, "=== Starting full recognition mode ===")
            
            // Create recognizer for full transcription
            val recognizer = Recognizer(currentModel, 16000.0f)
            recognizer.setMaxAlternatives(1)
            recognizer.setWords(true)
            
            Log.d(TAG, "Recognizer created for full transcription")
            
            // Increase timeout for longer phrases
            // Note: Vosk's internal timeout is ~10 seconds of silence by default
            
            // Start recording audio for potential Google fallback
            audioBuffer.clear()
            isRecordingAudio = true
            
            Log.d(TAG, "Starting full recognition listener...")
            startListening(recognizer, RecognitionMode.FULL_RECOGNITION)
            Log.d(TAG, "Full recognition listener started")
            
        } catch (e: Exception) {
            Log.e(TAG, "Failed to start full recognition mode", e)
            onError("Failed to start recognition: ${e.message}")
        }
    }
    
    /**
     * Start listening with the given recognizer and mode.
     */
    private fun startListening(recognizer: Recognizer, mode: RecognitionMode) {
        Log.d(TAG, "startListening called with mode: $mode")
        
        try {
            speechService = SpeechService(recognizer, 16000.0f)
            Log.d(TAG, "SpeechService created, about to start listening...")
        } catch (e: Exception) {
            Log.e(TAG, "Failed to create SpeechService", e)
            throw e
        }
        
        speechService?.startListening(object : RecognitionListener {
            override fun onPartialResult(hypothesis: String?) {
                Log.d(TAG, "onPartialResult: $hypothesis (mode: $mode)")
                hypothesis?.let { 
                    when (mode) {
                        RecognitionMode.WAKE_WORD -> checkForWakeWord(it)
                        RecognitionMode.FULL_RECOGNITION -> {
                            // Could show partial results in UI here
                            Log.d(TAG, "Partial recognition: $it")
                        }
                        RecognitionMode.STOPPED -> {}
                    }
                }
            }
            
            override fun onResult(hypothesis: String?) {
                Log.d(TAG, "onResult: $hypothesis (mode: $mode)")
                hypothesis?.let {
                    when (mode) {
                        RecognitionMode.WAKE_WORD -> checkForWakeWord(it)
                        RecognitionMode.FULL_RECOGNITION -> handleFullResult(it)
                        RecognitionMode.STOPPED -> {}
                    }
                }
            }
            
            override fun onFinalResult(hypothesis: String?) {
                Log.d(TAG, "onFinalResult: $hypothesis (mode: $mode)")
                if (mode == RecognitionMode.FULL_RECOGNITION) {
                    hypothesis?.let { handleFullResult(it) }
                    isRecordingAudio = false
                }
            }
            
            override fun onError(exception: Exception?) {
                Log.e(TAG, "Recognition error in mode $mode", exception)
                currentMode = RecognitionMode.STOPPED
                isRecordingAudio = false
                onError("Recognition error: ${exception?.message}")
            }
            
            override fun onTimeout() {
                Log.d(TAG, "Recognition timeout in mode: $mode")
                isRecordingAudio = false
                currentMode = RecognitionMode.STOPPED
                
                if (mode == RecognitionMode.WAKE_WORD) {
                    // Restart wake word detection automatically
                    startWakeWordMode()
                } else {
                    // Full recognition timeout - user might not have said anything yet
                    // Don't treat this as an error, just return to wake word mode
                    Log.w(TAG, "Full recognition timeout - no speech detected")
                    // Return to wake word mode instead of showing error
                    startWakeWordMode()
                }
            }
        })
        
        currentMode = mode
        Log.d(TAG, "Started listening in mode: $mode")
    }
    
    /**
     * Check if the hypothesis contains a wake word.
     */
    private fun checkForWakeWord(hypothesis: String) {
        try {
            val json = JSONObject(hypothesis)
            val text = json.optString("partial", json.optString("text", ""))
            
            if (text.isNotEmpty()) {
                val lowerText = text.trim().lowercase()
                
                // Check for wake words (get from preferences)
                val wakeWords = getWakeWords()
                for (wakeWord in wakeWords) {
                    if (lowerText.contains(wakeWord)) {
                        Log.i(TAG, "Wake word detected: $wakeWord")
                        stop()
                        onWakeWordDetected()
                        return
                    }
                }
            }
        } catch (e: Exception) {
            Log.e(TAG, "Error parsing hypothesis: $hypothesis", e)
        }
    }
    
    /**
     * Handle full recognition result.
     */
    private fun handleFullResult(hypothesis: String) {
        try {
            Log.d(TAG, "handleFullResult called with: $hypothesis")
            val json = JSONObject(hypothesis)
            
            // Vosk returns: { "alternatives": [{ "text": "...", "confidence": ... }] }
            // NOT: { "text": "..." } at top level
            var text = ""
            var confidence = 1.0f
            
            val alternatives = json.optJSONArray("alternatives")
            if (alternatives != null && alternatives.length() > 0) {
                val firstAlt = alternatives.getJSONObject(0)
                text = firstAlt.optString("text", "").trim()
                confidence = firstAlt.optDouble("confidence", 1.0).toFloat()
                
                // Normalize confidence (Vosk gives raw log-likelihood scores)
                // Typical range: 0-500+, normalize to 0-1
                if (confidence > 1.0f) {
                    confidence = Math.min(confidence / 500.0f, 1.0f)
                }
            }
            
            if (text.isNotEmpty()) {
                Log.d(TAG, "Full result: '$text' (confidence: $confidence)")
                stop()
                onTranscriptionResult(text, confidence)
            } else {
                Log.w(TAG, "Empty transcription result, hypothesis: $hypothesis")
                // Don't show error - just silently return to wake word mode
                // User might have just paused or not said anything
                stop()
                // Don't call onError - just go back to wake word mode
                // The onTimeout handler will take care of this
            }
            
        } catch (e: Exception) {
            Log.e(TAG, "Error parsing full result: $hypothesis", e)
            stop()
            onError("Failed to parse result")
        }
    }
    
    /**
     * Stop listening.
     */
    fun stop() {
        if (currentMode == RecognitionMode.STOPPED) {
            return
        }
        
        try {
            speechService?.stop()
            speechService?.shutdown()
            speechService = null
            currentMode = RecognitionMode.STOPPED
            isRecordingAudio = false
            Log.d(TAG, "Stopped listening")
        } catch (e: Exception) {
            Log.e(TAG, "Error stopping recognition", e)
        }
    }
    
    /**
     * Get saved audio buffer for Google fallback.
     */
    fun getSavedAudioBuffer(): List<ByteArray> {
        return audioBuffer.toList()
    }
    
    /**
     * Save audio buffer to temporary file for Google fallback.
     */
    fun saveAudioToFile(): File? {
        if (audioBuffer.isEmpty()) return null
        
        try {
            val tempFile = File(context.cacheDir, "vosk_audio_${System.currentTimeMillis()}.wav")
            // TODO: Write WAV header + audio data
            // For now, return null - Google fallback via re-recording
            return null
        } catch (e: Exception) {
            Log.e(TAG, "Failed to save audio", e)
            return null
        }
    }
    
    /**
     * Copy model from assets to internal storage.
     */
    private fun copyModelFromAssets(targetDir: File) {
        targetDir.mkdirs()
        
        val dirs = listOf("am", "conf", "graph", "ivector")
        
        for (dir in dirs) {
            val assetPath = "vosk-model/$dir"
            val targetSubDir = File(targetDir, dir)
            targetSubDir.mkdirs()
            
            val files = context.assets.list(assetPath) ?: continue
            for (file in files) {
                val assetFilePath = "$assetPath/$file"
                val targetFile = File(targetSubDir, file)
                
                val subFiles = context.assets.list(assetFilePath)
                if (subFiles != null && subFiles.isNotEmpty()) {
                    targetFile.mkdirs()
                    for (subFile in subFiles) {
                        copyAssetFile("$assetFilePath/$subFile", File(targetFile, subFile))
                    }
                } else {
                    copyAssetFile(assetFilePath, targetFile)
                }
            }
        }
        
        try {
            copyAssetFile("vosk-model/README", File(targetDir, "README"))
        } catch (e: Exception) {
            // README is optional
        }
    }
    
    /**
     * Copy a single file from assets.
     */
    private fun copyAssetFile(assetPath: String, targetFile: File) {
        context.assets.open(assetPath).use { input ->
            targetFile.outputStream().use { output ->
                input.copyTo(output)
            }
        }
    }
    
    /**
     * Cleanup resources.
     */
    fun destroy() {
        stop()
        model?.close()
        model = null
        audioBuffer.clear()
    }
    
    /**
     * Get current mode.
     */
    fun getCurrentMode(): RecognitionMode = currentMode
    
    /**
     * Check if model is initialized.
     */
    fun isInitialized(): Boolean = model != null
}
