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. Mezcla un diccionario base (assets) con las palabras que el
 * usuario escribe (aprendidas y persistidas). El peso ordena las sugerencias.
 */
class SuggestionEngine(context: Context, private val scope: CoroutineScope) {

    private val appContext = context.applicationContext
    private val store = appContext.trazoLearnedStore
    private val freq = ConcurrentHashMap<String, Int>()

    init {
        scope.launch(Dispatchers.IO) {
            loadBase()
            loadLearned()
        }
    }

    private fun loadBase() {
        runCatching {
            appContext.assets.open("dict_es.txt").bufferedReader().useLines { lines ->
                var i = 0
                lines.forEach { line ->
                    val w = line.trim().lowercase()
                    if (w.isNotEmpty()) {
                        freq.merge(w, maxOf(1, BASE_TOP - i)) { a, b -> maxOf(a, b) }
                        i++
                    }
                }
            }
        }
    }

    private suspend fun loadLearned() {
        runCatching {
            store.data.first().asMap().forEach { (key, value) ->
                if (key.name.startsWith(W) && value is Int) {
                    val w = key.name.removePrefix(W)
                    freq.merge(w, value * LEARN_INC) { a, b -> a + b }
                }
            }
        }
    }

    /** Completa el prefijo con las palabras de mayor peso (más cortas primero a igual peso). */
    fun suggest(prefix: String, limit: Int = 3): List<String> {
        val p = prefix.trim().lowercase()
        if (p.isEmpty()) return emptyList()
        return freq.entries.asSequence()
            .filter { it.key.length > p.length && it.key.startsWith(p) }
            .sortedWith(compareByDescending<Map.Entry<String, Int>> { it.value }.thenBy { it.key.length })
            .map { it.key }
            .take(limit)
            .toList()
    }

    /** Aprende/refuerza una palabra escrita por el usuario. */
    fun learn(raw: String) {
        val w = raw.trim().lowercase()
        if (w.length < 2 || !w.all { it.isLetter() }) return
        freq.merge(w, LEARN_INC) { a, _ -> a + LEARN_INC }
        scope.launch(Dispatchers.IO) {
            runCatching {
                store.edit { prefs ->
                    val key = intPreferencesKey(W + w)
                    prefs[key] = (prefs[key] ?: 0) + 1
                }
            }
        }
    }

    private companion object {
        const val BASE_TOP = 5000
        const val LEARN_INC = 400
        const val W = "w:"
    }
}