Suggestion

app/src/main/java/xyz/quaver/pupil/networking/HitomiHttpClient.kt

@@ -13,7 +13,9 @@ import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.withContext
 import xyz.quaver.pupil.hitomi.max_node_size
 import java.nio.ByteBuffer
+import java.nio.ByteOrder
 import java.nio.IntBuffer
+import java.nio.charset.Charset
 
 const val domain = "ltn.hitomi.la"
 const val galleryBlockExtension = ".html"
@@ -28,6 +30,11 @@ const val galleriesIndexDir = "galleriesindex"
 const val languagesIndexDir = "languagesindex"
 const val nozomiURLIndexDir = "nozomiurlindex"
 
+data class Suggestion(
+    val tag: SearchQuery.Tag,
+    val count: Int
+)
+
 fun IntBuffer.toSet(): Set<Int> {
     val result = mutableSetOf<Int>()
 
@@ -78,7 +85,7 @@ class HitomiHttpClient {
         }
 
         return Node.decodeNode(
-            getURLAtRange(url, address until (address+max_node_size))
+            getURLAtRange(url, address ..< address+max_node_size)
         )
     }
 
@@ -112,7 +119,38 @@ class HitomiHttpClient {
         return getURLAtRange(url, offset until (offset+length)).asIntBuffer()
     }
 
-    suspend fun getGalleryIDsFromNozomi(
+    private suspend fun getSuggestionsFromData(field: String, data: Node.Data): List<Suggestion> {
+        val url = "https://$domain/$indexDir/$field.${getTagIndexVersion()}.data"
+        val (offset, length) = data
+
+        check(data.length in 1..10000) { "Invalid length ${data.length}" }
+
+        val buffer = getURLAtRange(url, offset..<offset+length).order(ByteOrder.BIG_ENDIAN)
+
+        val numberOfSuggestions = buffer.int
+
+        check(numberOfSuggestions in 1 .. 100) { "Number of suggestions $numberOfSuggestions is too long" }
+
+        return buildList {
+            for (i in 0 ..< numberOfSuggestions) {
+                val namespaceLen = buffer.int
+                val namespace = ByteArray(namespaceLen).apply {
+                    buffer.get(this)
+                }.toString(charset("UTF-8"))
+
+                val tagLen = buffer.int
+                val tag = ByteArray(tagLen).apply {
+                    buffer.get(this)
+                }.toString(charset("UTF-8"))
+
+                val count = buffer.int
+
+                add(Suggestion(SearchQuery.Tag(namespace, tag), count))
+            }
+        }
+    }
+
+    private suspend fun getGalleryIDsFromNozomi(
         area: String?,
         tag: String,
         language: String
@@ -132,7 +170,7 @@ class HitomiHttpClient {
         return ByteBuffer.wrap(result).asIntBuffer()
     }
 
-    suspend fun getGalleryIDsForQuery(query: SearchQuery.Tag, language: String = "all"): IntBuffer = when (query.namespace) {
+    private suspend fun getGalleryIDsForQuery(query: SearchQuery.Tag, language: String = "all"): IntBuffer = when (query.namespace) {
         "female", "male" -> getGalleryIDsFromNozomi("tag", query.toString(), language)
         "language" -> getGalleryIDsFromNozomi(null, "index", query.tag)
         null -> {
@@ -146,14 +184,24 @@ class HitomiHttpClient {
         else -> getGalleryIDsFromNozomi(query.namespace, query.tag, language)
     }
 
-    suspend fun search(query: SearchQuery?): Set<Int> = when (query) {
+    suspend fun getSuggestionsForQuery(query: SearchQuery.Tag): Result<List<Suggestion>> = runCatching {
+        val field = query.namespace ?: "global"
+        val key = Node.Key(field)
+        val node = getNodeAtAddress(field, 0)
+        val data = bSearch(field, key, node)
+
+        data?.let { getSuggestionsFromData(field, data) } ?: emptyList()
+    }
+
+    suspend fun search(query: SearchQuery?): Result<Set<Int>> = runCatching {
+        when (query) {
             is SearchQuery.Tag -> getGalleryIDsForQuery(query).toSet()
             is SearchQuery.Not -> coroutineScope {
                 val allGalleries = async {
                     getGalleryIDsFromNozomi(null, "index", "all")
                 }
 
-            val queriedGalleries = search(query.query)
+                val queriedGalleries = search(query.query).getOrThrow()
 
                 val result = mutableSetOf<Int>()
 
@@ -172,7 +220,7 @@ class HitomiHttpClient {
             is SearchQuery.And -> coroutineScope {
                 val queries = query.queries.map { query ->
                     async {
-                    search(query)
+                        search(query).getOrThrow()
                     }
                 }
 
@@ -189,7 +237,7 @@ class HitomiHttpClient {
             is SearchQuery.Or -> coroutineScope {
                 val queries = query.queries.map { query ->
                     async {
-                    search(query)
+                        search(query).getOrThrow()
                     }
                 }
 
@@ -204,4 +252,5 @@ class HitomiHttpClient {
             }
             null -> getGalleryIDsFromNozomi(null, "index", "all").toSet()
         }
+    }
 }