package theorycrafter.utils

import kotlinx.coroutines.currentCoroutineContext
import kotlinx.coroutines.ensureActive
import java.util.*
import kotlin.math.max


/**
 * A utility for searching in a set items by a string.
 */
class StringSearch<T: Any>(


    /**
     * The minimum number of input characters that causes an actual search.
     *
     * When the search query contains fewer significant characters than this, a `null` value is returned.
     */
    private val minCharacters: Int = 1,


    /**
     * Whether the search will be case-sensitive.
     */
    private val isCaseSensitive: Boolean = false,


    /**
     * An optional function to sort suggestions with identical scores.
     */
    private val suggestionComparator: Comparator<in T>? = null,


    /**
     * Returns whether the given character delimits search terms.
     */
    private val isTermDelimiter: (Char) -> Boolean = { it == ' ' },


    /**
     * Returns whether the given character is to be filtered completely from the search terms and the query.
     */
    private val isFiltered: (Char) -> Boolean = { false },


) {


    /**
     * The trie assisting our search.
     */
    private val searchTrie = SearchTrie<SearchItem<T>>()


    /**
     * A buffer of [Char] we (re)use to avoid allocating one when building acronyms.
     */
    private val charArrayBuffer = CharArray(20)


    /**
     * Adds an item to the search, and returns an [AddedItem] which can be used to remove it.
     */
    fun addItem(
        item: T,
        text: String,
        config: SearchConfig = DefaultSearchConfig
    ): AddedItem {
        val searchText = if (isCaseSensitive) text else text.lowercase(Locale.ROOT)
        val terms = searchText.splitToTerms(isDelimiter = isTermDelimiter, isFiltered = isFiltered)

        val searchItemsAndTerms = mutableListOf<Pair<SearchItem<T>, String>>()
        fun addSearchItem(searchItem: SearchItem<T>, term: String) {
            searchTrie.add(searchItem, term)
            searchItemsAndTerms.add(Pair(searchItem, term))
        }

        // The term
        for (index in terms.indices) {
            val term = terms[index]
            val score = (
                config.wholeTermScore
                    - config.earlierTermsPreferenceFactor * index +
                    - config.shorterTermsPreferenceFactor * term.length
            )
            addSearchItem(SearchItem(item, score = score), term)
        }

        // The acronym
        if (config.includeAcronyms && (terms.size > 1)) {
            val acronym = terms
                .flatMap {
                    val alternate = config.alternateAcronymTerms[it]
                    if (alternate != null)
                        sequenceOf(it, alternate)
                    else
                        sequenceOf(it)
                }
                .map { it.first() }
                .joinToStringUsing(charArrayBuffer)

            val searchItem = SearchItem(item, score = config.acronymScore)
            for (subsequence in acronym.subsequences(minCharacters))
                addSearchItem(searchItem, subsequence)
        }

        return AddedItemImpl(searchItemsAndTerms)
    }


    /**
     * Queries for items using the given string.
     *
     * Returns `null` if the query is shorter than [minCharacters].
     */
    suspend fun query(query: String): List<T>? {
        val searchQuery = if (isCaseSensitive) query else query.lowercase(Locale.ROOT)

        val significantCharacters = searchQuery.count { !isTermDelimiter(it) && !isFiltered(it) }
        if (significantCharacters < minCharacters)
            return null

        // Split the query into terms, find all the matches for each term and then intersect all the matches,
        // so we are left with the set of items that it match all terms.

        val terms = searchQuery.splitToTerms(isDelimiter = isTermDelimiter, isFiltered = isFiltered)

        val searchResultComparator: Comparator<Map.Entry<T, Int>> =
            if (suggestionComparator == null)
                compareBy { 0 }
            else
                Comparator<Map.Entry<T, Int>> { o1, o2 ->
                    suggestionComparator.compare(o1.key, o2.key)
                }

        suspend fun <T> T.ensureNotCancelled() = this.also { currentCoroutineContext().ensureActive() }

        return terms
            .map { term ->
                ensureNotCancelled()
                searchTrie.search(term)
            }
            .map { matches ->
                ensureNotCancelled()
                // For each word that matched, take the highest score only
                matches
                    .groupingBy { it.item }
                    .fold(
                        initialValueSelector = { _, searchItem -> searchItem.score },
                    ) { _, bestScore, searchItem ->
                        max(searchItem.score, bestScore)
                    }
                    .entries
            }
            .ensureNotCancelled()
            .flatten()
            .ensureNotCancelled()
            .groupBy { it.key }
            .ensureNotCancelled()
            .filterValues { it.size == terms.size }  // Remove items that didn't match every word
            .ensureNotCancelled()
            .mapValues { entry ->
                entry.value.sumOf { it.value }  // Sum the scores for each word
            }
            .ensureNotCancelled()
            .entries
            .sortedWith(
                compareByDescending<Map.Entry<T, Int>> { it.value }  // Sort by score
                    .then(searchResultComparator)                    // then using user-provided comparator
            )
            .ensureNotCancelled()
            .map {
                it.key
            }
    }



    /**
     * Encapsulates the information required to remove an item from the search.
     */
    interface AddedItem {


        /**
         * Removes the item.
         */
        fun remove()


    }


    /**
     * Implements [AddedItem].
     */
    private inner class AddedItemImpl(
        val termsAndSearchItems: List<Pair<SearchItem<T>, String>>
    ) : AddedItem {


        /**
         * Whether this [AddedItem] has already been used to remove the item.
         */
        private var removed = false


        override fun remove() {
            if (removed)
                throw IllegalStateException("Item has already been removed")

            for ((searchItem, term) in termsAndSearchItems) {
                searchTrie.remove(searchItem, term)
            }
        }


    }


}


