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. Tres niveles, de mejor a peor:
     *  1. prefijo exacto con acentos (`info` → `informe`)              → +[EXACT_BONUS]
     *  2. prefijo sin acentos (`informacion` → `información`)          → peso base
     *  3. fuzzy: 1 errata en el prefijo (`infrome` → `informe`),
     *     solo si aún faltan sugerencias                              → −[FUZZY_PENALTY]
     * 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 np = fold(p)

        val candidates = HashMap<String, Int>()
        fun consider(w: String, weight: Int) {
            if (w.length <= p.length) return
            val score = when {
                w.startsWith(p) -> weight + EXACT_BONUS
                fold(w).startsWith(np) -> weight
                else -> return
            }
            if (score > (candidates[w] ?: Int.MIN_VALUE)) candidates[w] = score
        }
        fun scan(map: ConcurrentHashMap<String, Int>) = map.forEach { (w, weight) -> consider(w, 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 && (count >= PROMOTE || inBase(w)) &&
                (w.startsWith(p) || fold(w).startsWith(np))
            ) {
                candidates[w] = (candidates[w] ?: 0) + count * LEARN_INC
            }
        }

        // Nivel 3: tolera 1 errata en el prefijo cuando hay pocas coincidencias limpias.
        if (candidates.size < limit && np.length >= 3) {
            fun scanFuzzy(map: ConcurrentHashMap<String, Int>) {
                map.forEach { (w, weight) ->
                    if (w.length > p.length && w !in candidates) {
                        val nw = fold(w)
                        if (within1(np, nw.take(np.length)) || within1(np, nw.take(np.length + 1))) {
                            candidates[w] = weight - FUZZY_PENALTY
                        }
                    }
                }
            }
            if (useEs) scanFuzzy(baseEs)
            if (useEn) scanFuzzy(baseEn)
        }

        return candidates.entries
            .sortedWith(compareByDescending<Map.Entry<String, Int>> { it.value }.thenBy { it.key.length })
            .map { it.key }
            .take(limit)
    }

    /** Quita acentos/diéresis y plega ñ→n, ç→c para comparar de forma laxa. */
    private fun fold(s: String): String = buildString(s.length) {
        for (c in s) append(
            when (c) {
                'á', 'à', 'ä', 'â' -> 'a'
                'é', 'è', 'ë', 'ê' -> 'e'
                'í', 'ì', 'ï', 'î' -> 'i'
                'ó', 'ò', 'ö', 'ô' -> 'o'
                'ú', 'ù', 'ü', 'û' -> 'u'
                'ñ' -> 'n'
                'ç' -> 'c'
                else -> c
            },
        )
    }

    /** ¿[a] se transforma en [b] con como mucho una edición (sustituir/insertar/borrar)? */
    private fun within1(a: String, b: String): Boolean {
        val la = a.length
        val lb = b.length
        if (la - lb > 1 || lb - la > 1) return false
        var i = 0
        var j = 0
        var edits = 0
        while (i < la && j < lb) {
            if (a[i] == b[j]) {
                i++; j++
            } else {
                if (++edits > 1) return false
                when {
                    la > lb -> i++
                    la < lb -> j++
                    else -> { i++; j++ }
                }
            }
        }
        return edits + (la - i) + (lb - j) <= 1
    }

    /** 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 = 20000      // ~tamaño de cada diccionario; da pesos descendentes por frecuencia
        const val LEARN_INC = 2000      // cada uso pesa como subir ~2000 puestos de frecuencia
        const val EXACT_BONUS = 50000   // un prefijo exacto (con acentos) gana a cualquier base
        const val FUZZY_PENALTY = 100000 // una coincidencia con errata siempre va por debajo de las limpias
        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
    }
}