package com.milaboratory.mixcr.bam

import cc.redberry.pipe.util.asSequence
import com.milaboratory.core.io.sequence.PairedRead
import com.milaboratory.core.io.sequence.SingleRead
import com.milaboratory.core.io.sequence.fastq.PairedFastqWriter
import com.milaboratory.core.io.sequence.fastq.SingleFastqWriter
import com.milaboratory.util.TempFileManager.systemTempFolderDestination
import io.kotest.matchers.shouldBe
import org.junit.Test
import java.io.BufferedReader
import java.io.File
import java.io.FileReader
import java.nio.file.Path
import java.nio.file.Paths
import kotlin.io.path.readLines
import kotlin.streams.asSequence

@Suppress("PrivatePropertyName")
class BAMReaderTest {
    private val classLoader = javaClass.classLoader

    // BAM file generated by the following STAR command:
    // ./STAR --runMode alignReads --outSAMunmapped Within --outSAMtype BAM SortedByCoordinate --readFilesCommand cat
    // --genomeDir {path_to_gen}/STARgenome --outFileNamePrefix bam_output/
    // --readFilesIn {path_to_fastq}/fastq1.fastq {path_to_fastq}/fastq2.fastq
    private val BY_COORD_BAM = asPath("bam/sortedByCoord.bam")

    // BAM file generated by the following STAR command:
    // ./STAR --runMode alignReads --outSAMunmapped Within --outSAMtype BAM Unsorted --readFilesCommand cat
    // --genomeDir {path_to_gen}/STARgenome --outFileNamePrefix bam_output/
    // --readFilesIn {path_to_fastq}/fastq1.fastq {path_to_fastq}/fastq2.fastq
    private val UNSORTED_BAM = asPath("bam/unsorted.bam")

    // BAM file generated by the following samtools command:
    // samtools sort -n unsorted.bam -o sortedByName.bam
    private val BY_NAME_BAM = asPath("bam/sortedByName.bam")

    // Corresponding initial fastq files
    private val FASTQ1 = asPath("bam/fastq1.fastq")
    private val FASTQ2 = asPath("bam/fastq2.fastq")

    // BAM files with paired and unpaired (or paired without mate in file) reads at the same time
    private val UNPAIRED_BY_COORD_BAM = asPath("bam/unpairedSortedByCoord.bam")
    private val UNPAIRED_BY_NAME_BAM = asPath("bam/unpairedSortedByName.bam")

    // Corresponding fastq files from samtools or bamUtil
    private val UNPAIRED_FASTQ1 = asPath("bam/unpairedFastq1.fastq")
    private val UNPAIRED_FASTQ2 = asPath("bam/unpairedFastq2.fastq")
    private val UNPAIRED_FASTQ = asPath("bam/unpairedFastq.fastq")

    private fun fastqFileHash(fastq1: Path): Int {
        val fastq1Reader = BufferedReader(FileReader(fastq1.toFile()))
        var result = 0
        fastq1Reader.lines().forEach { line ->
            result = result xor line.hashCode()
        }
        return result
    }

    private fun fastqFilesHash(fastq1: Path, fastq2: Path): Int {
        val fastq1Reader = BufferedReader(FileReader(fastq1.toFile()))
        val fastq2Reader = BufferedReader(FileReader(fastq2.toFile()))
        var result = 0
        fastq1Reader.lines().asSequence().zip(fastq2Reader.lines().asSequence()).forEach { (line1, line2) ->
            result = result xor (line1 + line2).hashCode()
        }
        return result
    }

    private fun checkPlural(
        readers: List<Path>,
        fastq1Name: List<Path>,
        fastq2Name: List<Path>,
        unpairedFastqName: List<Path>?
    ) {
        val tempFileDest = systemTempFolderDestination("test")
        BAMReader(
            readers,
            dropNonVDJChromosomes = false,
            replaceWildcards = false,
            tempFileDest,
            referenceForCram = null
        ).use { converter ->
            val resF1 = File.createTempFile("my_r1", ".fastq")
            val resF2 = File.createTempFile("my_r2", ".fastq")
            val resFu = File.createTempFile("my_ru", ".fastq")
            val maxId = PairedFastqWriter(resF1, resF2).use { wr ->
                SingleFastqWriter(resFu).use { swr ->
                    converter
                        .onEach { read ->
                            when (read) {
                                is PairedRead -> wr.write(read)
                                is SingleRead -> swr.write(read)
                            }
                        }
                        .map { it.id }
                        .asSequence().max()
                }
            }
            var targetHash = 0
            for (i in fastq1Name.indices) {
                val fastq1File = fastq1Name[i]
                val fastq2File = fastq2Name[i]
                targetHash = targetHash xor fastqFilesHash(fastq1File, fastq2File)
            }
            unpairedFastqName?.forEach { fastqFile ->
                targetHash = targetHash xor fastqFileHash(fastqFile)
            }
            val resultHash = fastqFilesHash(resF1.toPath(), resF2.toPath()) xor fastqFileHash(resFu.toPath())
            resultHash shouldBe targetHash
            val linesInFastq = fastq1Name.sumOf { it.readLines().size } +
                    (unpairedFastqName?.sumOf { it.readLines().size } ?: 0)
            maxId + 1 shouldBe linesInFastq / 4
        }
    }

    private fun asPath(file: String): Path =
        Paths.get(requireNotNull(classLoader.getResource(file)).path)

    @Test
    fun sortedByCoordTest() {
        checkPlural(listOf(BY_COORD_BAM), listOf(FASTQ1), listOf(FASTQ2), null)
    }

    @Test
    fun sortedByNameTest() {
        checkPlural(listOf(BY_NAME_BAM), listOf(FASTQ1), listOf(FASTQ2), null)
    }

    @Test
    fun unsortedTest() {
        checkPlural(listOf(UNSORTED_BAM), listOf(FASTQ1), listOf(FASTQ2), null)
    }

    @Test
    fun unpairedSortedByCoordTest() {
        checkPlural(
            listOf(UNPAIRED_BY_COORD_BAM),
            listOf(UNPAIRED_FASTQ1),
            listOf(UNPAIRED_FASTQ2),
            listOf(UNPAIRED_FASTQ)
        )
    }

    @Test
    fun unpairedSortedByNameTest() {
        checkPlural(
            listOf(UNPAIRED_BY_NAME_BAM),
            listOf(UNPAIRED_FASTQ1),
            listOf(UNPAIRED_FASTQ2),
            listOf(UNPAIRED_FASTQ)
        )
    }

    @Test
    fun combinedTest() {
        checkPlural(
            listOf(UNPAIRED_BY_NAME_BAM, BY_COORD_BAM),
            listOf(UNPAIRED_FASTQ1, FASTQ1),
            listOf(UNPAIRED_FASTQ2, FASTQ2),
            listOf(UNPAIRED_FASTQ)
        )
    }
}
