/*
 * Copyright (c) 2014-2024, MiLaboratories Inc. All Rights Reserved
 *
 * Before downloading or accessing the software, please read carefully the
 * License Agreement available at:
 * https://github.com/milaboratory/mixcr/blob/develop/LICENSE
 *
 * By downloading or accessing the software, you accept and agree to be bound
 * by the terms of the License Agreement. If you do not want to agree to the terms
 * of the Licensing Agreement, you must not download or access the software.
 */
package com.milaboratory.mixcr.cli

import cc.redberry.pipe.OutputPort
import cc.redberry.pipe.util.flatten
import com.milaboratory.app.InputFileType
import com.milaboratory.app.ValidationException
import com.milaboratory.app.logger
import com.milaboratory.cli.POverridesBuilderOps
import com.milaboratory.mitool.tag.TagType
import com.milaboratory.mixcr.basictypes.ClnAReader
import com.milaboratory.mixcr.basictypes.ClnAWriter
import com.milaboratory.mixcr.basictypes.ClnsReader
import com.milaboratory.mixcr.basictypes.ClnsWriter
import com.milaboratory.mixcr.basictypes.Clone
import com.milaboratory.mixcr.basictypes.CloneReader
import com.milaboratory.mixcr.basictypes.CloneSet
import com.milaboratory.mixcr.basictypes.IOUtil
import com.milaboratory.mixcr.basictypes.IOUtil.MiXCRFileType.CLNA
import com.milaboratory.mixcr.basictypes.IOUtil.MiXCRFileType.CLNS
import com.milaboratory.mixcr.basictypes.MiXCRHeader
import com.milaboratory.mixcr.basictypes.VDJCAlignments
import com.milaboratory.mixcr.clonegrouping.CloneGroupingParams.Companion.mkGrouper
import com.milaboratory.mixcr.presets.AnalyzeCommandDescriptor
import com.milaboratory.mixcr.presets.AssembleContigsMixins
import com.milaboratory.mixcr.presets.MiXCRParamsBundle
import com.milaboratory.mixcr.util.Concurrency
import com.milaboratory.util.ReportUtil
import com.milaboratory.util.SmartProgressReporter
import com.milaboratory.util.TempFileManager
import io.repseq.core.VDJCLibraryRegistry
import picocli.CommandLine.Command
import picocli.CommandLine.Mixin
import picocli.CommandLine.Option
import picocli.CommandLine.Parameters
import java.nio.file.Path

object CommandAssembleCells {
    const val COMMAND_NAME = AnalyzeCommandDescriptor.assembleCells.name

    abstract class CmdBase : MiXCRCommandWithOutputs(), MiXCRPresetAwareCommand<CommandAssembleCellsParams> {
        @Option(
            names = ["-O"],
            description = ["Overrides for the clone grouping parameters."],
            paramLabel = CommonDescriptions.Labels.OVERRIDES,
            order = OptionsOrder.overrides
        )
        private var overrides: Map<String, String> = mutableMapOf()

        override val paramsResolver =
            object : MiXCRParamsResolver<CommandAssembleCellsParams>(MiXCRParamsBundle::assembleCells) {
                override fun POverridesBuilderOps<CommandAssembleCellsParams>.paramsOverrides() {
                    CommandAssembleCellsParams::algorithm jsonOverrideWith overrides
                }
            }
    }

    @Command(
        description = ["Group clones by cells. Required data with cell tags. All clones should be fully covered by the same feature."]
    )
    class Cmd : CmdBase() {
        @Parameters(
            description = ["Path to input file."],
            paramLabel = "clones.(clns|clna)",
            index = "0"
        )
        lateinit var inputFile: Path

        @Parameters(
            description = ["Path where to write output. Will have the same file type."],
            paramLabel = "grouped.(clns|clna)",
            index = "1"
        )
        lateinit var outputFile: Path

        @Mixin
        lateinit var reportOptions: ReportOptions

        @Mixin
        lateinit var resetPreset: ResetPresetOptions

        @Mixin
        lateinit var dontSavePresetOption: DontSavePresetOption

        @Mixin
        lateinit var useLocalTemp: UseLocalTempOption

        override val inputFiles
            get() = listOf(inputFile)

        override val outputFiles
            get() = listOf(outputFile)

        override fun validate() {
            ValidationException.requireTheSameFileType(
                inputFile,
                outputFile,
                InputFileType.CLNS,
                InputFileType.CLNA
            )
        }

