package com.sherpaonnxofflinestt.managers

import com.k2fsa.sherpa.onnx.SpeakerEmbeddingExtractor
import com.k2fsa.sherpa.onnx.SpeakerEmbeddingExtractorConfig
import com.sherpaonnxofflinestt.STTConstants
import com.sherpaonnxofflinestt.models.DiarizationConfig
import com.sherpaonnxofflinestt.models.SpeakerProfile
import java.io.File
import java.util.Collections
import java.util.concurrent.CopyOnWriteArrayList
import java.util.concurrent.locks.ReentrantLock
import kotlin.concurrent.withLock

/**
 * Manages speaker diarization using embedding extraction and cosine similarity.
 * Implements dual threshold logic for better accuracy:
 * - speakerThreshold (lenient): Used for speaker assignment
 * - embeddingThreshold (strict): Used for profile updates to prevent degradation
 *
 * NOTE: This logic is duplicated in iOS (SpeakerDiarizationManager.swift).
 * Key functions that must stay in sync:
 * - normalizeEmbedding(): L2 normalization to unit length
 * - cosineSimilarity(): Dot product for normalized vectors
 * - averageEmbeddings(): Average multiple embeddings with re-normalization
 * - validateEmbedding(): Quality scoring with variance check
 * - matchOrCreateSpeaker(): Dual threshold speaker assignment logic
 *
 * Any changes here MUST be mirrored in iOS to maintain platform parity.
 */
