package com.rizoma.trazo.data

import android.content.Context
import androidx.datastore.core.DataStore
import androidx.datastore.preferences.core.Preferences
import androidx.datastore.preferences.core.edit
import androidx.datastore.preferences.core.intPreferencesKey
import androidx.datastore.preferences.preferencesDataStore
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Dispatchers
import kotlinx.coroutines.flow.first
import kotlinx.coroutines.launch
import java.util.concurrent.ConcurrentHashMap

private val Context.trazoLearnedStore: DataStore<Preferences> by preferencesDataStore(name = "trazo_learned")

/**
 * Predicción por prefijo, bilingüe (es + en). Dos fuentes:
 *  - [base]: diccionarios empaquetados (assets), peso fijo por frecuencia.
 *  - [learned]: palabras que el usuario escribe, contadas y persistidas.
 *
 * Mecanismo "buenas sugerencias": una palabra del diccionario se refuerza al instante,
 * pero una palabra NUEVA solo entra a las sugerencias tras escribirse [PROMOTE] veces,
 * así los errores escritos una sola vez nunca contaminan. Las sugerencias salen en el
 * idioma del prefijo que escribes (ambos diccionarios conviven).
 */
class SuggestionEngine(context: Context, private val scope: CoroutineScope) {

    private val appContext = context.applicationContext
    private val store = appContext.trazoLearnedStore

    // Diccionarios separados por idioma para poder activarlos/desactivarlos.
    private val baseEs = ConcurrentHashMap<String, Int>()
    private val baseEn = ConcurrentHashMap<String, Int>()
    private val learned = ConcurrentHashMap<String, Int>()
    // Bigramas: palabra previa -> {palabra siguiente -> veces}. Alimenta "palabra siguiente".
    private val bigrams = ConcurrentHashMap<String, ConcurrentHashMap<String, Int>>()

    init {
        scope.launch(Dispatchers.IO) {
            loadFile("dict_es.txt", baseEs)
            loadFile("dict_en.txt", baseEn)
            loadPersisted()
        }
    }

    private fun loadFile(name: String, into: ConcurrentHashMap<String, Int>) {
        runCatching {
            appContext.assets.open(name).bufferedReader().useLines { lines ->
                var i = 0
                lines.forEach { line ->
                    val w = line.trim().lowercase()
                    if (w.isNotEmpty()) {
                        val weight = maxOf(1, BASE_TOP - i)
                        into.merge(w, weight) { a, b -> maxOf(a, b) }
                        i++
                    }
                }
            }
        }
    }

    private suspend fun loadPersisted() {
        runCatching {
            store.data.first().asMap().forEach { (key, value) ->
                val name = key.name
                if (value !is Int) return@forEach
                when {
                    name.startsWith(W) -> learned[name.removePrefix(W)] = value
                    name.startsWith(B) -> {
                        val rest = name.removePrefix(B)
                        val sep = rest.indexOf(SEP)
                        if (sep > 0 && sep < rest.length - 1) {
                            val prev = rest.substring(0, sep)
                            val next = rest.substring(sep + 1)
                            bigrams.getOrPut(prev) { ConcurrentHashMap() }[next] = value
                        }
                    }
                }
            }
        }
    }

    private fun inBase(w: String) = baseEs.containsKey(w) || baseEn.containsKey(w)

    /** Completa el prefijo en los idiomas activos; mayor peso primero, a igual peso la más corta. */
    fun suggest(prefix: String, useEs: Boolean, useEn: Boolean, limit: Int = 3): List<String> {
        val p = prefix.trim().lowercase()
        if (p.isEmpty()) return emptyList()

        val candidates = HashMap<String, Int>()
        fun scan(map: ConcurrentHashMap<String, Int>) {
            map.forEach { (w, weight) ->
                if (w.length > p.length && w.startsWith(p)) {
                    candidates[w] = maxOf(candidates[w] ?: 0, weight)
                }
            }
        }
        if (useEs) scan(baseEs)
        if (useEn) scan(baseEn)
        learned.forEach { (w, count) ->
            // Solo palabras nuevas ya "promovidas"; las del diccionario suman refuerzo.
            if (w.length > p.length && w.startsWith(p) && (count >= PROMOTE || inBase(w))) {
                candidates[w] = (candidates[w] ?: 0) + count * LEARN_INC
            }
        }
        return candidates.entries
            .sortedWith(compareByDescending<Map.Entry<String, Int>> { it.value }.thenBy { it.key.length })
            .map { it.key }
            .take(limit)
    }

    /** Palabras más probables tras [prev], según los bigramas aprendidos. */
    fun nextWords(prev: String, limit: Int = 3): List<String> {
        val a = clean(prev) ?: return emptyList()
        val map = bigrams[a] ?: return emptyList()
        return map.entries
            // Evita propagar erratas: la siguiente debe ser palabra conocida.
            .filter { inBase(it.key) || (learned[it.key] ?: 0) >= PROMOTE }
            .sortedByDescending { it.value }
            .map { it.key }
            .take(limit)
    }

    /** Aprende/refuerza una palabra escrita por el usuario (se promueve sola al llegar a PROMOTE). */
    fun learn(raw: String) {
        val w = clean(raw) ?: return
        learned.merge(w, 1) { a, _ -> a + 1 }
        scope.launch(Dispatchers.IO) {
            runCatching {
                store.edit { prefs ->
                    val key = intPreferencesKey(W + w)
                    prefs[key] = (prefs[key] ?: 0) + 1
                }
            }
        }
    }

    /** Registra que [next] siguió a [prev] (para la predicción de palabra siguiente). */
    fun learnBigram(prev: String, next: String) {
        val a = clean(prev) ?: return
        val b = clean(next) ?: return
        bigrams.getOrPut(a) { ConcurrentHashMap() }.merge(b, 1) { x, _ -> x + 1 }
        scope.launch(Dispatchers.IO) {
            runCatching {
                store.edit { prefs ->
                    val key = intPreferencesKey(B + a + SEP + b)
                    prefs[key] = (prefs[key] ?: 0) + 1
                }
            }
        }
    }

    private fun clean(raw: String): String? {
        val w = raw.trim().lowercase()
        return if (w.length >= 2 && w.all { it.isLetter() }) w else null
    }

    private companion object {
        const val BASE_TOP = 5000
        const val LEARN_INC = 400
        const val PROMOTE = 2
        const val W = "w:"
        const val B = "b:"
        const val SEP = '>'   // separador prev/next; seguro porque las palabras son solo letras
    }
}