        override fun run1() {
            val report = when (val fileType = IOUtil.extractFileType(inputFile)) {
                CLNS -> processClns()
                CLNA -> processClna()
                else -> throw ValidationException("Unsupported input type $fileType")
            }
            ReportUtil.writeReportToStdout(report)
            reportOptions.appendToFiles(report)
        }

        private fun validate(reader: CloneReader) {
            ValidationException.require(reader.header.allFullyCoveredBy != null) {
                "Input were processed by `${CommandAssembleContigs.COMMAND_NAME}` or `${CommandAnalyze.COMMAND_NAME}` without `${AssembleContigsMixins.SetContigAssemblingFeatures.CMD_OPTION}` option."
            }
        }

        private fun processClns(): AssembleCellsReport =
            ClnsReader(inputFile, VDJCLibraryRegistry.getDefault()).use { reader ->
                validate(reader)
                val reportBuilder = AssembleCellsReport.Builder()

                val result = calculateGroupIdForClones(reader.readCloneSet(), reader.header, reportBuilder)

                logger.progress { "Writing output" }
                ClnsWriter(outputFile).use { writer ->
                    writer.writeCloneSet(result)
                    reportBuilder.setFinishMillis(System.currentTimeMillis())
                    val report = reportBuilder.buildReport()
                    writer.setFooter(reader.footer.addStepReport(AnalyzeCommandDescriptor.assembleCells, report))
                    report
                }
            }

        private fun processClna(): AssembleCellsReport =
            ClnAReader(inputFile, VDJCLibraryRegistry.getDefault(), Concurrency.noMoreThan(4)).use { reader ->
                validate(reader)
                val reportBuilder = AssembleCellsReport.Builder()

                val result = calculateGroupIdForClones(reader.readCloneSet(), reader.header, reportBuilder)

                val tempDest = TempFileManager.smartTempDestination(outputFile, "", !useLocalTemp.value)
                ClnAWriter(outputFile, tempDest).use { writer ->
                    var newNumberOfAlignments: Long = 0
                    val allAlignmentsList = mutableListOf<OutputPort<VDJCAlignments>>()
                    for (clone in result) {
                        newNumberOfAlignments += reader.numberOfAlignmentsInClone(clone.id)
                        allAlignmentsList += reader.readAlignmentsOfClone(clone.id)
                    }

                    logger.progress { "Writing output" }
                    writer.writeClones(result)
                    writer.collateAlignments(allAlignmentsList.flatten(), newNumberOfAlignments)
                    reportBuilder.setFinishMillis(System.currentTimeMillis())
                    val report = reportBuilder.buildReport()
                    writer.setFooter(reader.footer.addStepReport(AnalyzeCommandDescriptor.assembleCells, report))
                    writer.writeAlignmentsAndIndex()
                    report
                }
            }

        private fun calculateGroupIdForClones(
            input: CloneSet,
            header: MiXCRHeader,
            reportBuilder: AssembleCellsReport.Builder
        ): CloneSet {
            ValidationException.require(input.tagsInfo.hasTagsWithType(TagType.Cell)) {
                "Input doesn't have cell tags"
            }
            ValidationException.require(input.none { it.group != null }) {
                "Input file already grouped by cells"
            }

            val paramsSpec = resetPreset.overridePreset(header.paramsSpec)
            val (_, cmdParams) = paramsResolver.resolve(paramsSpec)


            reportBuilder.setStartMillis(System.currentTimeMillis())
            reportBuilder.setInputFiles(inputFiles)
            reportBuilder.setOutputFiles(outputFiles)
            reportBuilder.commandLine = commandLineArguments

            val grouper = cmdParams.algorithm.mkGrouper<Clone>(
                input.tagsInfo,
                input.cloneSetInfo.assemblingFeatures
            )
            SmartProgressReporter.startProgressReport(grouper)
            val groupedClones = grouper.groupClones(input.clones)
            reportBuilder.grouperReport = grouper.getReport()

            logger.progress { "Resorting and recalculating ranks" }
            return CloneSet.Builder(
                groupedClones,
                input.usedGenes,
                input.header
                    .copy(calculatedCloneGroups = true)
                    .addStepParams(AnalyzeCommandDescriptor.assembleCells, cmdParams)
                    .copy(paramsSpec = dontSavePresetOption.presetToSave(paramsSpec))
            )
                .sort(input.ordering)
                // some clones split, need to recalculate
                .recalculateRanks()
                // total sums are not changed
                .withTotalCounts(input.counts)
                .build()
        }
    }
}
