#     Copyright (C) 2023  Coretex LLC

#     This file is part of Coretex.ai

#     This program is free software: you can redistribute it and/or modify
#     it under the terms of the GNU Affero General Public License as
#     published by the Free Software Foundation, either version 3 of the
#     License, or (at your option) any later version.

#     This program is distributed in the hope that it will be useful,
#     but WITHOUT ANY WARRANTY; without even the implied warranty of
#     MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#     GNU Affero General Public License for more details.

#     You should have received a copy of the GNU Affero General Public License
#     along with this program.  If not, see <https://www.gnu.org/licenses/>.

from typing import List
from pathlib import Path
from zipfile import ZipFile

import logging
import shutil
import gzip

from ... import folder_manager
from ...entities import TaskRun, CustomSample, CustomDataset
from ...networking import NetworkRequestError


def createSample(name: str, dataset: CustomDataset, path: Path, taskRun: TaskRun, stepName: str, retryCount: int = 0) -> CustomSample:
    try:
        sample = dataset.add(path, name)
    except NetworkRequestError:
        if retryCount < 3:
            logging.info(f">> [Coretex] Retry count: {retryCount}")
            return createSample(name, dataset, path, taskRun, stepName, retryCount + 1)

        raise

    taskRun.createQiimeArtifact(f"{stepName}/{name}", path)

    return sample


def compressGzip(source: Path, destination: Path, deleteSource: bool = False) -> None:
    logging.info(f"{source} -> {destination}")

    with gzip.open(destination, "w") as destinationFile:
        destinationFile.write(source.read_bytes())

    if deleteSource:
        source.unlink()


def sampleNumber(sample: CustomSample) -> int:
    return int(sample.name.split("-")[0])


def isFastqMPSample(sample: CustomSample) -> bool:
    sample.unzip()

    sequenceFileNames = ["forward.fastq", "forward.fastq.gz", "sequences.fastq", "sequences.fastq.gz"]
    barcodesFileNames = ["barcodes.fastq", "barcodes.fastq.gz"]

    folderContent = list(sample.load().folderContent)

    sequenceFilePresent = any([path.name in sequenceFileNames for path in folderContent])
    barcodesFilePresent = any([path.name in barcodesFileNames for path in folderContent])

    return sequenceFilePresent and barcodesFilePresent


def isFastqDPSample(sample: CustomSample) -> bool:
    sample.unzip()

    return any([path.suffix == ".fastq" for path in sample.load().folderContent])


def isImportedSample(sample: CustomSample) -> bool:
    sample.unzip()

    sampleContent = [path.name for path in sample.load().folderContent]
    return "multiplexed-sequences.qza" in sampleContent


def isDemultiplexedSample(sample: CustomSample) -> bool:
    sample.unzip()

    sampleContent = [path.name for path in sample.load().folderContent]
    return "demux.qza" in sampleContent


def isDenoisedSample(sample: CustomSample) -> bool:
    sample.unzip()

    sampleContent = [path.name for path in sample.load().folderContent]
    return (
        "table.qza" in sampleContent and
        "rep-seqs.qza" in sampleContent and
        "stats.qza" in sampleContent
    )


def isPhylogeneticTreeSample(sample: CustomSample) -> bool:
    sample.unzip()

    sampleContent = [path.name for path in sample.load().folderContent]
    return (
        "rooted-tree.qza" in sampleContent and
        "unrooted-tree.qza" in sampleContent and
        "aligned-rep-seqs.qza" in sampleContent and
        "masked-aligned-rep-seqs.qza" in sampleContent
    )


def getFastqMPSamples(dataset: CustomDataset) -> List[CustomSample]:
    # Multiplexed fastq data
    return dataset.getSamples(isFastqMPSample)


def getFastqDPSamples(dataset: CustomDataset) -> List[CustomSample]:
    # Demultiplexed fastq data
    return dataset.getSamples(isFastqDPSample)


def getImportedSamples(dataset: CustomDataset) -> List[CustomSample]:
    return dataset.getSamples(isImportedSample)


def getDemuxSamples(dataset: CustomDataset) -> List[CustomSample]:
    return dataset.getSamples(isDemultiplexedSample)


def getDenoisedSamples(dataset: CustomDataset) -> List[CustomSample]:
    return dataset.getSamples(isDenoisedSample)


def getPhylogeneticTreeSamples(dataset: CustomDataset) -> List[CustomSample]:
    return dataset.getSamples(isPhylogeneticTreeSample)


def getMetadata(sample: CustomSample) -> Path:
    metadataPathList = list(sample.path.glob("*.tsv"))
    if len(metadataPathList) != 1:
        raise RuntimeError(f">> [Coretex] Metadata sample must contain one .tsv file. Found {len(metadataPathList)}")

    return metadataPathList[0]


def isPairedEnd(sample: CustomSample) -> bool:
    # In order to determine whether we are dealing with paired-end
    # sequences, this function unzips the qiime artifact and
    # reads the metadata, looking for the second (type) row, which will have
    # "PairedEnd" somewhere if it's paired-end

    sampleTemp = folder_manager.createTempFolder("qzaSample")
    qzaPath = list(sample.path.iterdir())[0]

    with ZipFile(qzaPath, "r") as qzaFile:
        qzaFile.extractall(sampleTemp)

    metadataPath = list(sampleTemp.rglob("*metadata.yaml"))[0]

    with metadataPath.open("r") as metadata:
        pairedEnd = "PairedEnd" in metadata.readlines()[1]

    shutil.rmtree(sampleTemp)

    return pairedEnd
