package com.sherpaonnxofflinestt.managers

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

/**
 * Diarization segment result from pyannote processing
 */
data class DiarizationSegment(
    val startTimeMs: Long,
    val endTimeMs: Long,
    val speakerId: Int
)

/**
 * Configuration for pyannote-based speaker segmentation
 */
data class PyannoteSegmentationConfig(
    /** Path to pyannote segmentation model (model.onnx) */
    val segmentationModelPath: String = "",
    /** Path to speaker embedding model for clustering */
    val embeddingModelPath: String = "",
    /** Number of speakers (0 = auto-detect using threshold) */
    val numSpeakers: Int = 0,
    /** Distance threshold for clustering (used when numSpeakers = 0). Smaller = more clusters. */
    val clusteringThreshold: Float = STTConstants.DEFAULT_CLUSTERING_THRESHOLD,
    /** Minimum duration for speech segment (seconds) */
    val minDurationOn: Float = STTConstants.DEFAULT_MIN_DURATION_ON,
    /** Minimum duration for silence segment (seconds) */
    val minDurationOff: Float = STTConstants.DEFAULT_MIN_DURATION_OFF,
    /** Number of threads for inference */
    val numThreads: Int = STTConstants.DEFAULT_AUXILIARY_NUM_THREADS,
    /** ONNX provider: "cpu", "nnapi", "gpu" */
    val provider: String = STTConstants.DEFAULT_PROVIDER
)

/**
 * Manages speaker diarization using sherpa-onnx's native pyannote segmentation pipeline.
 *
 * This provides frame-by-frame speaker change detection using the pyannote-segmentation-3.0 model,
 * which is specifically designed for detecting:
 * - Speaker boundaries (when speakers change)
 * - Overlapped speech (when multiple speakers talk simultaneously)
 *
 * The pipeline combines:
 * 1. Pyannote segmentation model - Detects speaker activity per frame
 * 2. Speaker embedding extraction - Creates speaker signatures
 * 3. Fast clustering - Groups segments by speaker
 */
