"""
Classes for interacting with the (ResFinder/PointFinder/PlasmidFinder) databases used to detect AMR genes.
"""
import argparse
import logging
import sys
from os import path, mkdir

from staramr.SubCommand import SubCommand
from staramr.Utils import get_string_with_spacing
from staramr.databases.AMRDatabasesManager import AMRDatabasesManager
from staramr.databases.resistance.ARGDrugTable import ARGDrugTable
from staramr.exceptions.CommandParseException import CommandParseException
from staramr.exceptions.DatabaseErrorException import DatabaseErrorException
from staramr.exceptions.DatabaseNotFoundException import DatabaseNotFoundException

"""
Base class for interacting with a database.
"""

logger = logging.getLogger('AMRDatabaseHandler')


class Database(SubCommand):

    def __init__(self, subparser, script_name):
        """
        Builds a SubCommand for interacting with databases.
        :param subparser: The subparser to use.  Generated from argparse.ArgumentParser.add_subparsers().
        :param script_name: The name of the script being run.
        """
        super().__init__(subparser, script_name)

    def _setup_args(self, arg_parser):
        arg_parser = self._subparser.add_parser('db', help='Download ResFinder/PointFinder/PlasmidFinder databases')
        subparser = arg_parser.add_subparsers(dest='db_command',
                                              help='Subcommand for ResFinder/PointFinder/PlasmidFinder databases.')

        Build(subparser, self._script_name + " db")
        Update(subparser, self._script_name + " db")
        Info(subparser, self._script_name + " db")
        RestoreDefault(subparser, self._script_name + " db")

        return arg_parser

    def run(self, args):
        super(Database, self).run(args)

        if args.db_command is None:
            self._root_arg_parser.print_help()


"""
Class for building a new database.
"""


class Build(Database):

    def __init__(self, subparser, script_name):
        """
        Creates a SubCommand for building a new database.
        :param subparser: The subparser to use.  Generated from argparse.ArgumentParser.add_subparsers().
        :param script_name: The name of the script being run.
        """
        super().__init__(subparser, script_name)

    def _setup_args(self, arg_parser):
        name = self._script_name
        self._default_dir = AMRDatabasesManager.get_default_database_directory()
        epilog = ("Example:\n"
                  "\t" + name + " build\n"
                                "\t\tBuilds a new ResFinder/PointFinder/PlasmidFinder database under " + self._default_dir + " if it does not exist\n\n" +
                  "\t" + name + " build --dir databases\n" +
                  "\t\tBuilds a new ResFinder/PointFinder/PlasmidFinder database under databases/")

        arg_parser = self._subparser.add_parser('build',
                                                epilog=epilog,
                                                formatter_class=argparse.RawTextHelpFormatter,
                                                help='Downloads and builds databases in the given directory.')
        arg_parser.add_argument('--dir', action='store', dest='destination', type=str,
                                help='The directory to download the databases into [' + self._default_dir + '].',
                                default=self._default_dir, required=False)
        arg_parser.add_argument('--resfinder-commit', action='store', dest='resfinder_commit', type=str,
                                help='The specific git commit for the resfinder database [latest].', required=False)
        arg_parser.add_argument('--pointfinder-commit', action='store', dest='pointfinder_commit', type=str,
                                help='The specific git commit for the pointfinder database [latest].', required=False)
        arg_parser.add_argument('--plasmidfinder-commit', action='store', dest='plasmidfinder_commit', type=str,
                                help='The specific git commit for the plasmidfinder database [latest].', required=False)
        return arg_parser

    def run(self, args):
        super(Build, self).run(args)

        if path.exists(args.destination):
            if args.destination == self._default_dir:
                raise CommandParseException("Error, default destination [" + args.destination + "] already exists",
                                            self._root_arg_parser, print_help=True)
            else:
                raise CommandParseException("Error, destination [" + args.destination + "] already exists",
                                            self._root_arg_parser)
        else:
            mkdir(args.destination)

        if args.destination == AMRDatabasesManager.get_default_database_directory():
            database_repos = AMRDatabasesManager.create_default_manager().get_database_repos()
        else:
            database_repos = AMRDatabasesManager(args.destination).get_database_repos()
        database_repos.build({'resfinder': args.resfinder_commit, 'pointfinder': args.pointfinder_commit,
                              'plasmidfinder': args.plasmidfinder_commit})
        if not AMRDatabasesManager.is_database_repos_default_commits(database_repos):
            logger.warning(
                "Built non-default ResFinder/PointFinder/PlasmidFinder database version. This may lead to " +
                "differences in the detected AMR genes depending on how the database files are structured.")


