package com.example.myapplication

import android.app.NotificationChannel
import android.app.NotificationManager
import android.content.Context
import android.content.res.ColorStateList
import android.graphics.Color
import android.inputmethodservice.InputMethodService
import android.media.AudioManager
import android.os.Build
import android.os.Handler
import android.os.Looper
import android.util.Log
import android.view.Gravity
import android.view.View
import android.view.ViewGroup
import android.view.inputmethod.EditorInfo
import android.view.inputmethod.InputConnection
import android.widget.HorizontalScrollView
import android.widget.LinearLayout
import android.widget.PopupWindow
import android.widget.TextView
import androidx.annotation.ColorInt
import androidx.annotation.DrawableRes
import androidx.core.app.NotificationCompat
import com.example.myapplication.data.WordDictionary
import com.example.myapplication.data.LanguageModelLoader
import com.example.myapplication.theme.ThemeManager
import com.example.myapplication.keyboard.KeyboardEnvironment
import com.example.myapplication.keyboard.MainKeyboard
import com.example.myapplication.keyboard.NumberKeyboard
import com.example.myapplication.keyboard.SymbolKeyboard
import com.example.myapplication.keyboard.AiKeyboard

class MyInputMethodService : InputMethodService(), KeyboardEnvironment {

    // ====== 键盘视图 & 键盘对象 ======
    private var currentKeyboardView: View? = null

    private var mainKeyboardView: View? = null
    private var numberKeyboardView: View? = null
    private var symbolKeyboardView: View? = null
    private var aiKeyboardView: View? = null

    private var mainKeyboard: MainKeyboard? = null
    private var numberKeyboard: NumberKeyboard? = null
    private var symbolKeyboard: SymbolKeyboard? = null
    private var aiKeyboard: AiKeyboard? = null

    // 单键预览（只给建议栏等用，按键的预览已经放到各 Keyboard 里了）
    private var keyPreviewPopup: PopupWindow? = null

    private var lastWordForLM: String? = null  // 上一次输入的词

    private val wordDictionary = WordDictionary(this)          // 词库
    private var currentInput = StringBuilder()                 // 当前输入前缀
    private var completionSuggestions = emptyList<String>()    // 自动完成建议

    @Volatile private var isSpecialToken: BooleanArray = BooleanArray(0)

    private val suggestionStats by lazy { SuggestionStats(applicationContext) }
    private val specialTokens = setOf("<unk>", "<s>", "</s>")

    @Volatile private var bigramModel: com.example.myapplication.data.BigramModel? = null
    @Volatile private var word2id: Map<String, Int> = emptyMap()
    @Volatile private var id2word: List<String> = emptyList()
    @Volatile private var bigramReady: Boolean = false

    companion object {
        private const val TAG = "MyIME"
        private const val NOTIFICATION_CHANNEL_ID = "input_method_channel"
        private const val NOTIFICATION_ID = 1
    }

    // ===== KeyboardEnvironment 实现所需属性 =====

    override val ctx: Context
        get() = this

    // override val layoutInflater get() = super.getLayoutInflater()

    override val mainHandler: Handler = Handler(Looper.getMainLooper())

    override val audioManager: AudioManager by lazy {
        getSystemService(AUDIO_SERVICE) as AudioManager
    }


    // ===== 输入视图状态 =====

    private var isInputViewShownFlag = false

    // Shift 状态
    private var isShiftOn = false

    // 删除长按
    private var isDeleting = false
    private val repeatDelInitialDelay = 350L
    private val repeatDelInterval = 50L

    private val repeatDelRunnable = object : Runnable {
        override fun run() {
            if (!isDeleting) return
            deleteOne()
            mainHandler.postDelayed(this, repeatDelInterval)
        }
    }

    // ===== 颜色缓存 =====
    override var currentTextColor: ColorStateList =
        ColorStateList.valueOf(Color.parseColor("#000000"))//联想字体颜色
        private set

