package com.rosetta.messenger.network.connection

import com.rosetta.messenger.network.PacketSearch
import com.rosetta.messenger.network.SearchUser
import kotlinx.coroutines.suspendCancellableCoroutine
import kotlinx.coroutines.withTimeout
import java.util.LinkedHashSet
import java.util.Locale
import java.util.concurrent.ConcurrentHashMap
import kotlin.coroutines.resume

class PacketRouter(
    private val sendSearchPacket: (PacketSearch) -> Unit,
    private val privateHashProvider: () -> String?
) {
    private val userInfoCache = ConcurrentHashMap<String, SearchUser>()
    private val pendingResolves =
        ConcurrentHashMap<String, MutableList<kotlinx.coroutines.CancellableContinuation<SearchUser?>>>()
    private val pendingSearchQueries =
        ConcurrentHashMap<String, MutableList<kotlinx.coroutines.CancellableContinuation<List<SearchUser>>>>()

    private fun normalizeSearchQuery(value: String): String =
        value.trim().removePrefix("@").lowercase(Locale.ROOT)

    suspend fun onSearchPacket(packet: PacketSearch, onUserDiscovered: suspend (SearchUser) -> Unit) {
        if (packet.users.isNotEmpty()) {
            packet.users.forEach { user ->
                val normalizedUserPublicKey = user.publicKey.trim()
                userInfoCache[normalizedUserPublicKey] = user

                pendingResolves
                    .keys
                    .filter { it.equals(normalizedUserPublicKey, ignoreCase = true) }
                    .forEach { key ->
                        pendingResolves.remove(key)?.forEach { cont ->
                            try {
                                cont.resume(user)
                            } catch (_: Exception) {}
                        }
                    }

                onUserDiscovered(user)
            }
        }

        if (packet.search.isNotEmpty() && packet.users.none { it.publicKey == packet.search }) {
            pendingResolves.remove(packet.search)?.forEach { cont ->
                try {
                    cont.resume(null)
                } catch (_: Exception) {}
            }
        }

        if (packet.search.isNotEmpty()) {
            val rawQuery = packet.search.trim()
            val normalizedQuery = normalizeSearchQuery(rawQuery)
            val continuations =
                LinkedHashSet<kotlinx.coroutines.CancellableContinuation<List<SearchUser>>>()

            fun collectByKey(key: String) {
                if (key.isEmpty()) return
                pendingSearchQueries.remove(key)?.let { continuations.addAll(it) }
            }

            collectByKey(rawQuery)
            if (normalizedQuery.isNotEmpty() && normalizedQuery != rawQuery) {
                collectByKey(normalizedQuery)
            }

            if (continuations.isEmpty()) {
                val matchedByQuery =
                    pendingSearchQueries.keys.firstOrNull { pendingKey ->
                        pendingKey.equals(rawQuery, ignoreCase = true) ||
                            normalizeSearchQuery(pendingKey) == normalizedQuery
                    }
                if (matchedByQuery != null) collectByKey(matchedByQuery)
            }

            if (continuations.isEmpty() && packet.users.isNotEmpty()) {
                val responseUsernames =
                    packet.users
                        .map { normalizeSearchQuery(it.username) }
                        .filter { it.isNotEmpty() }
                        .toSet()
                if (responseUsernames.isNotEmpty()) {
                    val matchedByUsers =
                        pendingSearchQueries.keys.firstOrNull { pendingKey ->
                            val normalizedPending = normalizeSearchQuery(pendingKey)
                            normalizedPending.isNotEmpty() &&
                                responseUsernames.contains(normalizedPending)
                        }
                    if (matchedByUsers != null) collectByKey(matchedByUsers)
                }
            }

            continuations.forEach { cont ->
                try {
                    cont.resume(packet.users)
                } catch (_: Exception) {}
            }
        }
    }

    fun getCachedUserName(publicKey: String): String? {
        val cached = userInfoCache[publicKey] ?: return null
        return cached.title.ifEmpty { cached.username }.ifEmpty { null }
    }

    fun getCachedUserInfo(publicKey: String): SearchUser? = userInfoCache[publicKey]

    fun getCachedUserByUsername(username: String): SearchUser? {
        val normalizedUsername = normalizeSearchQuery(username)
        if (normalizedUsername.isEmpty()) return null
        return userInfoCache.values.firstOrNull { cached ->
            normalizeSearchQuery(cached.username) == normalizedUsername
        }
    }

    suspend fun resolveUserName(publicKey: String, timeoutMs: Long = 3000): String? {
        if (publicKey.isEmpty()) return null

        userInfoCache[publicKey]?.let { cached ->
            val name = cached.title.ifEmpty { cached.username }
            if (name.isNotEmpty()) return name
        }

        val privateHash = privateHashProvider()?.takeIf { it.isNotBlank() } ?: return null

        return try {
            withTimeout(timeoutMs) {
                suspendCancellableCoroutine { cont ->
                    pendingResolves.getOrPut(publicKey) { mutableListOf() }.add(cont)

                    cont.invokeOnCancellation {
                        pendingResolves[publicKey]?.remove(cont)
                        if (pendingResolves[publicKey]?.isEmpty() == true) {
                            pendingResolves.remove(publicKey)
                        }
                    }

                    val packet = PacketSearch().apply {
                        this.privateKey = privateHash
                        this.search = publicKey
                    }
                    sendSearchPacket(packet)
                }
            }?.let { user -> user.title.ifEmpty { user.username }.ifEmpty { null } }
        } catch (_: Exception) {
            pendingResolves.remove(publicKey)
            null
        }
    }

    suspend fun resolveUserInfo(publicKey: String, timeoutMs: Long = 3000): SearchUser? {
        if (publicKey.isEmpty()) return null

        userInfoCache[publicKey]?.let { return it }
        val privateHash = privateHashProvider()?.takeIf { it.isNotBlank() } ?: return null

        return try {
            withTimeout(timeoutMs) {
                suspendCancellableCoroutine { cont ->
                    pendingResolves.getOrPut(publicKey) { mutableListOf() }.add(cont)
                    cont.invokeOnCancellation {
                        pendingResolves[publicKey]?.remove(cont)
                        if (pendingResolves[publicKey]?.isEmpty() == true) {
                            pendingResolves.remove(publicKey)
                        }
                    }
                    val packet = PacketSearch().apply {
                        this.privateKey = privateHash
                        this.search = publicKey
                    }
                    sendSearchPacket(packet)
                }
            }
        } catch (_: Exception) {
            pendingResolves.remove(publicKey)
            null
        }
    }

    suspend fun searchUsers(query: String, timeoutMs: Long = 3000): List<SearchUser> {
        val normalizedQuery = normalizeSearchQuery(query)
        if (normalizedQuery.isEmpty()) return emptyList()

        val privateHash = privateHashProvider()?.takeIf { it.isNotBlank() } ?: return emptyList()

        val cachedMatches =
            userInfoCache.values.filter { cached ->
                normalizeSearchQuery(cached.username) == normalizedQuery && cached.publicKey.isNotBlank()
            }
        if (cachedMatches.isNotEmpty()) {
            return cachedMatches.distinctBy { it.publicKey }
        }

        return try {
            withTimeout(timeoutMs) {
                suspendCancellableCoroutine { cont ->
                    pendingSearchQueries.getOrPut(normalizedQuery) { mutableListOf() }.add(cont)

                    cont.invokeOnCancellation {
                        pendingSearchQueries[normalizedQuery]?.remove(cont)
                        if (pendingSearchQueries[normalizedQuery]?.isEmpty() == true) {
                            pendingSearchQueries.remove(normalizedQuery)
                        }
                    }

                    val packet = PacketSearch().apply {
                        this.privateKey = privateHash
                        this.search = normalizedQuery
                    }
                    sendSearchPacket(packet)
                }
            }
        } catch (_: Exception) {
            pendingSearchQueries.remove(normalizedQuery)
            emptyList()
        }
    }
}