class PyannoteSegmentationManager(
    private val eventManager: EventEmissionManager
) {
    private var diarizer: OfflineSpeakerDiarization? = null
    private var config: PyannoteSegmentationConfig = PyannoteSegmentationConfig()

    var isEnabled: Boolean = false
        private set

    /**
     * Initialize the pyannote-based diarization pipeline.
     *
     * @param segmentationModelPath Path to pyannote model directory (containing model.onnx)
     * @param embeddingModelPath Path to speaker embedding model (for clustering)
     * @param config Configuration options
     * @return true if initialization successful
     */
    fun initialize(
        segmentationModelPath: String,
        embeddingModelPath: String,
        segmentationConfig: PyannoteSegmentationConfig = PyannoteSegmentationConfig()
    ): Boolean {
        this.config = segmentationConfig.copy(
            segmentationModelPath = segmentationModelPath,
            embeddingModelPath = embeddingModelPath
        )

        // Validate model paths
        val segModelFile = File(segmentationModelPath)
        val embModelFile = File(embeddingModelPath)

        if (!segModelFile.exists()) {
            STTLogger.Pyannote.w("Segmentation model not found: $segmentationModelPath")
            return false
        }
        if (!embModelFile.exists()) {
            STTLogger.Pyannote.w("Embedding model not found: $embeddingModelPath")
            return false
        }

        return try {
            val startTime = System.currentTimeMillis()

            // Configure pyannote segmentation model
            val pyannoteConfig = OfflineSpeakerSegmentationPyannoteModelConfig(
                model = segmentationModelPath
            )

            val segmentationModelConfig = OfflineSpeakerSegmentationModelConfig(
                pyannote = pyannoteConfig,
                numThreads = config.numThreads,
                debug = false,
                provider = config.provider
            )

            // Configure speaker embedding extractor
            val embeddingConfig = SpeakerEmbeddingExtractorConfig(
                model = embeddingModelPath,
                numThreads = config.numThreads,
                debug = false,
                provider = config.provider
            )

            // Configure clustering
            val clusteringConfig = FastClusteringConfig(
                numClusters = config.numSpeakers,
                threshold = config.clusteringThreshold
            )

            // Build diarization config
            val diarizationConfig = OfflineSpeakerDiarizationConfig(
                segmentation = segmentationModelConfig,
                embedding = embeddingConfig,
                clustering = clusteringConfig,
                minDurationOn = config.minDurationOn,
                minDurationOff = config.minDurationOff
            )

            // Create diarizer
            diarizer = OfflineSpeakerDiarization(assetManager = null, config = diarizationConfig)
            isEnabled = true

            val initTime = System.currentTimeMillis() - startTime
            STTLogger.Pyannote.i(
                "Pyannote diarization initialized in ${initTime}ms " +
                "(numSpeakers=${config.numSpeakers}, threshold=${config.clusteringThreshold})")
            true
        } catch (e: Exception) {
            STTLogger.Pyannote.e("Failed to initialize: ${e.message}", e)
            isEnabled = false
            false
        }
    }

    /**
     * Get the expected sample rate for the diarizer.
     * @return Sample rate in Hz (typically 16000)
     */
    fun getSampleRate(): Int {
        return diarizer?.sampleRate() ?: 16000
    }

    /**
     * Process audio samples and return diarization segments.
     *
     * This processes the entire audio buffer at once using the pyannote model
     * to detect speaker boundaries and assign speaker IDs.
     *
     * @param samples Audio samples (mono, float, -1 to 1 range)
     * @param sampleRate Sample rate in Hz (should match getSampleRate())
     * @return List of diarization segments with speaker assignments, or null on error
     */
    fun process(samples: FloatArray, sampleRate: Int): List<DiarizationSegment>? {
        val diar = diarizer
        if (!isEnabled || diar == null) {
            STTLogger.Pyannote.w("Process called but diarizer not initialized")
            return null
        }

        if (samples.isEmpty()) {
            STTLogger.Pyannote.w("Empty samples array")
            return null
        }

        val durationMs = (samples.size.toLong() * 1000) / sampleRate
        STTLogger.Pyannote.i(
            "[PYANNOTE DEBUG] Processing ${samples.size} samples (${durationMs}ms) @ ${sampleRate}Hz")

        return try {
            val startTime = System.currentTimeMillis()

            // Process through diarization pipeline
            val segments = diar.process(samples)

            val processingTime = System.currentTimeMillis() - startTime
            STTLogger.Pyannote.i(
                "[PYANNOTE DEBUG] Processed in ${processingTime}ms, found ${segments.size} segments")

            // Convert to our segment format
            // OfflineSpeakerDiarizationSegment properties: start (Float), end (Float), speaker (Int)
            val result = segments.map { segment ->
                DiarizationSegment(
                    startTimeMs = (segment.start * 1000).toLong(),
                    endTimeMs = (segment.end * 1000).toLong(),
                    speakerId = segment.speaker
                )
            }

            // Log segment details
            if (result.isNotEmpty()) {
                val speakers = result.map { it.speakerId }.distinct().sorted()
                STTLogger.Pyannote.i("[PYANNOTE DEBUG] Speakers detected: $speakers")

                result.forEachIndexed { index, seg ->
                    STTLogger.Pyannote.d(
                        "[PYANNOTE DEBUG] Segment $index: ${seg.startTimeMs}-${seg.endTimeMs}ms -> Speaker ${seg.speakerId}")
                }

                // Emit multi-speaker event if multiple speakers detected
                if (speakers.size > 1) {
                    val windowPairs = result.map { Pair(it.startTimeMs.toInt(), it.speakerId) }
                    eventManager.sendMultiSpeakerDetected(speakers, windowPairs)
                }
            } else {
                STTLogger.Pyannote.w("[PYANNOTE DEBUG] No segments detected")
            }

            result
        } catch (e: Exception) {
            STTLogger.Pyannote.e("Processing error: ${e.message}", e)
            null
        }
    }

    /**
     * Process audio and return the primary speaker ID.
     *
     * This is a simplified interface that returns just the dominant speaker
     * for compatibility with the existing diarization flow.
     *
     * @param samples Audio samples
     * @param sampleRate Sample rate
     * @return Primary speaker ID (most speaking time), or -1 if none detected
     */
    fun processToPrimarySpeaker(samples: FloatArray, sampleRate: Int): Int {
        val segments = process(samples, sampleRate) ?: return -1
        if (segments.isEmpty()) return -1

        // Calculate speaking time per speaker
        val speakingTime = mutableMapOf<Int, Long>()
        segments.forEach { seg ->
            val duration = seg.endTimeMs - seg.startTimeMs
            speakingTime[seg.speakerId] = (speakingTime[seg.speakerId] ?: 0) + duration
        }

        // Return speaker with most speaking time
        return speakingTime.maxByOrNull { it.value }?.key ?: -1
    }

    /**
     * Set the expected number of speakers.
     * Call this before process() if you know how many speakers to expect.
     * Setting to -1 enables auto-detection using the clustering threshold.
     */
    fun setNumSpeakers(numSpeakers: Int) {
        val diar = diarizer ?: return
        // Update config with new clustering settings
        val newConfig = OfflineSpeakerDiarizationConfig(
            clustering = FastClusteringConfig(
                numClusters = numSpeakers,
                threshold = config.clusteringThreshold
            )
        )
        diar.setConfig(newConfig)
        STTLogger.Pyannote.i("Set numSpeakers to $numSpeakers")
    }

    /**
     * Set the clustering threshold for auto speaker count detection.
     * Only used when numSpeakers = -1.
     * Smaller values = more clusters (more speakers detected).
     */
    fun setClusteringThreshold(threshold: Float) {
        val diar = diarizer ?: return
        // Update config with new clustering settings
        val newConfig = OfflineSpeakerDiarizationConfig(
            clustering = FastClusteringConfig(
                numClusters = config.numSpeakers,
                threshold = threshold
            )
        )
        diar.setConfig(newConfig)
        STTLogger.Pyannote.i("Set clustering threshold to $threshold")
    }

    /**
     * Check if the manager is initialized
     */
    fun isInitialized(): Boolean = diarizer != null

    /**
     * Release all resources
     */
    fun release() {
        diarizer?.release()
        diarizer = null
        isEnabled = false
        STTLogger.Pyannote.i("Pyannote segmentation released")
    }
}