import Foundation
import Accelerate

/// Speaker profile for diarization tracking
struct SpeakerProfile {
    let speakerId: Int
    var embedding: [Float]
    var numSegments: Int
    var status: String // "pending" or "confirmed"
    var totalSpeechDurationMs: Int64
    var lastSeenTimestamp: TimeInterval
}

/// Diarization configuration (uses STTConstants for defaults)
struct DiarizationConfig {
    var speakerThreshold: Float = STTConstants.defaultSpeakerThreshold
    var embeddingThreshold: Float = STTConstants.defaultEmbeddingThreshold
    var minSpeechDurationMs: Int = STTConstants.defaultMinSpeechDurationMs
    var minEmbeddingUpdateMs: Int = STTConstants.defaultMinEmbeddingUpdateMs
    var maxSpeakers: Int = STTConstants.defaultMaxSpeakers
    var embeddingAverageWeight: Float = STTConstants.defaultEmbeddingAverageWeight
    var onboardingSegments: Int = STTConstants.defaultOnboardingSegments
    var minSpeechForNewSpeaker: Int = STTConstants.defaultMinSpeechForNewSpeaker
}

/// Result of diarization processing
struct DiarizationResult {
    let speakerId: Int
    let status: String
    let justConfirmed: Bool
    let totalSpeakers: Int
    let confidence: Float
    let embeddingQuality: Float
}

/// Callback protocol for diarization events
protocol SpeakerDiarizationDelegate: AnyObject {
    func onSpeakerUpdate(_ result: DiarizationResult)
}

/// Manages speaker diarization using embedding extraction.
/// Extracted from SherpaOnnxOfflineStt.swift to reduce file size.
///
/// NOTE: This logic is duplicated in Android (DiarizationManager.kt).
/// Any changes here MUST be mirrored in Android to maintain platform parity.
class SpeakerDiarizationManager {
    // Speaker state
    private(set) var speakerProfiles: [SpeakerProfile] = []
    private(set) var currentSpeakerId: Int = -1
    private var nextSpeakerId: Int = 1

    // Configuration
    private(set) var config: DiarizationConfig
    private(set) var isEnabled: Bool = false

    // Thread safety
    private let lock = NSLock()

    // Metrics
    private(set) var lastConfidence: Float = 0
    private(set) var lastEmbeddingQuality: Float = 0

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

    // New speaker buffering: accumulate evidence before creating new speaker
    private var candidateEmbeddings: [[Float]] = []
    private var candidateSpeechMs: Int64 = 0
    private var lastMatchedSpeakerId: Int = -1

    // Delegate for events
    weak var delegate: SpeakerDiarizationDelegate?

    init(config: DiarizationConfig = DiarizationConfig()) {
        self.config = config
    }

    /// Enable diarization
    func enable() {
        lock.lock()
        defer { lock.unlock() }
        isEnabled = true
    }

    /// Disable diarization
    func disable() {
        lock.lock()
        defer { lock.unlock() }
        isEnabled = false
    }

    /// Update configuration
    func updateConfig(_ newConfig: DiarizationConfig) {
        lock.lock()
        defer { lock.unlock() }
        config = newConfig
    }

    /// Get current speaker count
    func getSpeakerCount() -> Int {
        lock.lock()
        defer { lock.unlock() }
        return speakerProfiles.count
    }

    /// Reset all speakers
    func resetSpeakers() {
        lock.lock()
        defer { lock.unlock() }

        speakerProfiles.removeAll()
        currentSpeakerId = -1
        nextSpeakerId = 1
        lastConfidence = 0
        lastEmbeddingQuality = 0
        onboardingRemaining = 0
        candidateEmbeddings.removeAll()
        candidateSpeechMs = 0
        lastMatchedSpeakerId = -1
    }

    /// Merge two speakers
    func mergeSpeakers(sourceId: Int, targetId: Int) -> Bool {
        lock.lock()
        defer { lock.unlock() }

        guard let sourceIndex = speakerProfiles.firstIndex(where: { $0.speakerId == sourceId }),
              let targetIndex = speakerProfiles.firstIndex(where: { $0.speakerId == targetId }) else {
            return false
        }

        // Merge stats into target
        speakerProfiles[targetIndex].numSegments += speakerProfiles[sourceIndex].numSegments
        speakerProfiles[targetIndex].totalSpeechDurationMs += speakerProfiles[sourceIndex].totalSpeechDurationMs

        // Average embeddings
        let alpha: Float = 0.5
        for i in 0..<speakerProfiles[targetIndex].embedding.count {
            speakerProfiles[targetIndex].embedding[i] =
                alpha * speakerProfiles[targetIndex].embedding[i] + (1 - alpha) * speakerProfiles[sourceIndex].embedding[i]
        }
        speakerProfiles[targetIndex].embedding = normalizeEmbedding(speakerProfiles[targetIndex].embedding)

        // Remove source speaker
        speakerProfiles.remove(at: sourceIndex)

        // Update current speaker if it was the source
        if currentSpeakerId == sourceId {
            currentSpeakerId = targetId
        }

        return true
    }