class DiarizationManager(
    private val eventManager: EventEmissionManager
) {
    // Thread-safe extractor access
    @Volatile
    private var speakerExtractor: SpeakerEmbeddingExtractor? = null

    // Thread-safe collections for concurrent access
    private val speakerProfiles = Collections.synchronizedList(mutableListOf<SpeakerProfile>())
    private val audioBuffer = CopyOnWriteArrayList<FloatArray>()

    // Locks for critical sections
    private val embeddingLock = ReentrantLock()
    private val stateLock = ReentrantLock()

    @Volatile
    private var currentSpeakerId: Int = -1

    @Volatile
    private var nextSpeakerId: Int = 1

    @Volatile
    var isEnabled: Boolean = false
        private set

    // Configuration
    private var config: DiarizationConfig = DiarizationConfig()

    // Last match metrics for event emission
    @Volatile
    private var lastConfidence: Float = 0f
    @Volatile
    private var lastEmbeddingQuality: Float = 0f

    // Onboarding mode: first N segments always go to Speaker 1
    @Volatile
    private var onboardingRemaining: Int = 0

    // New speaker buffering: accumulate evidence before creating new speaker
    private val candidateEmbeddings = Collections.synchronizedList(mutableListOf<FloatArray>())
    @Volatile
    private var candidateSpeechMs: Long = 0
    @Volatile
    private var lastMatchedSpeakerId: Int = -1

    // Sub-segment diarization settings (using constants)
    private val windowSizeMs: Int = STTConstants.DIARIZATION_WINDOW_SIZE_MS
    private val windowStepMs: Int = STTConstants.DIARIZATION_WINDOW_STEP_MS

    // Last detected speakers in segment (for multi-speaker reporting)
    @Volatile
    private var lastSegmentSpeakers: List<Int> = emptyList()

    /**
     * Initialize the speaker embedding extractor
     * @param modelPath Path to speaker embedding model
     * @param diarizationConfig Configuration for diarization
     * @return true if initialization successful
     */
    fun initialize(
        modelPath: String,
        diarizationConfig: DiarizationConfig = DiarizationConfig()
    ): Boolean {
        android.util.Log.i("DiarizationManager", "[DIAR INIT] Starting initialization with model: $modelPath")

        if (modelPath.isEmpty()) {
            android.util.Log.w("DiarizationManager", "[DIAR INIT] Model path is empty!")
            return false
        }

        val modelFile = File(modelPath)
        if (!modelFile.exists()) {
            android.util.Log.w("DiarizationManager", "[DIAR INIT] Model file does not exist: $modelPath")
            return false
        }

        android.util.Log.i("DiarizationManager", "[DIAR INIT] Model file exists, size: ${modelFile.length()} bytes")

        return try {
            this.config = diarizationConfig

            android.util.Log.i("DiarizationManager", "[DIAR INIT] Creating SpeakerEmbeddingExtractorConfig...")
            val extractorConfig = SpeakerEmbeddingExtractorConfig(
                model = modelPath,
                numThreads = 2,
                debug = false,
                provider = "cpu"
            )

            android.util.Log.i("DiarizationManager", "[DIAR INIT] Creating SpeakerEmbeddingExtractor...")
            speakerExtractor = SpeakerEmbeddingExtractor(assetManager = null, config = extractorConfig)

            if (speakerExtractor == null) {
                android.util.Log.e("DiarizationManager", "[DIAR INIT] SpeakerEmbeddingExtractor is NULL after creation!")
                isEnabled = false
                return false
            }

            isEnabled = true
            android.util.Log.i("DiarizationManager",
                "[DIAR INIT] SUCCESS! Speaker diarization initialized (speakerThreshold=${config.speakerThreshold}, " +
                "embeddingThreshold=${config.embeddingThreshold}, maxSpeakers=${config.maxSpeakers})")
            true
        } catch (e: Throwable) {
            // Catch Throwable to also catch Error (e.g., UnsatisfiedLinkError)
            android.util.Log.e("DiarizationManager", "[DIAR INIT] FAILED: ${e.javaClass.name}: ${e.message}", e)
            isEnabled = false
            false
        }
    }

    /**
     * Validate embedding quality.
     * Returns quality score 0-1, or -1 if invalid.
     * Note: Embeddings may not be unit-normalized depending on the model.
     */
    private fun validateEmbedding(embedding: FloatArray): Float {
        // Check for empty embedding
        if (embedding.isEmpty()) {
            android.util.Log.w("DiarizationManager", "Empty embedding")
            return -1f
        }

        // Calculate norm and check for invalid values
        var norm = 0f
        var hasValidValues = false
        for (value in embedding) {
            // Check for NaN or Infinite values
            if (value.isNaN() || value.isInfinite()) {
                android.util.Log.w("DiarizationManager", "Embedding contains NaN or Infinite values")
                return -1f
            }
            norm += value * value
            if (value != 0f) hasValidValues = true
        }
        norm = kotlin.math.sqrt(norm)

        // Check for near-zero norm (silent/invalid audio)
        if (norm < 0.001f || !hasValidValues) {
            android.util.Log.w("DiarizationManager", "Embedding norm too small: $norm")
            return -1f
        }

        // Quality score based on variance of embedding values
        // Higher variance = more distinctive = higher quality
        val mean = embedding.sum() / embedding.size
        var variance = 0f
        for (value in embedding) {
            variance += (value - mean) * (value - mean)
        }
        variance /= embedding.size

        // Normalize quality to 0-1 range (typical variance is 0.01-0.1 for good embeddings)
        val quality = kotlin.math.min(variance * 10f, 1.0f)

        android.util.Log.d("DiarizationManager", "Embedding norm=${String.format("%.2f", norm)}, variance=${String.format("%.4f", variance)}, quality=${String.format("%.3f", quality)}")
        return quality
    }

    /**
     * Add audio chunk to the diarization buffer.
     * Thread-safe: uses CopyOnWriteArrayList for concurrent access.
     * Call this during speech segments.
     */
    @Synchronized
    fun addAudioChunk(samples: FloatArray) {
        if (isEnabled) {
            audioBuffer.add(samples.copyOf())
            // Log periodically (every 10 chunks)
            if (audioBuffer.size % 10 == 0) {
                val totalSamples = audioBuffer.sumOf { it.size }
                android.util.Log.d("DiarizationManager", "[DIAR DEBUG] addAudioChunk: buffer now has ${audioBuffer.size} chunks, ${totalSamples} samples")
            }
        } else {
            // Log only occasionally to avoid spam (every 50 chunks)
            if (audioBuffer.size % 50 == 0) {
                android.util.Log.w("DiarizationManager", "[DIAR DEBUG] addAudioChunk: isEnabled=false, not buffering (${audioBuffer.size} discarded)")
            }
        }
    }

    /**
     * Clear the audio buffer (call on speech start).
     * Thread-safe: synchronized method.
     */
    @Synchronized
    fun clearBuffer() {
        val previousSize = audioBuffer.size
        audioBuffer.clear()
        if (previousSize > 0) {
            android.util.Log.d("DiarizationManager", "[DIAR DEBUG] clearBuffer: cleared $previousSize chunks")
        }
    }

    /**
     * Process accumulated audio buffer for speaker identification.
     * Thread-safe: uses embeddingLock for embedding operations.
     * Uses sub-segment processing to detect multiple speakers within a single VAD segment.
     * @param sampleRate Audio sample rate
     * @return Primary speaker ID if identified, -1 otherwise
     */
    @Synchronized
    fun processBuffer(sampleRate: Int): Int {
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] processBuffer called, buffer has ${audioBuffer.size} chunks")

        if (!isEnabled) {
            android.util.Log.w("DiarizationManager", "[DIAR DEBUG] processBuffer: isEnabled=false, skipping")
            return -1
        }
        val extractor = speakerExtractor
        if (extractor == null) {
            android.util.Log.w("DiarizationManager", "[DIAR DEBUG] processBuffer: speakerExtractor is NULL, skipping")
            return -1
        }

        val totalSamples = audioBuffer.sumOf { it.size }
        val minSamples = (config.minSpeechDurationMs * sampleRate) / 1000

        val totalDurationMs = (totalSamples.toLong() * 1000) / sampleRate
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Processing: ${audioBuffer.size} chunks, ${totalSamples} samples (${totalDurationMs}ms), minSamples=$minSamples")

        if (totalSamples < minSamples) {
            android.util.Log.d("DiarizationManager", "Audio too short, skipping")
            audioBuffer.clear()
            return -1
        }

        // Combine all audio chunks - create a snapshot of the buffer
        val bufferSnapshot = audioBuffer.toList()
        val allSamples = FloatArray(totalSamples)
        var offset = 0
        bufferSnapshot.forEach { buffer ->
            buffer.copyInto(allSamples, offset)
            offset += buffer.size
        }
        audioBuffer.clear()

        // Calculate window sizes in samples
        val windowSizeSamples = (windowSizeMs * sampleRate) / 1000
        val windowStepSamples = (windowStepMs * sampleRate) / 1000

        // If audio is short enough, process as single segment (original behavior)
        if (totalSamples <= windowSizeSamples * 1.5) {
            android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Short segment - single speaker detection")
            return processSingleSegment(allSamples, sampleRate, totalDurationMs)
        }

        // Sub-segment processing for longer audio
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Long segment (${totalDurationMs}ms) - multi-speaker detection with ${windowSizeMs}ms windows")

        val windowSpeakers = mutableListOf<Pair<Int, Int>>()  // (windowStartMs, speakerId)
        var windowStart = 0

        while (windowStart + windowSizeSamples <= totalSamples) {
            val windowEnd = windowStart + windowSizeSamples
            val windowSamples = allSamples.copyOfRange(windowStart, windowEnd)
            val windowStartMs = (windowStart.toLong() * 1000) / sampleRate

            val speakerId = processWindow(windowSamples, sampleRate)
            if (speakerId > 0) {
                windowSpeakers.add(Pair(windowStartMs.toInt(), speakerId))
                android.util.Log.d("DiarizationManager", "[DIAR DEBUG] Window @${windowStartMs}ms -> Speaker $speakerId")
            }

            windowStart += windowStepSamples
        }

        // Process any remaining audio at the end
        if (windowStart < totalSamples && totalSamples - windowStart >= minSamples) {
            val remainingSamples = allSamples.copyOfRange(windowStart, totalSamples)
            val windowStartMs = (windowStart.toLong() * 1000) / sampleRate
            val speakerId = processWindow(remainingSamples, sampleRate)
            if (speakerId > 0) {
                windowSpeakers.add(Pair(windowStartMs.toInt(), speakerId))
                android.util.Log.d("DiarizationManager", "[DIAR DEBUG] Final window @${windowStartMs}ms -> Speaker $speakerId")
            }
        }

        // Analyze results
        if (windowSpeakers.isEmpty()) {
            android.util.Log.w("DiarizationManager", "[DIAR DEBUG] No speakers detected in any window")
            return -1
        }

        // Get unique speakers and their counts
        val speakerCounts = windowSpeakers.groupBy { it.second }.mapValues { it.value.size }
        val uniqueSpeakers = speakerCounts.keys.toList().sorted()
        lastSegmentSpeakers = uniqueSpeakers

        // Find dominant speaker (most windows)
        val dominantSpeaker = speakerCounts.maxByOrNull { it.value }?.key ?: -1
        val firstSpeaker = windowSpeakers.firstOrNull()?.second ?: dominantSpeaker

        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] ===== MULTI-SPEAKER RESULT =====")
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Detected ${uniqueSpeakers.size} speaker(s): $uniqueSpeakers")
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Speaker counts: $speakerCounts")
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] First speaker: $firstSpeaker, Dominant: $dominantSpeaker")

        // Report speaker changes
        if (uniqueSpeakers.size > 1) {
            android.util.Log.i("DiarizationManager", "[DIAR DEBUG] !!! MULTIPLE SPEAKERS DETECTED IN SEGMENT !!!")
            // Emit event for multi-speaker detection
            eventManager.sendMultiSpeakerDetected(uniqueSpeakers, windowSpeakers.map { it.first to it.second })
        }

        currentSpeakerId = firstSpeaker
        return firstSpeaker
    }

    /**
     * Process a single audio window and return matched speaker ID.
     * Thread-safe with proper resource cleanup using try-finally.
     */
    private fun processWindow(samples: FloatArray, sampleRate: Int): Int {
        val extractor = speakerExtractor ?: return -1

        return embeddingLock.withLock {
            var stream: com.k2fsa.sherpa.onnx.OnlineStream? = null
            try {
                stream = extractor.createStream()
                stream.acceptWaveform(samples, sampleRate)

                if (!extractor.isReady(stream)) {
                    return@withLock -1
                }

                val embedding = extractor.compute(stream)

                val quality = validateEmbedding(embedding)
                if (quality < 0) return@withLock -1

                // Match against existing speakers (don't create new ones in window processing)
                matchSpeakerOnly(embedding)
            } catch (e: Exception) {
                android.util.Log.w("DiarizationManager", "Window processing error: ${e.message}")
                -1
            } finally {
                stream?.release()
            }
        }
    }

    /**
     * Match embedding against existing speakers without creating new ones
     */
    private fun matchSpeakerOnly(embedding: FloatArray): Int {
        val normalizedEmbedding = normalizeEmbedding(embedding)

        var bestMatch = -1
        var bestSimilarity = 0f

        speakerProfiles.forEach { profile ->
            val similarity = cosineSimilarity(normalizedEmbedding, profile.embedding)
            if (similarity > bestSimilarity && similarity >= config.speakerThreshold) {
                bestSimilarity = similarity
                bestMatch = profile.speakerId
            }
        }

        return if (bestMatch > 0) bestMatch else lastMatchedSpeakerId
    }

    /**
     * Process short segment as single speaker (original behavior).
     * Thread-safe with proper resource cleanup using try-finally.
     */
    private fun processSingleSegment(allSamples: FloatArray, sampleRate: Int, speechDurationMs: Long): Int {
        val extractor = speakerExtractor ?: return -1

        return embeddingLock.withLock {
            var stream: com.k2fsa.sherpa.onnx.OnlineStream? = null
            try {
                stream = extractor.createStream()
                stream.acceptWaveform(allSamples, sampleRate)

                if (!extractor.isReady(stream)) {
                    return@withLock -1
                }

                val embedding = extractor.compute(stream)

                val quality = validateEmbedding(embedding)
                lastEmbeddingQuality = if (quality >= 0) quality else 0f

                if (quality < 0) return@withLock -1

                val (speakerId, _) = matchOrCreateSpeaker(embedding, speechDurationMs)
                currentSpeakerId = speakerId
                lastSegmentSpeakers = listOf(speakerId)
                speakerId
            } catch (e: Exception) {
                android.util.Log.e("DiarizationManager", "Error: ${e.message}", e)
                -1
            } finally {
                stream?.release()
            }
        }
    }

    /**
     * Match embedding against existing speakers or create new speaker.
     * Thread-safe: operates under stateLock.
     * Uses dual threshold logic:
     * - speakerThreshold for matching (lenient)
     * - embeddingThreshold for profile update (strict)
     *
     * Also implements:
     * - Onboarding mode: First N segments always go to Speaker 1
     * - New speaker buffering: Require X seconds of unmatched speech before creating new speaker
     *
     * @return Pair of (speakerId, justConfirmed)
     */
    private fun matchOrCreateSpeaker(embedding: FloatArray, speechDurationMs: Long): Pair<Int, Boolean> = stateLock.withLock {
        // Normalize the embedding for comparison
        val normalizedEmbedding = normalizeEmbedding(embedding)

        // Case 1: No speakers exist yet - create Speaker 1 and start onboarding
        if (speakerProfiles.isEmpty()) {
            val newId = nextSpeakerId++
            speakerProfiles.add(SpeakerProfile(
                speakerId = newId,
                embedding = normalizedEmbedding.copyOf(),
                numSegments = 1,
                status = "pending",
                totalSpeechDurationMs = speechDurationMs,
                lastSeenTimestamp = System.currentTimeMillis()
            ))

            // Initialize onboarding counter (remaining = total - 1 because we just used one)
            onboardingRemaining = config.onboardingSegments - 1
            lastMatchedSpeakerId = newId
            lastConfidence = 1.0f

            android.util.Log.i("DiarizationManager", "Created Speaker $newId, onboarding started (${onboardingRemaining} remaining)")
            eventManager.sendSpeakerUpdate(newId, "pending", false, speakerProfiles.size, 1.0f, lastEmbeddingQuality)
            return Pair(newId, false)
        }

        // Case 2: Onboarding mode - always assign to Speaker 1
        if (onboardingRemaining > 0) {
            val speaker1 = speakerProfiles.firstOrNull()
            if (speaker1 != null) {
                speaker1.numSegments++
                speaker1.totalSpeechDurationMs += speechDurationMs
                speaker1.lastSeenTimestamp = System.currentTimeMillis()

                // Update embedding with running average during onboarding
                val alpha = config.embeddingAverageWeight
                for (i in speaker1.embedding.indices) {
                    speaker1.embedding[i] = (1 - alpha) * speaker1.embedding[i] + alpha * normalizedEmbedding[i]
                }
                // Re-normalize after averaging
                val updatedNorm = normalizeEmbedding(speaker1.embedding)
                for (i in speaker1.embedding.indices) {
                    speaker1.embedding[i] = updatedNorm[i]
                }

                onboardingRemaining--
                lastMatchedSpeakerId = speaker1.speakerId
                lastConfidence = 1.0f

                val justConfirmed = speaker1.numSegments == 2
                if (justConfirmed) {
                    speaker1.status = "confirmed"
                }

                android.util.Log.i("DiarizationManager", "Onboarding: assigned to Speaker ${speaker1.speakerId} (${onboardingRemaining} remaining)")
                eventManager.sendSpeakerUpdate(
                    speaker1.speakerId,
                    speaker1.status,
                    justConfirmed,
                    speakerProfiles.size,
                    lastConfidence,
                    lastEmbeddingQuality
                )
                return Pair(speaker1.speakerId, justConfirmed)
            }
        }

        // Case 3: Normal matching logic
        var bestMatch = -1
        var bestSimilarity = 0f
        var bestSpeakerId = -1

        speakerProfiles.forEachIndexed { index, profile ->
            val similarity = cosineSimilarity(normalizedEmbedding, profile.embedding)
            android.util.Log.i("DiarizationManager", "[DIAR DEBUG] vs Speaker ${profile.speakerId}: similarity=${String.format("%.3f", similarity)} (threshold=${config.speakerThreshold})")
            if (similarity > bestSimilarity) {
                bestSimilarity = similarity
                bestSpeakerId = profile.speakerId
                if (similarity >= config.speakerThreshold) {
                    bestMatch = index
                }
            }
        }

        lastConfidence = bestSimilarity

        val matchStatus = if (bestMatch >= 0) "MATCHED Speaker $bestSpeakerId" else "NO MATCH (best=$bestSpeakerId at ${String.format("%.3f", bestSimilarity)}, need >=${config.speakerThreshold})"
        android.util.Log.i("DiarizationManager", "[DIAR DEBUG] Result: $matchStatus")

        return if (bestMatch >= 0) {
            // Match found - update existing speaker
            val profile = speakerProfiles[bestMatch]
            profile.numSegments++
            profile.totalSpeechDurationMs += speechDurationMs
            profile.lastSeenTimestamp = System.currentTimeMillis()

            // Clear candidate buffer - we matched an existing speaker
            candidateEmbeddings.clear()
            candidateSpeechMs = 0
            lastMatchedSpeakerId = profile.speakerId

            // Only update embedding if similarity exceeds stricter threshold
            // AND speech duration is sufficient
            val shouldUpdateEmbedding = bestSimilarity >= config.embeddingThreshold
                && speechDurationMs >= config.minEmbeddingUpdateMs

            if (shouldUpdateEmbedding) {
                // Update embedding with weighted running average (using normalized embedding)
                val alpha = config.embeddingAverageWeight
                for (i in profile.embedding.indices) {
                    profile.embedding[i] = (1 - alpha) * profile.embedding[i] + alpha * normalizedEmbedding[i]
                }
                // Re-normalize after averaging
                val updatedNorm = normalizeEmbedding(profile.embedding)
                for (i in profile.embedding.indices) {
                    profile.embedding[i] = updatedNorm[i]
                }
                android.util.Log.d("DiarizationManager", "Updated embedding for speaker ${profile.speakerId}")
            } else {
                android.util.Log.d("DiarizationManager",
                    "Skipped embedding update (similarity=${String.format("%.3f", bestSimilarity)}, " +
                    "threshold=${config.embeddingThreshold}, duration=${speechDurationMs}ms)")
            }

            val justConfirmed = profile.numSegments == 2
            if (justConfirmed) {
                profile.status = "confirmed"
            }

            eventManager.sendSpeakerUpdate(
                profile.speakerId,
                profile.status,
                justConfirmed,
                speakerProfiles.size,
                lastConfidence,
                lastEmbeddingQuality
            )
            Pair(profile.speakerId, justConfirmed)
        } else {
            // No match found - buffer candidate for new speaker
            candidateEmbeddings.add(normalizedEmbedding.copyOf())
            candidateSpeechMs += speechDurationMs

            android.util.Log.d("DiarizationManager",
                "Buffering candidate: ${candidateSpeechMs}ms / ${config.minSpeechForNewSpeaker}ms required")

            if (candidateSpeechMs >= config.minSpeechForNewSpeaker) {
                // Enough evidence accumulated - create new speaker from averaged embeddings
                val averagedEmbedding = averageEmbeddings(candidateEmbeddings)
                candidateEmbeddings.clear()
                candidateSpeechMs = 0

                // Check if we need to evict a speaker (LRU)
                if (speakerProfiles.size >= config.maxSpeakers) {
                    evictLRUSpeaker()
                }

                // Create new speaker with averaged embedding
                val newId = nextSpeakerId++
                speakerProfiles.add(SpeakerProfile(
                    speakerId = newId,
                    embedding = averagedEmbedding,
                    numSegments = 1,
                    status = "pending",
                    totalSpeechDurationMs = speechDurationMs,
                    lastSeenTimestamp = System.currentTimeMillis()
                ))

                lastMatchedSpeakerId = newId

                android.util.Log.i("DiarizationManager", "Created Speaker $newId from buffered segments (total: ${speakerProfiles.size})")
                eventManager.sendSpeakerUpdate(newId, "pending", false, speakerProfiles.size, 0f, lastEmbeddingQuality)
                Pair(newId, false)
            } else {
                // Not enough evidence yet - return last matched speaker
                val fallbackSpeakerId = if (lastMatchedSpeakerId > 0) lastMatchedSpeakerId else speakerProfiles.firstOrNull()?.speakerId ?: -1
                android.util.Log.d("DiarizationManager",
                    "Returning fallback speaker $fallbackSpeakerId while buffering")
                Pair(fallbackSpeakerId, false)
            }
        }
    }

    /**
     * Average multiple embeddings into one
     */
    private fun averageEmbeddings(embeddings: List<FloatArray>): FloatArray {
        if (embeddings.isEmpty()) return FloatArray(0)
        if (embeddings.size == 1) return embeddings[0].copyOf()

        val dim = embeddings[0].size
        val result = FloatArray(dim)

        for (embedding in embeddings) {
            for (i in 0 until dim) {
                result[i] += embedding[i]
            }
        }

        val count = embeddings.size.toFloat()
        for (i in 0 until dim) {
            result[i] /= count
        }

        // Normalize the averaged embedding
        return normalizeEmbedding(result)
    }

    /**
     * Evict the least recently used speaker to make room for a new one
     */
    private fun evictLRUSpeaker() {
        if (speakerProfiles.isEmpty()) return

        // Find speaker with oldest lastSeenTimestamp that is NOT the current speaker
        var lruIndex = -1
        var oldestTimestamp = Long.MAX_VALUE

        speakerProfiles.forEachIndexed { index, profile ->
            if (profile.speakerId != currentSpeakerId && profile.lastSeenTimestamp < oldestTimestamp) {
                oldestTimestamp = profile.lastSeenTimestamp
                lruIndex = index
            }
        }

        if (lruIndex >= 0) {
            val evicted = speakerProfiles.removeAt(lruIndex)
            android.util.Log.i("DiarizationManager",
                "Evicted speaker ${evicted.speakerId} (last seen: ${System.currentTimeMillis() - oldestTimestamp}ms ago)")
        }
    }

    /**
     * L2 normalize an embedding to unit length
     */
    private fun normalizeEmbedding(embedding: FloatArray): FloatArray {
        var norm = 0f
        for (value in embedding) {
            norm += value * value
        }
        norm = kotlin.math.sqrt(norm)

        if (norm < 0.001f) return embedding

        val normalized = FloatArray(embedding.size)
        for (i in embedding.indices) {
            normalized[i] = embedding[i] / norm
        }
        return normalized
    }

    /**
     * Calculate cosine similarity between two L2-normalized embeddings
     * For normalized vectors, this is simply the dot product
     */
    private fun cosineSimilarity(a: FloatArray, b: FloatArray): Float {
        if (a.size != b.size) return 0f

        var dot = 0f
        for (i in a.indices) {
            dot += a[i] * b[i]
        }
        return dot
    }

    /**
     * Get current speaker ID
     */
    fun getCurrentSpeakerId(): Int = currentSpeakerId

    /**
     * Get total speaker count
     */
    fun getSpeakerCount(): Int = speakerProfiles.size

    /**
     * Get last match confidence
     */
    fun getLastConfidence(): Float = lastConfidence

    /**
     * Get last embedding quality
     */
    fun getLastEmbeddingQuality(): Float = lastEmbeddingQuality

    /**
     * Get speakers detected in the last segment (for multi-speaker reporting)
     */
    fun getLastSegmentSpeakers(): List<Int> = lastSegmentSpeakers

    /**
     * Reset all speaker profiles.
     * Thread-safe: synchronized method.
     */
    @Synchronized
    fun resetSpeakers() {
        speakerProfiles.clear()
        currentSpeakerId = -1
        nextSpeakerId = 1
        audioBuffer.clear()
        lastConfidence = 0f
        lastEmbeddingQuality = 0f
        onboardingRemaining = 0
        candidateEmbeddings.clear()
        candidateSpeechMs = 0
        lastMatchedSpeakerId = -1
        android.util.Log.i("DiarizationManager", "Speakers reset")
    }

    // =====================
    // Management APIs
    // =====================

    /**
     * Merge source speaker into target speaker.
     * Thread-safe: synchronized method.
     * Combines embeddings (weighted by segment count) and removes source.
     * @return true if successful
     */
    @Synchronized
    fun mergeSpeakers(sourceId: Int, targetId: Int): Boolean {
        val sourceIndex = speakerProfiles.indexOfFirst { it.speakerId == sourceId }
        val targetIndex = speakerProfiles.indexOfFirst { it.speakerId == targetId }

        if (sourceIndex < 0 || targetIndex < 0) {
            android.util.Log.w("DiarizationManager", "mergeSpeakers failed: source=$sourceId ($sourceIndex), target=$targetId ($targetIndex)")
            return false
        }

        val source = speakerProfiles[sourceIndex]
        val target = speakerProfiles[targetIndex]

        // Average embeddings (weighted by segment count)
        val totalSegments = source.numSegments + target.numSegments
        val sourceWeight = source.numSegments.toFloat() / totalSegments
        val targetWeight = target.numSegments.toFloat() / totalSegments

        val mergedEmbedding = FloatArray(target.embedding.size)
        for (i in target.embedding.indices) {
            mergedEmbedding[i] = targetWeight * target.embedding[i] + sourceWeight * source.embedding[i]
        }

        // Normalize result
        val normalized = normalizeEmbedding(mergedEmbedding)
        normalized.copyInto(target.embedding)

        // Merge stats
        target.numSegments = totalSegments
        target.totalSpeechDurationMs += source.totalSpeechDurationMs
        if (source.lastSeenTimestamp > target.lastSeenTimestamp) {
            target.lastSeenTimestamp = source.lastSeenTimestamp
        }

        // Remove source
        speakerProfiles.removeAt(sourceIndex)

        // Update current speaker if needed
        if (currentSpeakerId == sourceId) {
            currentSpeakerId = targetId
        }
        if (lastMatchedSpeakerId == sourceId) {
            lastMatchedSpeakerId = targetId
        }

        // Clear candidate buffers - context has changed after merge
        candidateEmbeddings.clear()
        candidateSpeechMs = 0

        // Reset onboarding to help stabilize merged profile
        onboardingRemaining = config.onboardingSegments

        android.util.Log.i("DiarizationManager", "Merged Speaker $sourceId into Speaker $targetId (total: ${speakerProfiles.size}), reset onboarding")
        return true
    }

    /**
     * Remove a speaker profile.
     * Thread-safe: synchronized method.
     * @return true if speaker existed and was removed
     */
    @Synchronized
    fun removeSpeaker(speakerId: Int): Boolean {
        val index = speakerProfiles.indexOfFirst { it.speakerId == speakerId }
        if (index < 0) {
            android.util.Log.w("DiarizationManager", "removeSpeaker failed: speaker $speakerId not found")
            return false
        }

        speakerProfiles.removeAt(index)

        if (currentSpeakerId == speakerId) {
            currentSpeakerId = speakerProfiles.firstOrNull()?.speakerId ?: -1
        }
        if (lastMatchedSpeakerId == speakerId) {
            lastMatchedSpeakerId = speakerProfiles.firstOrNull()?.speakerId ?: -1
        }

        android.util.Log.i("DiarizationManager", "Removed Speaker $speakerId (total: ${speakerProfiles.size})")
        return true
    }

    /**
     * Get all speaker profiles info.
     * Thread-safe: creates a snapshot of profiles under synchronization.
     * @return List of speaker info maps
     */
    @Synchronized
    fun getSpeakerProfiles(): List<Map<String, Any>> {
        // Create a snapshot to avoid concurrent modification
        return speakerProfiles.toList().map { profile ->
            mapOf(
                "speakerId" to profile.speakerId,
                "status" to profile.status,
                "segmentCount" to profile.numSegments,
                "totalSpeechMs" to profile.totalSpeechDurationMs,
                "lastSeenTimestamp" to profile.lastSeenTimestamp
            )
        }
    }

    /**
     * Check if diarization is initialized
     */
    fun isInitialized(): Boolean = speakerExtractor != null

    /**
     * Release diarization resources.
     * Thread-safe: synchronized method that acquires embedding lock.
     */
    @Synchronized
    fun release() {
        embeddingLock.withLock {
            speakerExtractor?.release()
            speakerExtractor = null
        }
        speakerProfiles.clear()
        audioBuffer.clear()
        currentSpeakerId = -1
        nextSpeakerId = 1
        isEnabled = false
        lastConfidence = 0f
        lastEmbeddingQuality = 0f
        onboardingRemaining = 0
        candidateEmbeddings.clear()
        candidateSpeechMs = 0
        lastMatchedSpeakerId = -1
        android.util.Log.i("DiarizationManager", "Diarization released")
    }
}