/**
 * Returns an optimized function that, given a [Char], returns whether it is in the given string.
 */
fun String.containsFunction(): (Char) -> Boolean {
    when {
        this.isEmpty() -> return { _ -> false }
        this.length == 1 -> {
            val c = this[0]
            return { it == c }
        }
        this.length < 5 -> return { indexOf(it) >= 0 }
        else -> {
            val charSet = this.mapTo(mutableSetOf()){ it }
            return charSet::contains
        }
    }
}


/**
 * A search item and its score.
 */
private data class SearchItem<out T>(
    val item: T,
    val score: Int
)


/**
 * Converts a sequence of characters to a [String], using the given buffer if possible.
 */
private fun List<Char>.joinToStringUsing(buffer: CharArray): String {
    val charArray = if (size <= buffer.size) buffer else CharArray(size)
    this.forEachIndexed { index, c ->
        charArray[index] = c
    }
    return String(charArray, 0, size)
}


/**
 * Returns all the given string's sub-sequences (in the mathematical sense - the characters need not be consecutive).
 */
private fun String.subsequences(minLength: Int): Sequence<String> {
    if (length > 32)
        throw IllegalArgumentException("String length should be less then 32")
    return sequence {
        val string = this@subsequences
        val len = string.length
        val chars = CharArray(len)

        val max = 1 shl length
        for (bits in 1..max) {
            if (bits.countOneBits() < minLength)
                continue
            var ri = 0
            for (i in 0 until len){
                if ((bits and (1 shl i)) != 0)
                    chars[ri++] = string[i]
            }
            yield(String(chars, offset = 0, length = ri))
        }
    }
}


/**
 * Splits a string into search terms.
 */
private inline fun String.splitToTerms(
    isDelimiter: (Char) -> Boolean,
    isFiltered: (Char) -> Boolean,
): List<String> = buildList {
    val string = this@splitToTerms
    var startIndex = 0
    for (endIndex in 0..string.length){
        if ((endIndex == length) || isDelimiter(string[endIndex])) {
            if (endIndex > startIndex) {
                val term = string.substring(startIndex, endIndex = endIndex)
                    .filter { !isFiltered(it) }
                if (term.isNotBlank())
                    add(term)
            }
            startIndex = endIndex+1
        }
    }
}


/**
 * The search scoring parameters.
 */
data class SearchConfig (


    /**
     * Whether to include acronyms in the search.
     */
    val includeAcronyms: Boolean = false,


    /**
     * A map of alternate acronym terms.
     */
    val alternateAcronymTerms: Map<String, String> = emptyMap(),


    /**
     * The base score we assign to a [SearchItem] matching a prefix of a whole word/term of the search string.
     */
    val wholeTermScore: Int = 100_000,


    /**
     * The base score we assign to a [SearchItem] matching the acronym of the search string.
     */
    val acronymScore: Int = 1000,


    /**
     * The factor of the preference for matching terms closer to the beginning of the search string.
     *
     * This is multiplied by the index of the term and subtracted from the score.
     */
    val earlierTermsPreferenceFactor: Int = 50,


    /**
     * The factor of preference for matching shorter terms.
     *
     * This is multiplied by the length of the term and subtracted from the score.
     */
    val shorterTermsPreferenceFactor: Int = 1


)


/**
 * The default search configuration.
 */
val DefaultSearchConfig = SearchConfig()