package com.sherpaonnxofflinestt.managers

import android.util.Log
import com.k2fsa.sherpa.onnx.OfflinePunctuation
import com.k2fsa.sherpa.onnx.OfflinePunctuationConfig
import com.k2fsa.sherpa.onnx.OfflinePunctuationModelConfig
import java.io.File

/**
 * Manages offline punctuation restoration using sherpa-onnx.
 * Supports CT-Transformer models for adding punctuation to unpunctuated text.
 */
class PunctuationManager {
    companion object {
        private const val TAG = "PunctuationManager"
    }

    private var punctuation: OfflinePunctuation? = null
    private var _isEnabled: Boolean = false

    val isEnabled: Boolean
        get() = _isEnabled && punctuation != null

    /**
     * Initialize the punctuation model.
     *
     * @param modelPath Path to the CT-Transformer ONNX model file
     * @param numThreads Number of threads for inference (default: 2)
     * @param provider ONNX provider: "cpu", "nnapi", or "gpu" (default: "cpu")
     * @return true if initialization successful
     */
    fun initialize(
        modelPath: String,
        numThreads: Int = 2,
        provider: String = "cpu"
    ): Boolean {
        return try {
            // Validate model file exists
            if (!File(modelPath).exists()) {
                Log.e(TAG, "Model file not found: $modelPath")
                return false
            }

            val modelConfig = OfflinePunctuationModelConfig(
                ctTransformer = modelPath,
                numThreads = numThreads,
                debug = false,
                provider = provider
            )

            val config = OfflinePunctuationConfig(model = modelConfig)

            punctuation = OfflinePunctuation(config = config)
            _isEnabled = true

            Log.i(TAG, "Punctuation initialized successfully: $modelPath")
            true
        } catch (e: Exception) {
            Log.e(TAG, "Failed to initialize punctuation: ${e.message}", e)
            punctuation = null
            _isEnabled = false
            false
        }
    }

    /**
     * Check if the punctuation model is initialized.
     */
    fun isInitialized(): Boolean = punctuation != null

    /**
     * Add punctuation to the given text.
     *
     * @param text Unpunctuated text (e.g., "bonjour comment allez vous")
     * @return Punctuated text (e.g., "Bonjour, comment allez-vous?")
     */
    fun addPunctuation(text: String): String {
        val punct = punctuation
        if (punct == null) {
            Log.w(TAG, "Punctuation not initialized, returning original text")
            return text
        }

        if (text.isBlank()) {
            return text
        }

        return try {
            val startTime = System.currentTimeMillis()
            val result = punct.addPunctuation(text)
            val elapsed = System.currentTimeMillis() - startTime
            Log.d(TAG, "Punctuation added in ${elapsed}ms: \"$text\" -> \"$result\"")
            result
        } catch (e: Exception) {
            Log.e(TAG, "Failed to add punctuation: ${e.message}", e)
            text
        }
    }

    /**
     * Enable or disable punctuation processing.
     *
     * @param enabled Whether to enable punctuation
     * @return The new enabled state
     */
    fun setEnabled(enabled: Boolean): Boolean {
        if (enabled && punctuation == null) {
            Log.w(TAG, "Cannot enable punctuation - not initialized")
            return false
        }
        _isEnabled = enabled
        Log.i(TAG, "Punctuation ${if (enabled) "enabled" else "disabled"}")
        return _isEnabled
    }

    /**
     * Release the punctuation model and free resources.
     */
    fun release() {
        punctuation?.release()
        punctuation = null
        _isEnabled = false
        Log.i(TAG, "Punctuation released")
    }
}