package org.loculus.backend.service.submission

import com.github.luben.zstd.Zstd
import org.loculus.backend.api.GeneticSequence
import org.loculus.backend.api.Organism
import org.loculus.backend.api.ProcessedData
import org.loculus.backend.api.SubmittedData
import org.loculus.backend.config.BackendConfig
import org.springframework.stereotype.Service
import java.nio.charset.StandardCharsets
import java.util.Base64

data class CompressedSequence(val compressedSequence: String, val compressionDictId: Int?)

enum class CompressionAlgorithm(val extension: String) {
    NONE(""),
    ZSTD(".zst"),
    XZ(".xz"),
    GZIP(".gz"),
    ZIP(".zip"),
    BZIP2(".bz2"),
    LZMA(".lzma"),
}

@Service
class CompressionService(
    private val compressionDictService: CompressionDictService,
    private val backendConfig: BackendConfig,
) {

    fun compressSubmittedSequence(sequenceData: GeneticSequence, organism: Organism) = compress(
        sequenceData,
        compressionDictService.getDictForUnalignedSequence(organism),
    )

    private fun compressNucleotideSequence(
        uncompressedSequence: GeneticSequence,
        segmentName: String,
        organism: Organism,
    ): CompressedSequence = compress(
        uncompressedSequence,
        compressionDictService.getDictForSegmentOrGene(organism, segmentName),
    )

    private fun decompressNucleotideSequence(compressedSequence: CompressedSequence): GeneticSequence = decompress(
        compressedSequence,
    )

    private fun compressAminoAcidSequence(
        uncompressedSequence: String,
        gene: String,
        organism: Organism,
    ): CompressedSequence = compress(
        uncompressedSequence,
        compressionDictService.getDictForSegmentOrGene(organism, gene),
    )

    private fun decompressAminoAcidSequence(compressedSequence: CompressedSequence): GeneticSequence = decompress(
        compressedSequence,
    )

    fun decompressSequencesInSubmittedData(submittedData: SubmittedData<CompressedSequence>) = SubmittedData(
        submittedData.metadata,
        submittedData
            .unalignedNucleotideSequences.mapValues {
                when (val compressedSequence = it.value) {
                    null -> null
                    else -> decompressNucleotideSequence(compressedSequence)
                }
            },
        submittedData.files,
    )

    fun compressSequencesInSubmittedData(submittedData: SubmittedData<GeneticSequence>, organism: Organism) =
        SubmittedData(
            submittedData.metadata,
            submittedData
                .unalignedNucleotideSequences.mapValues { (_, sequenceData) ->
                    when (sequenceData) {
                        null -> null
                        else -> compressSubmittedSequence(sequenceData, organism)
                    }
                },
            submittedData.files,
        )

    fun decompressSequencesInProcessedData(processedData: ProcessedData<CompressedSequence>) = ProcessedData(
        processedData.metadata,
        processedData
            .unalignedNucleotideSequences.mapValues { (_, sequenceData) ->
                when (sequenceData) {
                    null -> null
                    else -> decompressNucleotideSequence(sequenceData)
                }
            },
        processedData.alignedNucleotideSequences.mapValues { (_, sequenceData) ->
            when (sequenceData) {
                null -> null
                else -> decompressNucleotideSequence(sequenceData)
            }
        },
        processedData.nucleotideInsertions,
        processedData.alignedAminoAcidSequences.mapValues { (_, sequenceData) ->
            when (sequenceData) {
                null -> null
                else -> decompressAminoAcidSequence(sequenceData)
            }
        },
        processedData.aminoAcidInsertions,
        processedData.sequenceNameToFastaId,
        processedData.files,
    )

    fun compressSequencesInProcessedData(processedData: ProcessedData<String>, organism: Organism) = ProcessedData(
        processedData.metadata,
        processedData
            .unalignedNucleotideSequences.mapValues { (segmentName, sequenceData) ->
                when (sequenceData) {
                    null -> null
                    else -> compressNucleotideSequence(sequenceData, segmentName, organism)
                }
            },
        processedData.alignedNucleotideSequences.mapValues { (segmentName, sequenceData) ->
            when (sequenceData) {
                null -> null
                else -> compressNucleotideSequence(sequenceData, segmentName, organism)
            }
        },
        processedData.nucleotideInsertions,
        processedData.alignedAminoAcidSequences.mapValues { (gene, sequenceData) ->
            when (sequenceData) {
                null -> null
                else -> compressAminoAcidSequence(sequenceData, gene, organism)
            }
        },
        processedData.aminoAcidInsertions,
        processedData.sequenceNameToFastaId,
        processedData.files,
    )

    private fun compress(sequence: GeneticSequence, dictEntry: DictEntry?): CompressedSequence {
        val input = sequence.toByteArray(StandardCharsets.UTF_8)
        val compressBound = Zstd.compressBound(input.size.toLong()).toInt()
        val outputBuffer = ByteArray(compressBound)

        val compressionReturnCode: Long = if (dictEntry == null) {
            Zstd.compress(outputBuffer, input, backendConfig.zstdCompressionLevel)
        } else {
            Zstd.compress(outputBuffer, input, dictEntry.dict, backendConfig.zstdCompressionLevel)
        }

        if (Zstd.isError(compressionReturnCode)) {
            throw RuntimeException("Zstd compression failed: error code $compressionReturnCode")
        }

        return CompressedSequence(
            compressedSequence = Base64.getEncoder()
                .encodeToString(outputBuffer.copyOfRange(0, compressionReturnCode.toInt())),
            compressionDictId = dictEntry?.id,
        )
    }

    private fun decompress(compressedSequence: CompressedSequence): String {
        val compressed = Base64.getDecoder().decode(compressedSequence.compressedSequence)
        val decompressedSize = Zstd.getFrameContentSize(compressed)
        if (Zstd.isError(decompressedSize)) {
            throw RuntimeException("reading Zstd decompressed size failed: error code $decompressedSize")
        }

        val decompressedBuffer = ByteArray(decompressedSize.toInt())
        val decompressionReturnCode: Long = if (compressedSequence.compressionDictId == null) {
            Zstd.decompress(decompressedBuffer, compressed)
        } else {
            val dictionary = compressionDictService.getDictById(compressedSequence.compressionDictId)
            Zstd.decompress(decompressedBuffer, compressed, dictionary)
        }
        if (Zstd.isError(decompressionReturnCode)) {
            throw RuntimeException("Zstd decompression failed: error code $decompressionReturnCode")
        }
        return String(decompressedBuffer, 0, decompressionReturnCode.toInt(), StandardCharsets.UTF_8)
    }
}