    /// Remove a specific speaker
    func removeSpeaker(speakerId: Int) -> Bool {
        lock.lock()
        defer { lock.unlock() }

        guard let index = speakerProfiles.firstIndex(where: { $0.speakerId == speakerId }) else {
            return false
        }

        speakerProfiles.remove(at: index)

        if currentSpeakerId == speakerId {
            currentSpeakerId = speakerProfiles.first?.speakerId ?? -1
        }

        return true
    }

    /// Get speaker profiles as dictionary array
    func getSpeakerProfiles() -> [[String: Any]] {
        lock.lock()
        defer { lock.unlock() }

        return speakerProfiles.map { profile in
            return [
                "speakerId": profile.speakerId,
                "status": profile.status,
                "segmentCount": profile.numSegments,
                "totalSpeechMs": profile.totalSpeechDurationMs,
                "lastSeenTimestamp": profile.lastSeenTimestamp
            ]
        }
    }

    /// Process an embedding and assign/create speaker
    /// Returns the assigned speaker ID
    func processEmbedding(_ embedding: [Float], speechDurationMs: Int64) -> Int {
        guard isEnabled else { return -1 }

        // Validate embedding
        let quality = validateEmbedding(embedding)
        lastEmbeddingQuality = quality >= 0 ? quality : 0

        if quality < 0 {
            STTLogger.w("Diarization", "Invalid embedding, skipping")
            return currentSpeakerId
        }

        // Normalize embedding for comparison
        let normalizedEmbedding = normalizeEmbedding(embedding)

        lock.lock()
        defer { lock.unlock() }

        // Case 1: No speakers exist yet - create Speaker 1 and start onboarding
        if speakerProfiles.isEmpty {
            return createFirstSpeaker(embedding: normalizedEmbedding, speechDurationMs: speechDurationMs)
        }

        // Case 2: Onboarding mode - always assign to Speaker 1
        if onboardingRemaining > 0 {
            return handleOnboarding(embedding: normalizedEmbedding, speechDurationMs: speechDurationMs)
        }

        // Case 3: Normal matching logic
        return matchOrCreateSpeaker(embedding: normalizedEmbedding, speechDurationMs: speechDurationMs)
    }

    // MARK: - Private Methods

    private func createFirstSpeaker(embedding: [Float], speechDurationMs: Int64) -> Int {
        let newId = nextSpeakerId
        nextSpeakerId += 1

        speakerProfiles.append(SpeakerProfile(
            speakerId: newId,
            embedding: embedding,
            numSegments: 1,
            status: "pending",
            totalSpeechDurationMs: speechDurationMs,
            lastSeenTimestamp: Date().timeIntervalSince1970
        ))
        currentSpeakerId = newId

        // Initialize onboarding counter
        onboardingRemaining = config.onboardingSegments - 1
        lastMatchedSpeakerId = newId
        lastConfidence = 1.0

        STTLogger.i("Diarization", "Created Speaker \(newId), onboarding started (\(onboardingRemaining) remaining)")

        notifyDelegate(speakerId: newId, status: "pending", justConfirmed: false)

        return newId
    }