    override var currentBorderColor: ColorStateList =
        ColorStateList.valueOf(Color.parseColor("#1A000000"))
        private set

    override var currentBackgroundColor: ColorStateList =
        ColorStateList.valueOf(Color.TRANSPARENT)
        private set

    // 副字符映射表（主键盘上滑）
    private val swipeAltMap: Map<Char, Char> = mapOf(
        'q' to '1',
        'w' to '2',
        'e' to '3',
        'r' to '4',
        't' to '5',
        'y' to '6',
        'u' to '7',
        'i' to '8',
        'o' to '9',
        'p' to '0',

        'a' to '!',
        's' to '@',
        'd' to '#',
        'f' to '$',
        'g' to '%',
        'h' to '^',
        'j' to '&',
        'k' to '*',
        'l' to '(',

        'z' to ')',
        'x' to '-',
        'c' to '+',
        'v' to '=',
        'b' to '?',
        'n' to '/',
        'm' to '.'
    )

    // ================= 生命周期 =================

    override fun onCreate() {
        super.onCreate()

        ThemeManager.ensureBuiltInThemesInstalled(this)
        ThemeManager.init(this)

        // 异步加载词典与 bigram 模型
        Thread {
            // 1) Trie 词典
            try {
                wordDictionary.loadIfNeeded()
            } catch (e: Throwable) {
                Log.w(TAG, "Trie load failed: ${e.message}", e)
            }

            // 2) Bigram 模型
            try {
                val m = LanguageModelLoader.load(this)

                require(m.biRowptr.size == m.vocab.size + 1) {
                    "biRowptr size ${m.biRowptr.size} != vocab.size+1 ${m.vocab.size + 1}"
                }
                require(m.biCols.size == m.biLogp.size) {
                    "biCols size ${m.biCols.size} != biLogp size ${m.biLogp.size}"
                }

                bigramModel = m

                val map = HashMap<String, Int>(m.vocab.size * 2)
                m.vocab.forEachIndexed { idx, w -> map[w.lowercase()] = idx }
                word2id = map
                id2word = m.vocab

                isSpecialToken = BooleanArray(id2word.size)
                for (i in id2word.indices) {
                    if (specialTokens.contains(id2word[i])) {
                        isSpecialToken[i] = true
                    }
                }

                bigramReady = true
            } catch (e: Throwable) {
                bigramReady = false
                Log.w(TAG, "Bigram load failed: ${e.message}", e)
            }
        }.start()

        createNotificationChannelIfNeeded()
        tryStartForegroundSafe()
    }

