package org.loculus.backend.service.submission

import mu.KotlinLogging
import org.jetbrains.exposed.sql.transactions.transaction
import org.loculus.backend.api.Organism
import org.loculus.backend.config.BackendConfig
import org.loculus.backend.service.submission.dbtables.CompressionDictionariesTable
import org.loculus.backend.service.submission.dbtables.CompressionDictionaryEntity
import org.loculus.backend.utils.DateProvider
import org.springframework.stereotype.Service
import java.security.MessageDigest
import java.util.concurrent.ConcurrentHashMap

private val log = KotlinLogging.logger { }

class DictEntry(val id: Int, val dict: ByteArray)

/**
 * An abstraction over the compression_dictionaries_table.
 * Caches the contents in memory to avoid repeated DB lookups.
 */
@Service
class CompressionDictService(private val backendConfig: BackendConfig, private val dateProvider: DateProvider) {
    private data class DictKey(val organism: Organism, val segmentOrGene: String)

    private data class DictCaches(
        val byOrganismAndName: Map<DictKey, DictEntry>,
        val unalignedByOrganism: Map<Organism, DictEntry>,
        val dictsById: ConcurrentHashMap<Int, ByteArray>,
    )

    private val caches: DictCaches by lazy {
        // Must be lazy to make sure the caches are only populated after the Flyway migration has run
        // The Spring bean is created before Flyway runs, i.e. we must not read from the DB when creating this class
        populateCaches()
    }

    /**
     * Main responsibility: Make sure that all dictionaries that might be used for new data exist in the database.
     *
     * Also, already populates the caches.
     */
    private fun populateCaches(): DictCaches {
        log.info { "Populating compression dictionary caches" }

        val dictCache = mutableMapOf<DictKey, DictEntry>()
        val unalignedDictCache = mutableMapOf<Organism, DictEntry>()
        val cacheById = ConcurrentHashMap<Int, ByteArray>()

        transaction {
            backendConfig.organisms.forEach { (organismString, instanceConfig) ->
                val organism = Organism(organismString)
                val nucleotideSequences = instanceConfig.referenceGenome.nucleotideSequences
                val genes = instanceConfig.referenceGenome.genes
                for (referenceSequence in nucleotideSequences + genes) {
                    val reference = referenceSequence.sequence
                    val description = "${organism.name} - ${referenceSequence.name}"
                    val dictId = getDictIdOrInsertNewEntry(reference, description)

                    val dict = reference.toByteArray()

                    val dictKey = DictKey(organism = organism, segmentOrGene = referenceSequence.name)
                    dictCache[dictKey] = DictEntry(dictId, dict)
                    cacheById[dictId] = dict
                }

                // No need to concatenate nothing, it'll never be used
                if (nucleotideSequences.isNotEmpty()) {
                    val concatenatedNucleotideSequences = nucleotideSequences
                        .map { it.sequence }
                        .sortedBy { it }
                        .joinToString("")
                    val dictId = getDictIdOrInsertNewEntry(
                        concatenatedNucleotideSequences,
                        "${organism.name} - concatenated nucleotide sequences",
                    )
                    val unalignedDict = concatenatedNucleotideSequences.toByteArray()
                    val dictEntry = DictEntry(dictId, unalignedDict)

                    unalignedDictCache[organism] = dictEntry
                    cacheById[dictId] = unalignedDict
                }
            }
        }

        log.info {
            "Populated compression dictionary caches: ${dictCache.size} byOrganismAndName, " +
                "${unalignedDictCache.size} unalignedByOrganism, ${cacheById.size} dictsById"
        }

        return DictCaches(
            byOrganismAndName = dictCache.toMap(),
            unalignedByOrganism = unalignedDictCache.toMap(),
            dictsById = cacheById,
        )
    }

    /**
     * Get dictionary for a specific segment or gene (used when compressing processed sequences)
     */
    fun getDictForSegmentOrGene(organism: Organism, segmentOrGene: String): DictEntry? =
        caches.byOrganismAndName[DictKey(organism = organism, segmentOrGene = segmentOrGene)]

    /**
     * Get dictionary for unaligned sequences (used when compressing submitted sequences)
     */
    fun getDictForUnalignedSequence(organism: Organism): DictEntry? = caches.unalignedByOrganism[organism]

    /**
     * Get dictionary by ID (used when decompressing sequences)
     */
    fun getDictById(id: Int): ByteArray {
        val cachedDict = caches.dictsById[id]
        if (cachedDict != null) {
            return cachedDict
        }

        return transaction {
            val dict = CompressionDictionaryEntity.findById(id)
                ?.dictContents
                ?: throw RuntimeException("Did not find compression dictionary with id $id")
            caches.dictsById[id] = dict
            dict
        }
    }

    private fun getDictIdOrInsertNewEntry(dict: String, description: String): Int {
        val hash = computeHash(dict)

        val existingId = CompressionDictionaryEntity.find { CompressionDictionariesTable.hashColumn eq hash }
            .firstOrNull()
            ?.id
            ?.value

        if (existingId != null) {
            return existingId
        }

        return CompressionDictionaryEntity
            .new {
                this.hash = hash
                this.dictContents = dict.toByteArray()
                this.description = description
                this.createdAt = dateProvider.getCurrentDateTime()
            }
            .also {
                log.debug {
                    "Inserted new dict entry: id ${it.id.value}, $description, for dict ${dict.substring(0..10)}..."
                }
            }
            .id
            .value
    }

    private fun computeHash(input: String): String {
        val hashFunction = MessageDigest.getInstance("SHA-256")
        val digest = hashFunction.digest(input.toByteArray())
        return digest.joinToString("") { "%02x".format(it) }
    }
}