    private func handleOnboarding(embedding: [Float], speechDurationMs: Int64) -> Int {
        guard let speaker1Index = speakerProfiles.indices.first else {
            return currentSpeakerId
        }

        speakerProfiles[speaker1Index].numSegments += 1
        speakerProfiles[speaker1Index].totalSpeechDurationMs += speechDurationMs
        speakerProfiles[speaker1Index].lastSeenTimestamp = Date().timeIntervalSince1970

        // Update embedding with running average during onboarding
        let alpha = config.embeddingAverageWeight
        for i in 0..<speakerProfiles[speaker1Index].embedding.count {
            speakerProfiles[speaker1Index].embedding[i] =
                (1 - alpha) * speakerProfiles[speaker1Index].embedding[i] + alpha * embedding[i]
        }
        // Re-normalize after averaging
        speakerProfiles[speaker1Index].embedding = normalizeEmbedding(speakerProfiles[speaker1Index].embedding)

        onboardingRemaining -= 1
        lastMatchedSpeakerId = speakerProfiles[speaker1Index].speakerId
        currentSpeakerId = speakerProfiles[speaker1Index].speakerId
        lastConfidence = 1.0

        let justConfirmed = speakerProfiles[speaker1Index].numSegments == 2
        if justConfirmed {
            speakerProfiles[speaker1Index].status = "confirmed"
        }

        STTLogger.i("Diarization", "Onboarding: assigned to Speaker \(speakerProfiles[speaker1Index].speakerId) (\(onboardingRemaining) remaining)")

        notifyDelegate(
            speakerId: speakerProfiles[speaker1Index].speakerId,
            status: speakerProfiles[speaker1Index].status,
            justConfirmed: justConfirmed
        )

        return speakerProfiles[speaker1Index].speakerId
    }

    private func matchOrCreateSpeaker(embedding: [Float], speechDurationMs: Int64) -> Int {
        var bestMatch = -1
        var bestSimilarity: Float = 0

        for (index, profile) in speakerProfiles.enumerated() {
            let similarity = cosineSimilarity(embedding, profile.embedding)
            STTLogger.d("Diarization", "vs Speaker \(profile.speakerId) similarity = \(String(format: "%.3f", similarity))")

            if similarity > bestSimilarity && similarity >= config.speakerThreshold {
                bestSimilarity = similarity
                bestMatch = index
            }
        }

        lastConfidence = bestSimilarity

        STTLogger.d("Diarization", "bestMatch=\(bestMatch), similarity=\(String(format: "%.3f", bestSimilarity)), threshold=\(config.speakerThreshold)")

        if bestMatch >= 0 {
            return updateExistingSpeaker(at: bestMatch, embedding: embedding, similarity: bestSimilarity, speechDurationMs: speechDurationMs)
        } else {
            return bufferNewSpeaker(embedding: embedding, speechDurationMs: speechDurationMs)
        }
    }

    private func updateExistingSpeaker(at index: Int, embedding: [Float], similarity: Float, speechDurationMs: Int64) -> Int {
        speakerProfiles[index].numSegments += 1
        speakerProfiles[index].totalSpeechDurationMs += speechDurationMs
        speakerProfiles[index].lastSeenTimestamp = Date().timeIntervalSince1970
        currentSpeakerId = speakerProfiles[index].speakerId

        // Clear candidate buffer - we matched an existing speaker
        candidateEmbeddings.removeAll()
        candidateSpeechMs = 0
        lastMatchedSpeakerId = speakerProfiles[index].speakerId

        // Only update embedding if similarity exceeds stricter threshold
        let shouldUpdateEmbedding = similarity >= config.embeddingThreshold
            && speechDurationMs >= Int64(config.minEmbeddingUpdateMs)

        if shouldUpdateEmbedding {
            let alpha = config.embeddingAverageWeight
            for i in 0..<speakerProfiles[index].embedding.count {
                speakerProfiles[index].embedding[i] =
                    (1 - alpha) * speakerProfiles[index].embedding[i] + alpha * embedding[i]
            }
            speakerProfiles[index].embedding = normalizeEmbedding(speakerProfiles[index].embedding)
            STTLogger.d("Diarization", "Updated embedding for speaker \(speakerProfiles[index].speakerId)")
        }

        let justConfirmed = speakerProfiles[index].numSegments == 2
        if justConfirmed {
            speakerProfiles[index].status = "confirmed"
        }

        notifyDelegate(
            speakerId: speakerProfiles[index].speakerId,
            status: speakerProfiles[index].status,
            justConfirmed: justConfirmed
        )

        return speakerProfiles[index].speakerId
    }