    private fun createNotificationChannelIfNeeded() {
        if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O) {
            val channel = NotificationChannel(
                NOTIFICATION_CHANNEL_ID,
                "输入法服务",
                NotificationManager.IMPORTANCE_LOW
            ).apply { description = "保持输入法服务运行" }

            val manager = getSystemService(Context.NOTIFICATION_SERVICE) as NotificationManager
            manager.createNotificationChannel(channel)
        }
    }

    private fun tryStartForegroundSafe() {
        try {
            val notification = NotificationCompat.Builder(this, NOTIFICATION_CHANNEL_ID)
                .setContentTitle("输入法服务")
                .setContentText("输入法正在运行")
                .setSmallIcon(android.R.drawable.ic_dialog_info)
                .setPriority(NotificationCompat.PRIORITY_LOW)
                .setOngoing(true)
                .build()

            startForeground(NOTIFICATION_ID, notification)
            Log.i(TAG, "startForeground success")
        } catch (e: Exception) {
            Log.w(TAG, "startForeground failed: ${e.message}")
        }
    }

    // 创建输入视图：此处只负责选择主键盘
    override fun onCreateInputView(): View {
        val keyboard = ensureMainKeyboard()
        currentKeyboardView = keyboard.rootView
        mainKeyboardView = keyboard.rootView
        return keyboard.rootView
    }

    override fun onStartInputView(info: EditorInfo?, restarting: Boolean) {
        super.onStartInputView(info, restarting)
        isInputViewShownFlag = true
        showMainKeyboard()
    }

    override fun onFinishInputView(finishingInput: Boolean) {
        super.onFinishInputView(finishingInput)
        isInputViewShownFlag = false
        stopRepeatDelete()

        mainHandler.postDelayed({
            if (!isInputViewShownFlag) {
                try {
                    hideWindow()
                } catch (e: Exception) {
                    Log.w(TAG, "hideWindow failed: ${e.message}")
                }
            }
        }, 50)
    }

    override fun onDestroy() {
        super.onDestroy()
        stopRepeatDelete()
    }

    // ================= KeyboardEnvironment：键盘切换 =================

    private fun ensureMainKeyboard(): MainKeyboard {
        if (mainKeyboard == null) {
            mainKeyboard = MainKeyboard(
                env = this,
                swipeAltMap = swipeAltMap,
                onToggleShift = {
                    isShiftOn = !isShiftOn
                    isShiftOn
                }
            )
            mainKeyboardView = mainKeyboard!!.rootView
    
            // 这里再给删除键挂长按连删（用你自己的 attachRepeatDeleteInternal）
            val delId = resources.getIdentifier("key_del", "id", packageName)
            mainKeyboardView?.findViewById<View?>(delId)?.let { attachRepeatDeleteInternal(it) }
        }
        return mainKeyboard!!
    }

    private fun ensureNumberKeyboard(): NumberKeyboard {
        if (numberKeyboard == null) {
            numberKeyboard = NumberKeyboard(this)
            numberKeyboardView = numberKeyboard!!.rootView
    
            //  数字键盘删除键 key_del
            val delId = resources.getIdentifier("key_del", "id", packageName)
            numberKeyboardView?.findViewById<View?>(delId)?.let { attachRepeatDeleteInternal(it) }
        }
        return numberKeyboard!!
    }
    
    private fun ensureSymbolKeyboard(): SymbolKeyboard {
        if (symbolKeyboard == null) {
            symbolKeyboard = SymbolKeyboard(this)
            symbolKeyboardView = symbolKeyboard!!.rootView
    
            //  符号键盘删除键 key_backspace
            val delId = resources.getIdentifier("key_backspace", "id", packageName)
            symbolKeyboardView?.findViewById<View?>(delId)?.let { attachRepeatDeleteInternal(it) }
        }
        return symbolKeyboard!!
    }

    private fun ensureAiKeyboard(): AiKeyboard {
        if (aiKeyboard == null) {
            aiKeyboard = AiKeyboard(this)
            aiKeyboardView = aiKeyboard!!.rootView
        }
        return aiKeyboard!!
    }

    override fun showMainKeyboard() {
        val kb = ensureMainKeyboard()
        currentKeyboardView = kb.rootView
        setInputView(kb.rootView)
        kb.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
    }
    
    override fun showNumberKeyboard() {
        val kb = ensureNumberKeyboard()
        currentKeyboardView = kb.rootView
        setInputView(kb.rootView)
        kb.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
    }
    
    override fun showSymbolKeyboard() {
        val kb = ensureSymbolKeyboard()
        currentKeyboardView = kb.rootView
        setInputView(kb.rootView)
        kb.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
    }
    
    override fun showAiKeyboard() {
        val kb = ensureAiKeyboard()
        currentKeyboardView = kb.rootView
        setInputView(kb.rootView)
        kb.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
    }
    

    // ================== 文本输入核心逻辑 ==================

    /** 获取上一个完整英文词（纯字母），统一小写 */
    private fun getPrevWordBeforeCursor(maxLen: Int = 128): String? {
        val before = currentInputConnection
            ?.getTextBeforeCursor(maxLen, 0)
            ?.toString()
            ?: return null

        // 如果最后一个字符是字母，说明词不完整，返回 null
        if (before.isNotEmpty() && before.last().isLetter()) return null

        val toks = before
            .replace(Regex("[^A-Za-z]"), " ")
            .trim()
            .split(Regex("\\s+"))

        if (toks.isEmpty()) return null

        val word = toks.last().lowercase()
        if (word.isEmpty()) return null
        return word
    }

    // 提交一个字符（原 sendKey）
    override fun commitKey(c: Char) {
        val ic = currentInputConnection ?: return

        val toSend = if (isShiftOn && c in 'a'..'z') c.uppercaseChar() else c
        ic.commitText(toSend.toString(), 1)

        when {
            c.isLetter() -> {
                val prefix = getCurrentWordPrefix()
                updateCompletionsAndRender(prefix)
            }
            c == ' ' || !c.isLetter() -> {
                updateCompletionsAndRender(prefix = "")
            }
        }

        playKeyClick()
    }

    // 删除一个字符（原 handleBackspace）
    override fun deleteOne() {
        val ic = currentInputConnection ?: return

        ic.deleteSurroundingText(1, 0)

        val prefix = getCurrentWordPrefix()
        updateCompletionsAndRender(prefix)

        playKeyClick()
    }

    // 发送（原 performSendAction）
    override fun performSendAction() {
        val ic = currentInputConnection ?: return

        val handled = ic.performEditorAction(EditorInfo.IME_ACTION_SEND)
        if (!handled) {
            ic.commitText("\n", 1)
        }
        playKeyClick()
    }

    // 按键音效
    override fun playKeyClick() {
        try {
            audioManager.playSoundEffect(AudioManager.FX_KEY_CLICK)
        } catch (_: Throwable) {
        }
    }

    // ================== 补全 & 联想 ==================

    override fun getCurrentWordPrefix(maxLen: Int): String {
        val before = currentInputConnection?.getTextBeforeCursor(maxLen, 0)?.toString() ?: ""
        val match = Regex("[A-Za-z]+$").find(before)
        return (match?.value ?: "").lowercase()
    }

    // 统一处理补全/联想
    override fun updateCompletionsAndRender(prefix: String) {
        currentInput.clear()
        currentInput.append(prefix)

        val lastWord = getPrevWordBeforeCursor()

        Thread {
            val list = try {
                if (prefix.isEmpty()) {
                    if (lastWord == null) {
                        emptyList()
                    } else {
                        suggestWithBigram("", lastWord, topK = 20)
                    }
                } else {
                    val fromBi = suggestWithBigram(prefix, lastWord, topK = 20)
                    if (fromBi.isNotEmpty()) {
                        fromBi.filter { it != prefix }
                    } else {
                        wordDictionary.wordTrie.startsWith(prefix)
                            .take(20)
                            .filter { it != prefix }
                    }
                }
            } catch (_: Throwable) {
                if (prefix.isNotEmpty()) {
                    wordDictionary.wordTrie.startsWith(prefix)
                        .take(20)
                        .filterNot { it == prefix }
                } else {
                    emptyList()
                }
            }

            mainHandler.post {
                completionSuggestions = suggestionStats.sortByCount(list.distinct().take(20))
                showCompletionSuggestions()
            }
        }.start()
    }

    // 显示自动完成建议（布局不变）
    private fun showCompletionSuggestions() {
        mainHandler.post {
            val suggestionsView =
                mainKeyboardView?.findViewById<LinearLayout>(R.id.completion_suggestions)

            val suggestions = (0..20).map { i ->
                mainKeyboardView?.findViewById<TextView>(
                    resources.getIdentifier("suggestion_$i", "id", packageName)
                )
            }

            // suggestion_0 显示 prefix
            val prefix = getCurrentWordPrefix()
            suggestions[0]?.text = prefix
            suggestions[0]?.visibility = if (prefix.isEmpty()) View.GONE else View.VISIBLE
            suggestions[0]?.setOnClickListener {
                insertCompletion(prefix)
            }

            // suggestion_1.. 按 completionSuggestions 填充
            suggestions.drop(1).forEachIndexed { index, textView ->
                textView?.text = completionSuggestions.getOrNull(index) ?: ""
                if (index < completionSuggestions.size) {
                    textView?.setOnClickListener {
                        suggestionStats.incClick(completionSuggestions[index])
                        insertCompletion(completionSuggestions[index])
                    }
                } else {
                    textView?.setOnClickListener(null)
                }
            }

            scrollSuggestionsToStart()
        }
    }

    // 自动滚回到最左边
    private fun scrollSuggestionsToStart() {
        val sv = mainKeyboardView?.findViewById<HorizontalScrollView>(R.id.completion_scroll)
        sv?.post { sv.fullScroll(View.FOCUS_LEFT) }
    }

    // 插入一个补全词
    private fun insertCompletion(word: String) {
        val ic = currentInputConnection ?: return

        ic.beginBatchEdit()
        try {
            val len = currentInput.length
            if (len > 0) ic.deleteSurroundingText(len, 0)
            ic.commitText("$word ", 1)
        } finally {
            ic.endBatchEdit()
        }

        currentInput.clear()
        // 补完一个词后，直接联想下一个词
        updateCompletionsAndRender(prefix = "")
    }

    // ================== 长按删除 ==================

    // 真正实现逻辑（基本照搬你原来的 attachRepeatDelete）
    private fun attachRepeatDeleteInternal(view: View) {
        view.setOnLongClickListener {
            if (!isDeleting) {
                isDeleting = true
                mainHandler.postDelayed(repeatDelRunnable, repeatDelInitialDelay)
                deleteOne() // 首次立刻删一次
            }
            true
        }
        view.setOnTouchListener { _, event ->
            when (event.actionMasked) {
                android.view.MotionEvent.ACTION_UP,
                android.view.MotionEvent.ACTION_CANCEL,
                android.view.MotionEvent.ACTION_OUTSIDE -> stopRepeatDelete()
            }
            false
        }
    }

    private fun stopRepeatDelete() {
        if (isDeleting) {
            isDeleting = false
            mainHandler.removeCallbacks(repeatDelRunnable)
        }
    }

    // ================== 主题 / 样式接口（保留原 API） ==================

    // 键盘背景大图（LinearLayout R.id.background 用的）
    fun setKeyboardBackground(@DrawableRes resId: Int) {
        mainHandler.post {
            mainKeyboardView
                ?.findViewById<LinearLayout>(R.id.background)
                ?.setBackgroundResource(resId)

            numberKeyboardView
                ?.findViewById<LinearLayout>(R.id.background)
                ?.setBackgroundResource(resId)

            symbolKeyboardView
                ?.findViewById<LinearLayout>(R.id.background)
                ?.setBackgroundResource(resId)

            aiKeyboardView
                ?.findViewById<LinearLayout>(R.id.background)
                ?.setBackgroundResource(resId)
        }
    }

    // ===== 文字颜色 =====

    fun setKeyTextColorInt(@ColorInt colorInt: Int) {
        setKeyTextColor(ColorStateList.valueOf(colorInt))
    }

    fun setKeyTextColor(colorStateList: ColorStateList) {
        currentTextColor = colorStateList
        mainHandler.post { applyThemeToAllKeyboards() }
    }

    fun setKeyTextColorWithStates() {
        val states = arrayOf(
            intArrayOf(android.R.attr.state_pressed),
            intArrayOf(android.R.attr.state_activated),
            intArrayOf()
        )

        val colors = intArrayOf(
            Color.parseColor("#222222"), // pressed
            Color.parseColor("#2563EB"), // activated
            Color.BLACK                  // default
        )

        setKeyTextColor(ColorStateList(states, colors))
    }

    // ===== 边框颜色 =====

    fun setKeyBorderColorInt(@ColorInt colorInt: Int) {
        setKeyBorderColor(ColorStateList.valueOf(colorInt))
    }

    fun setKeyBorderColor(colorStateList: ColorStateList) {
        currentBorderColor = colorStateList
        mainHandler.post { applyThemeToAllKeyboards() }
    }

    // ===== 背景颜色（如果你有用纯色背景的话） =====

    fun setKeyBackgroundColor(colorStateList: ColorStateList) {
        currentBackgroundColor = colorStateList
        mainHandler.post { applyThemeToAllKeyboards() }
    }

    fun setKeyBackgroundColorInt(@ColorInt colorInt: Int) {
        setKeyBackgroundColor(ColorStateList.valueOf(colorInt))
    }

    private fun applyThemeToAllKeyboards() {
        mainKeyboard?.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
        numberKeyboard?.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
        symbolKeyboard?.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
        aiKeyboard?.applyTheme(currentTextColor, currentBorderColor, currentBackgroundColor)
    }

    // ================== bigram & 联想实现 ==================

    private fun suggestWithBigram(prefix: String, lastWord: String?, topK: Int = 20): List<String> {
        val m = bigramModel
        if (m == null || !bigramReady) {
            return if (prefix.isNotEmpty()) {
                wordDictionary.wordTrie.startsWith(prefix).take(topK)
            } else {
                emptyList()
            }
        }

        val pf = prefix.lowercase()
        val last = lastWord?.lowercase()
        val lastId = last?.let { word2id[it] }

        if (lastId != null && lastId >= 0 && lastId + 1 < m.biRowptr.size) {
            val start = m.biRowptr[lastId]
            val end = m.biRowptr[lastId + 1]

            if (start in 0..end && end <= m.biCols.size) {

                // 带前缀过滤
                val buf = ArrayList<Pair<String, Float>>(maxOf(0, end - start))
                var i = start
                while (i < end) {
                    val nextId = m.biCols[i]
                    if (nextId in id2word.indices && !isSpecialToken[nextId]) {
                        val w = id2word[nextId]
                        if (pf.isEmpty() || w.startsWith(pf)) {
                            buf.add(w to m.biLogp[i])
                        }
                    }
                    i++
                }
                if (buf.isNotEmpty()) return topKByScore(buf, topK)

                // 无前缀兜底
                val allBuf = ArrayList<Pair<String, Float>>(maxOf(0, end - start))
                i = start
                while (i < end) {
                    val nextId = m.biCols[i]
                    if (nextId in id2word.indices && !isSpecialToken[nextId]) {
                        allBuf.add(id2word[nextId] to m.biLogp[i])
                    }
                    i++
                }
                if (allBuf.isNotEmpty()) return topKByScore(allBuf, topK)
            }
        }

        // —— 无上文 或 无出边 —— 
        return if (pf.isNotEmpty()) {
            wordDictionary.wordTrie.startsWith(pf).take(topK)
        } else {
            unigramTopKFiltered(topK)
        }
    }

    private fun unigramTopKFiltered(topK: Int): List<String> {
        val m = bigramModel ?: return emptyList()
        if (!bigramReady) return emptyList()

        val heap = java.util.PriorityQueue<Pair<String, Float>>(topK.coerceAtLeast(1)) { a, b ->
            a.second.compareTo(b.second)
        }

        var i = 0
        val n = id2word.size
        while (i < n) {
            if (!isSpecialToken[i]) {
                heap.offer(id2word[i] to m.uniLogp[i])
                if (heap.size > topK) heap.poll()
            }
            i++
        }

        val out = ArrayList<String>(heap.size)
        while (heap.isNotEmpty()) out.add(heap.poll().first)
        out.reverse()
        return out
    }

    private fun topKByScore(pairs: List<Pair<String, Float>>, k: Int): List<String> {
        val heap = java.util.PriorityQueue<Pair<String, Float>>(k.coerceAtLeast(1)) { a, b ->
            a.second.compareTo(b.second)
        }
        for (p in pairs) {
            heap.offer(p)
            if (heap.size > k) heap.poll()
        }
        val out = ArrayList<String>(heap.size)
        while (heap.isNotEmpty()) out.add(heap.poll().first)
        out.reverse()
        return out
    }
}
