#!/usr/bin/env python
import argparse
import shutil
from importlib.resources import files
from pathlib import Path
from typing import Optional, Sequence

from metacarp.app.utils.snakemakeutils import parse_illumina_pairs
from metacarp.camel_gene_detection.app.command.command import Command
from metacarp.camel_gene_detection.app.snakemake.snakepipelineutils import SnakePipelineUtils
from metacarp.camel_gene_detection.app.loggers import initialize_logging, logger

class GmmScreening:
    """
    Main class to run the GMM screening.
    """

    def __init__(self, args: Optional[Sequence[str]] = None) -> None:
        """
        Initializes the main class.
        :param args: Arguments (optional)
        :return: None
        """
        self._args = GmmScreening._parse_arguments(args)
        self._path_tsv_out = self._args.output.absolute() / self._args.output_tsv
        if self._path_tsv_out.exists():
            self._path_tsv_out.unlink()

    def run(self) -> None:
        """
        Runs the main script.
        :return: None
        """
        logger.info('Starting GMM screening workflow')
        self._args.dir_working.mkdir(parents=True, exist_ok=True)
        self._check_dependencies()

        # Run Snakemake workflow
        path_snakefile = Path(str(files('metacarp').joinpath('resources/snakefiles/GMM_screening.smk')))
        path_config = self.__create_config_file()

        SnakePipelineUtils.run_snakemake(
            path_snakefile, path_config, [self._path_tsv_out], self._args.dir_working, self._args.threads)
        logger.info(f'Workflow finished successfully, output available in: {self._args.output}')

        # Copy the log file if it exists
        if (self._args.dir_working / 'GMM_screening.log').exists():
            shutil.copyfile(self._args.dir_working / 'GMM_screening.log', self._args.output / 'GMM_screening.log')

    @staticmethod
    def _parse_arguments(args: Optional[Sequence[str]]) -> argparse.Namespace:
        """
        Parses the command line arguments.
        :param args: Arguments
        :return: Parsed arguments.
        """
        parser = argparse.ArgumentParser()

        # Input and output
        parser.add_argument('--ilmn-in', help='Directory with Illumina input FASTQ files')
        parser.add_argument('--ont-in', help='Directory with ONT input FASTQ files')

        parser.add_argument(
            '--dir-working', type=Path, help='Working directory', default=Path.cwd())
        parser.add_argument('--output', required=True, type=Path, help='Output directory')
        parser.add_argument('--output-tsv', type=Path, help='Output report name', default='summary_all_samples.tsv')

        # Parameters
        parser.add_argument('--db', help='GMM DB - either "junctions" or "genes-vectors" (default: junctions)', type=Path,
                            default=str(files('metacarp').joinpath('resources/DB_GMM/V6/junctions')))
        parser.add_argument('--min-identity', type=int, default=90, help='Minimal %% identity with template to report kma hit')
        parser.add_argument('--min-coverage', type=int, default=90, help='Minimum %% coverage of template to report kma hit')

        # Threads
        parser.add_argument('--threads', type=int, default=4)

        return parser.parse_args(args)

    def __create_config_file(self) -> Path:
        """
        Creates the config file for the Snakemake workflow.
        :return: Path to config file
        """
        if str(self._args.db) == 'junctions':
            db_fasta = str(files('metacarp').joinpath('resources/DB_GMM/V6/junctions/GMM_junctions-clustered_80.fasta'))
        elif str(self._args.db) == 'genes-vectors':
            db_fasta = Path(str(files('metacarp').joinpath('resources/DB_GMM/V6/genes-vectors/genes-vectors-clustered_80.fasta')))
        else:
            raise FileNotFoundError(
                f'No valid database name provided, either choose "junctions" or "genes-vectors"')

        input_dict = {}
        if self._args.ilmn_in is not None:
            input_dict.update(parse_illumina_pairs(Path(self._args.ilmn_in)))
        if self._args.ont_in is not None:
            for fq_file in sorted(Path(self._args.ont_in).glob('*.fastq.gz')):
                input_dict[fq_file.name.split('.')[0]] = {'fq': str(fq_file.absolute()), 'tech': 'ont'}
        if len(input_dict) == 0:
            raise FileNotFoundError(
                f'No input datasets found (searched in Illumina={self._args.ilmn_in}, ONT={self._args.ont_in})')

        config_data = {
            'db': db_fasta,
            'threads': self._args.threads,
            'filters': {
                'min_identity': self._args.min_identity,
                'min_coverage': self._args.min_coverage,
            },
            'input': input_dict,
            'output': {
                'tsv': str(self._path_tsv_out),
                'dir': str(self._args.output.absolute()),
                'dir_working': str(self._args.dir_working.absolute()),
            },
        }
        return Path(SnakePipelineUtils.generate_config_file(config_data, self._args.dir_working.absolute()))

    def _check_dependencies(self) -> None:
        """
        Checks if the required dependencies are available.
        :return: None
        """
        commands = {
            'kma': 'kma -v',
            'fastp': 'fastp -v'
        }
        logger.info(f'Checking dependencies')
        for tool, command in commands.items():
            command = Command(command)
            command.run(self._args.dir_working, disable_logging=True)
            if not command.returncode == 0:
                raise RuntimeError(f"Dependency '{tool}' not available")
            logger.info(f"{tool}: OK")

def main() -> None:
    """
    Runs the main workflow.
    :return: None
    """
    initialize_logging()
    main_ = GmmScreening()
    main_.run()

if __name__ == '__main__':
    main()