    private func bufferNewSpeaker(embedding: [Float], speechDurationMs: Int64) -> Int {
        candidateEmbeddings.append(embedding)
        candidateSpeechMs += speechDurationMs

        STTLogger.d("Diarization", "Buffering candidate: \(candidateSpeechMs)ms / \(config.minSpeechForNewSpeaker)ms required")

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

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

            // Create new speaker with averaged embedding
            let newId = nextSpeakerId
            nextSpeakerId += 1
            speakerProfiles.append(SpeakerProfile(
                speakerId: newId,
                embedding: averagedEmbedding,
                numSegments: 1,
                status: "pending",
                totalSpeechDurationMs: speechDurationMs,
                lastSeenTimestamp: Date().timeIntervalSince1970
            ))
            currentSpeakerId = newId
            lastMatchedSpeakerId = newId

            STTLogger.i("Diarization", "Created Speaker \(newId) from buffered segments (total: \(speakerProfiles.count))")
            notifyDelegate(speakerId: newId, status: "pending", justConfirmed: false)

            return newId
        } else {
            // Not enough evidence yet - return last matched speaker
            let fallbackSpeakerId = lastMatchedSpeakerId > 0 ? lastMatchedSpeakerId : (speakerProfiles.first?.speakerId ?? -1)
            currentSpeakerId = fallbackSpeakerId
            STTLogger.d("Diarization", "Returning fallback speaker \(fallbackSpeakerId) while buffering")
            return fallbackSpeakerId
        }
    }

    /// Evict the least recently used speaker to make room for a new one
    private func evictLRUSpeaker() {
        guard !speakerProfiles.isEmpty else { return }

        // Find speaker with oldest lastSeenTimestamp that is NOT the current speaker
        var lruIndex = -1
        var oldestTimestamp: TimeInterval = .greatestFiniteMagnitude

        for (index, profile) in speakerProfiles.enumerated() {
            if profile.speakerId != currentSpeakerId && profile.lastSeenTimestamp < oldestTimestamp {
                oldestTimestamp = profile.lastSeenTimestamp
                lruIndex = index
            }
        }

        if lruIndex >= 0 {
            let evicted = speakerProfiles.remove(at: lruIndex)
            let agoMs = Int((Date().timeIntervalSince1970 - oldestTimestamp) * 1000)
            STTLogger.i("Diarization", "Evicted speaker \(evicted.speakerId) (last seen: \(agoMs)ms ago)")
        }
    }

    private func notifyDelegate(speakerId: Int, status: String, justConfirmed: Bool) {
        let result = DiarizationResult(
            speakerId: speakerId,
            status: status,
            justConfirmed: justConfirmed,
            totalSpeakers: speakerProfiles.count,
            confidence: lastConfidence,
            embeddingQuality: lastEmbeddingQuality
        )
        delegate?.onSpeakerUpdate(result)
    }

    // MARK: - Embedding Utilities

    /// Validate an embedding and return quality score (0-1), or -1 if invalid
    func validateEmbedding(_ embedding: [Float]) -> Float {
        guard !embedding.isEmpty else { return -1 }

        // Check for NaN or Inf
        for value in embedding {
            if value.isNaN || value.isInfinite {
                return -1
            }
        }

        // Check norm (should not be near zero)
        var norm: Float = 0
        vDSP_dotpr(embedding, 1, embedding, 1, &norm, vDSP_Length(embedding.count))
        norm = sqrt(norm)

        guard norm > STTConstants.minEmbeddingNorm else { return -1 }

        // Calculate variance as quality indicator
        let mean = embedding.reduce(0, +) / Float(embedding.count)
        var variance: Float = 0
        for value in embedding {
            variance += (value - mean) * (value - mean)
        }
        variance /= Float(embedding.count)

        // Higher variance = more distinctive = better quality
        let quality = min(variance * STTConstants.embeddingQualityFactor, 1.0)
        return quality
    }

    /// Average multiple embeddings into one
    func averageEmbeddings(_ embeddings: [[Float]]) -> [Float] {
        guard !embeddings.isEmpty else { return [] }
        guard embeddings.count > 1 else { return embeddings[0] }

        let dim = embeddings[0].count
        var result = [Float](repeating: 0, count: dim)

        for embedding in embeddings {
            for i in 0..<dim {
                result[i] += embedding[i]
            }
        }

        let count = Float(embeddings.count)
        for i in 0..<dim {
            result[i] /= count
        }

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

    /// L2 normalize an embedding to unit length
    func normalizeEmbedding(_ embedding: [Float]) -> [Float] {
        var norm: Float = 0
        vDSP_dotpr(embedding, 1, embedding, 1, &norm, vDSP_Length(embedding.count))
        norm = sqrt(norm)

        guard norm > STTConstants.minEmbeddingNorm else { return embedding }

        var normalized = [Float](repeating: 0, count: embedding.count)
        var divisor = norm
        vDSP_vsdiv(embedding, 1, &divisor, &normalized, 1, vDSP_Length(embedding.count))
        return normalized
    }

    /// Calculate cosine similarity between two L2-normalized embeddings
    /// For normalized vectors, this is simply the dot product
    func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
        guard a.count == b.count else { return 0 }

        var dot: Float = 0
        vDSP_dotpr(a, 1, b, 1, &dot, vDSP_Length(a.count))
        return dot
    }
}