"""
Class for updating an existing database.
"""


class Update(Database):

    def __init__(self, subparser, script_name):
        """
        Creates a SubCommand for updating an existing database.
        :param subparser: The subparser to use.  Generated from argparse.ArgumentParser.add_subparsers().
        :param script_name: The name of the script being run.
        """
        super().__init__(subparser, script_name)

    def _setup_args(self, arg_parser):
        self._default_dir = AMRDatabasesManager.get_default_database_directory()
        name = self._script_name
        epilog = ("Example:\n"
                  "\t" + name + " update databases/\n"
                                "\t\tUpdates the ResFinder/PointFinder/PlasmidFinder database under databases/\n\n" +
                  "\t" + name + " update -d\n" +
                  "\t\tUpdates the default ResFinder/PointFinder/PlasmidFinder database under " + self._default_dir)
        arg_parser = self._subparser.add_parser('update',
                                                epilog=epilog,
                                                formatter_class=argparse.RawTextHelpFormatter,
                                                help='Updates databases in the given directories.')

        arg_parser.add_argument('-d', '--update-default', action='store_true', dest='update_default',
                                help='Updates default database directory (' + self._default_dir + ').', required=False)
        arg_parser.add_argument('--resfinder-commit', action='store', dest='resfinder_commit', type=str,
                                help='The specific git commit for the resfinder database [latest].', required=False)
        arg_parser.add_argument('--pointfinder-commit', action='store', dest='pointfinder_commit', type=str,
                                help='The specific git commit for the pointfinder database [latest].', required=False)
        arg_parser.add_argument('--plasmidfinder-commit', action='store', dest='plasmidfinder_commit', type=str,
                                help='The specific git commit for the plasmidfinder database [latest].', required=False)
        arg_parser.add_argument('directories', nargs='*')

        return arg_parser

    def run(self, args):
        super(Update, self).run(args)

        if len(args.directories) == 0:
            if not args.update_default:
                raise CommandParseException("Must pass at least one directory to update, or use '--update-default'",
                                            self._root_arg_parser,
                                            print_help=True)
            else:
                try:
                    database_repos = AMRDatabasesManager.create_default_manager().get_database_repos(
                        force_use_git=True)

                    database_repos.update(
                        {'resfinder': args.resfinder_commit, 'pointfinder': args.pointfinder_commit,
                         'plasmidfinder': args.plasmidfinder_commit})

                    if not AMRDatabasesManager.is_database_repos_default_commits(database_repos):
                        logger.warning(
                            "Updated to non-default ResFinder/PointFinder/PlasmidFinder database version. This may lead to " +
                            "differences in the detected AMR genes depending on how the database files are structured.")
                except DatabaseErrorException as e:
                    logger.error(
                        "Could not update default database. Please try restoring with 'staramr db restore-default'")
                    raise e
        else:
            for directory in args.directories:
                database_repos = AMRDatabasesManager(directory).get_database_repos()
                database_repos.update({'resfinder': args.resfinder_commit, 'pointfinder': args.pointfinder_commit,
                                       'plasmidfinder': args.plasmidfinder_commit})
                if not AMRDatabasesManager.is_database_repos_default_commits(database_repos):
                    logger.warning(
                        "Updated to non-default ResFinder/PointFinder/PlasmidFinder database version [%s]. This may lead to " +
                        "differences in the detected AMR genes depending on how the database files are structured.",
                        directory)


"""
Class for restoring the default database.
"""


class RestoreDefault(Database):

    def __init__(self, subparser, script_name):
        """
        Creates a SubCommand for restoring to the default database.
        :param subparser: The subparser to use.  Generated from argparse.ArgumentParser.add_subparsers().
        :param script_name: The name of the script being run.
        """
        super().__init__(subparser, script_name)

    def _setup_args(self, arg_parser):
        name = self._script_name
        epilog = ("Example:\n"
                  "\t" + name + " restore-default\n"
                                "\t\tRestores the default ResFinder/PointFinder/PlasmidFinder database\n\n")
        arg_parser = self._subparser.add_parser('restore-default',
                                                epilog=epilog,
                                                formatter_class=argparse.RawTextHelpFormatter,
                                                help='Restores the default ResFinder/PointFinder/PlasmidFinder databases.')

        arg_parser.add_argument('-f', '--force', action='store_true', dest='force',
                                help='Force restore without asking for confirmation.', required=False)
        return arg_parser

    def _confirm_restore(self):
        """
        Confirms with the user whether or not to restore the database directory.
        :return: True if should restore, False otherwise.
        """
        confirmed = False
        while not confirmed:
            response = str(input(
                "Restore the default ResFinder/PointFinder/PlasmidFinder databases (Y/N)? ").lower().strip())
            if response == 'y' or response == 'yes':
                return True
            elif response == 'n' or response == 'no':
                return False

    def run(self, args):
        super(RestoreDefault, self).run(args)

        database_manager = AMRDatabasesManager.create_default_manager()

        if not args.force:
            response = self._confirm_restore()
        else:
            response = True

        if response:
            database_manager.restore_default()


"""
Class for getting information from an existing database.
"""


class Info(Database):

    def __init__(self, subparser, script_name):
        """
        Creates a SubCommand for printing information about a database.
        :param subparser: The subparser to use.  Generated from argparse.ArgumentParser.add_subparsers().
        :param script_dir: The directory containing the main application script.
        :param script_name: The name of the script being run.
        """
        super().__init__(subparser, script_name)

    def _setup_args(self, arg_parser):
        name = self._script_name
        default_dir = AMRDatabasesManager.get_default_database_directory()
        epilog = ("Example:\n"
                  "\t" + name + " info\n"
                                "\t\tPrints information about the default database in " + default_dir + "\n\n" +
                  "\t" + name + " info databases\n" +
                  "\t\tPrints information on the database stored in databases/")
        arg_parser = self._subparser.add_parser('info',
                                                epilog=epilog,
                                                formatter_class=argparse.RawTextHelpFormatter,
                                                help='Prints information on databases in the given directories.')
        arg_parser.add_argument('directories', nargs='*')

        return arg_parser

    def run(self, args):
        super(Info, self).run(args)

        arg_drug_table = ARGDrugTable()

        if len(args.directories) == 0:
            database_repos = AMRDatabasesManager.create_default_manager().get_database_repos()
            if not AMRDatabasesManager.is_database_repos_default_commits(database_repos):
                logger.warning(
                    "Using non-default ResFinder/PointFinder/PlasmidFinder database versions. This may lead to differences in the detected " +
                    "AMR genes depending on how the database files are structured.")

            try:
                database_info = database_repos.info()
                database_info.update(arg_drug_table.get_resistance_table_info())
                sys.stdout.write(get_string_with_spacing(database_info))
            except DatabaseNotFoundException as e:
                logger.error("No database found. Perhaps try restoring the default with 'staramr db restore-default'")
        else:
            for directory in args.directories:
                try:
                    database_repos = AMRDatabasesManager(directory).get_database_repos()
                    if not AMRDatabasesManager.is_database_repos_default_commits(database_repos):
                        logger.warning(
                            "Using non-default ResFinder/PointFinder/PlasmidFinder database version [%s]. This may lead to " +
                            "differences in the detected AMR genes depending on how the database files are structured.",
                            directory)

                    database_info = database_repos.info()
                    database_info.update(arg_drug_table.get_resistance_table_info())
                    sys.stdout.write(get_string_with_spacing(database_info))
                except DatabaseNotFoundException as e:
                    logger.error("Database not found in [%s]. Perhaps try building with 'staramr db build --dir %s'",
                                 directory, directory)
