import datetime
import gzip
import os
import logging
import re
import shutil

from abc import ABC, abstractmethod
from collections import defaultdict

from utils.input_output import is_path_writable_or_can_be_created
from utils.helper_funcs import recursively_iter_dict
from utils.databases.superclasses.database_superclasses import DbYamlInteractor

logger = logging.getLogger(__name__)
logger.addHandler(logging.NullHandler())


class Migrator(ABC):
    def __init__(self, lst_yaml_objs):
        self.lst_yaml_objs = lst_yaml_objs
        self.assert_no_db_is_below_from_version()
        self.yaml_objs_to_migrate = self.get_yamls_to_migrate()
        self.migrated_yaml_objs = self.get_yamls_already_satisfying_target_version()
        self.files = self.get_files_to_change()

    def assert_no_db_is_below_from_version(self):
        for y in self.lst_yaml_objs:
            if y.get_db_schema_version() < self.from_version:
                raise ValueError("Provided Database has schema version lower than this migrator can deal with: \n"
                                 + f"\trequired version: {self.from_version}\n"
                                 + f"\tdatabase version {y.get_db_schema_version()}\n"
                                 + f"\tdatabase: {y.get_yaml()}\n")

    def get_yamls_to_migrate(self):
        yamls_to_migrate = []
        for y in self.lst_yaml_objs:
            if y.get_db_schema_version() < self.to_version:
                yamls_to_migrate.append(y)
        return yamls_to_migrate

    def get_yamls_already_satisfying_target_version(self):
        yamls = []
        for y in self.lst_yaml_objs:
            if y.get_db_schema_version() >= self.to_version:
                yamls.append(y)
        return yamls

    def get_files_to_change(self):
        set_of_files = set()
        for method in self.list_of_ordered_migration_methods:
            set_of_files.update(method(dry_run=True))
        return set_of_files

    def migrate(self):
        logger.warning(f"Executing database migration from schema version {self.from_version} to {self.to_version}.")
        protected_files = self.get_protected_files()
        if protected_files:
            raise PermissionError(
                "Write access required to the following files for database style migration.\n"
                + "Please make the following files writeable and try again:\n"
                + '\n'.join(protected_files)
            )
        for method in self.list_of_ordered_migration_methods:
            method()
        self.bump_version_of_yaml_objs_to_migrate()
        self.write_yamls_of_yaml_objs_to_migrate()
        self.migrated_yaml_objs += self.yaml_objs_to_migrate
        return self.migrated_yaml_objs

    @property
    @abstractmethod
    def from_version(self):
        return None

    @property
    @abstractmethod
    def to_version(self):
        return None

    @property
    @abstractmethod
    def list_of_ordered_migration_methods(self):
        return []

    def get_protected_files(self):
        lst_protected_files = []
        for f in set(self.files):
            if not is_path_writable_or_can_be_created(f):
                lst_protected_files.append(f)
        return lst_protected_files

    def bump_version_of_yaml_objs_to_migrate(self):
        for y in self.yaml_objs_to_migrate:
            y._config["db_schema_version"] = self.to_version

    def write_yamls_of_yaml_objs_to_migrate(self):
        for y in self.yaml_objs_to_migrate:
            y.write_yaml(force=True)


def test_fail_of_migrator_without_method_overriding():
    db = ""
    import pytest
    with pytest.raises(TypeError):
        Migrator(db)


class MigrateV00toV01(Migrator):
    """
    Remove non-gzipped map files (*.map) from file system
    adjust paths in affected yamls (*.map -> *.map.gz)
    add "db_schema_version: 1" to yamls
    """
    from_version = 0
    to_version = 1

    @property
    def list_of_ordered_migration_methods(self):
        return [
            self.remove_old_map_files,
            self.update_db_config_map_entries
        ]

    def remove_old_map_files(self, dry_run=False):
        manipulated_files = []
        for yaml_object in self.yaml_objs_to_migrate:
            db_config = yaml_object.get_config_deepcopy()
            map_files = self.get_all_db_associated_map_files(db_config)
            for map_file in map_files:
                if self._is_plain_map_entry(map_file):
                    if os.path.exists(map_file):
                        if not dry_run:
                            os.remove(map_file)
                        manipulated_files.append(map_file)
        return manipulated_files

    def get_all_db_associated_map_files(self, db_config):
        map_files = []
        for k, v in recursively_iter_dict(db_config):
            if self._is_map_entry(k):
                map_files.append(v)
        return map_files

    def _is_map_entry(self, config_key):
        if config_key.endswith('map'):
            return True
        else:
            return False

    def _is_plain_map_entry(self, map_file):
        if type(map_file) is str:
            if map_file.endswith('.map'):
                return True
        else:
            raise ValueError('Got element of type "{}", can only process str'.format(type(map_file)))
        return False

    def update_db_config_map_entries(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            map_entry_keys = [k for k in yaml_obj_config if self._is_map_entry(k)]
            if map_entry_keys:
                adjusted_files.append(yaml_obj.get_yaml())
                if not dry_run:
                    for map_key in map_entry_keys:
                        map_path = yaml_obj_config[map_key]
                        yaml_obj._config[map_key] = self.replace_map_by_map_gz(map_path)
        return adjusted_files

    @staticmethod
    def replace_map_by_map_gz(map_path):
        regex = re.compile(r".map$")
        new_map_path = regex.sub(".map.gz", map_path)
        return new_map_path


def test_replace_map_by_map_gz():
    assert MigrateV00toV01.replace_map_by_map_gz("wasd.map") == "wasd.map.gz"


class MigrateV01toV02(Migrator):
    """
    Remove skip.yaml
    bump "db_schema_version: 2" in yaml
    """
    from_version = 1
    to_version = 2

    @property
    def list_of_ordered_migration_methods(self):
        return [
            self.remove_skip_yaml
        ]

    def remove_skip_yaml(self, dry_run=False):
        adjusted_files = []
        lst_skip_yaml_objs = [y for y in self.yaml_objs_to_migrate if y.get_yaml().endswith("skip.yaml")]
        if len(lst_skip_yaml_objs) == 1:
            skip_yaml_obj = lst_skip_yaml_objs[0]
            skip_yaml_file = skip_yaml_obj.get_yaml()
            adjusted_files.append(skip_yaml_file)
            if not dry_run:
                os.remove(skip_yaml_file)
                self.yaml_objs_to_migrate = [y for y in self.yaml_objs_to_migrate
                                             if not y.get_yaml().endswith("skip.yaml")]
        elif len(lst_skip_yaml_objs) > 1:
            raise ValueError(
                "Expected one skip.yaml db config but got multiple:\n"
                + "\n".join(lst_skip_yaml_objs)
            )
        return adjusted_files


class MigrateV02toV03(Migrator):
    """
    Add accession2taxid yaml
    Add uniprot_tax yaml
    Change taxdump scope to 'helper_db'
    Add required_dbs section to nr and uniprot yamls
    Rename database folders to match db_name and db_version property
    add "fasta_gz" relative path to ncbi + uniprot yamls
    add "hmmlib" relative path to tigrfams + pfams yamls
    change absolute paths in eggnog to relative paths and move the to 'files' section in yaml
    bump "db_schema_version: 3" in yaml
    """
    from_version = 2
    to_version = 3

    @property
    def list_of_ordered_migration_methods(self):
        return [
            self.add_accession2taxid_yaml,
            self.add_uniprot_tax_yaml,
            self.change_taxdump_scope_to_helper_db,
            self.add_taxdump_tar_to_taxdump_yaml_files,
            self.add_required_dbs_to_nr_and_uniprot,
            self.add_files_section_to_eggnog,
            self.remove_non_relative_paths,
            self.rename_db_folder_to_yaml_attributes_and_adjust_yaml_path,
            self.add_unprocessed_sequence_lib_to_yaml
        ]

    def add_accession2taxid_yaml(self, dry_run=False):
        adjusted_files = []
        uniprot_ncbi_yaml_objs = self.get_yamls_objs_matching_db_types(
            ["ncbi_nr", "uniprot_complete", "uniprot_sp", "uniprot_tr"])
        set_acc2tax_files = set()
        for y in uniprot_ncbi_yaml_objs:
            y_cfg = y.get_config_deepcopy()
            lst_acc2tax_files = y_cfg["acc2tax"]
            set_acc2tax_files.update(lst_acc2tax_files)

        db_dir_dict = defaultdict(list)
        for f in set_acc2tax_files:
            db_dir_dict[os.path.dirname(f)].append(f)
        for db_dir, files in db_dir_dict.items():
            yaml_path = os.path.join(db_dir, "accession2taxid.yaml")
            adjusted_files.append(yaml_path)

            # build yaml dict:
            db_config_dict = {}
            db_config_dict['name'] = "NCBI accession to taxid mapping database"
            db_config_dict['type'] = "accession2taxid"
            version_dir = os.path.basename(db_dir)
            version = self.convert_version_dir_to_datetime(version_dir)
            db_config_dict["version"] = version
            db_config_dict["comment"] = f"test database of the original database downloaded at {version}"
            db_config_dict["files"] = {os.path.basename(f).split(".")[0] + "2taxid": os.path.basename(f)
                                       for f in set_acc2tax_files}
            db_config_dict["scope"] = "helper_db"
            db_config_dict["db_schema_version"] = 2
            if not dry_run:
                yaml_obj = DbYamlInteractor(yaml=yaml_path, db_yaml_dict=db_config_dict)
                self.yaml_objs_to_migrate.append(yaml_obj)
                self.lst_yaml_objs.append(yaml_obj)
                for f in db_config_dict["files"].values():
                    path = os.path.join(db_dir, f)
                    assert os.path.exists(path), f"Required file does not exist: {path}"
        return adjusted_files

    def get_yamls_objs_matching_db_types(self, lst_of_db_types):
        matching_yaml_objs = []
        for y in self.lst_yaml_objs:
            y_cfg = y.get_config_deepcopy()
            db_type = y_cfg['type']
            if db_type in lst_of_db_types:
                matching_yaml_objs.append(y)
        return matching_yaml_objs

    def convert_version_dir_to_datetime(self, version_dir):
        return datetime.date(int(version_dir[:4]), int(version_dir[4:6]), int(version_dir[6:]))

    def add_uniprot_tax_yaml(self, dry_run=False):
        adjusted_files = []
        uniprot_yaml_objs = self.get_yamls_objs_matching_db_types(
            ["uniprot_complete", "uniprot_sp", "uniprot_tr"])
        set_uniprot_tax_files = set()
        full_path_file_dict = {}
        for y in uniprot_yaml_objs:
            y_cfg = y.get_config_deepcopy()
            dir_path = os.path.dirname(y_cfg["taxmap"])
            idmapping_path = os.path.join(dir_path, 'idmapping.dat.gz')
            full_path_file_dict['idmap'] = idmapping_path
            uniprot_tax_txt_path = os.path.join(dir_path, "uniprot_tax.txt")
            full_path_file_dict['uniprot_tax_txt'] = uniprot_tax_txt_path
            set_uniprot_tax_files.update([idmapping_path, uniprot_tax_txt_path])

        assert len(set_uniprot_tax_files) == 2, \
            ((f"Expected two uniprot_tax db_files, got: {len(set_uniprot_tax_files)}\n\t"
             + "\n\t".join(set_uniprot_tax_files)))

        # build yaml_path
        db_dir = os.path.commonpath([os.path.dirname(f) for f in set_uniprot_tax_files])
        assert all(db_dir == os.path.dirname(f) for f in set_uniprot_tax_files), \
            ("All uniprot_tax db_files must be located in the same subdirectory\n\t"
             + "\n\t".join(set_uniprot_tax_files))
        yaml_path = os.path.join(db_dir, "uniprot_tax.yaml")
        adjusted_files.append(yaml_path)

        # build yaml dict:
        db_config_dict = {}
        db_config_dict['name'] = "Uniprot accession to taxid and taxid to tax-name mapping database"
        db_config_dict['type'] = "uniprot_tax"
        version_dir = os.path.basename(db_dir)
        version = self.convert_version_dir_to_datetime(version_dir)
        db_config_dict["version"] = version
        db_config_dict["comment"] = f"test database of the original database downloaded at {version}"
        db_config_dict["files"] = {key: os.path.basename(path) for key, path in full_path_file_dict.items()}
        db_config_dict["scope"] = "helper_db"
        db_config_dict["db_schema_version"] = 2
        if not dry_run:
            yaml_obj = DbYamlInteractor(yaml=yaml_path, db_yaml_dict=db_config_dict)
            self.yaml_objs_to_migrate.append(yaml_obj)
            self.lst_yaml_objs.append(yaml_obj)
            for f in db_config_dict["files"].values():
                path = os.path.join(db_dir, f)
                assert os.path.exists(path), f"Required file does not exist: {path}"
        return adjusted_files

    def change_taxdump_scope_to_helper_db(self, dry_run=False):
        adjusted_files = []
        lst_yaml_objs = self.get_yamls_objs_matching_db_types(['taxdump'])
        for yaml_obj in lst_yaml_objs:
            if not dry_run:
                yaml_obj._config["scope"] = "helper_db"
            adjusted_files.append(yaml_obj.get_yaml())
        return adjusted_files

    def add_taxdump_tar_to_taxdump_yaml_files(self, dry_run=False):
        adjusted_files = []
        lst_taxdump_yaml_obj = self.get_yamls_objs_matching_db_types(['taxdump'])
        for taxdump_yaml_obj in lst_taxdump_yaml_obj:
            taxdump_cfg = taxdump_yaml_obj.get_config_deepcopy()
            taxdump_version = taxdump_cfg["version"]
            set_taxdump_file_paths = set()
            for db_type in ["ncbi_nr", "uniprot_complete", "uniprot_sp", "uniprot_tr"]:
                try:
                    y_obj = self.get_one_yaml_obj_matching_db_type_and_version(db_type, taxdump_version)
                except ValueError as e:
                    logger.warning(f"No matching db of type {db_type} found for taxdump yaml {taxdump_yaml_obj.get_yaml()}")
                    continue
                set_taxdump_file_paths.add(y_obj._config["taxdump"])
            assert len(set_taxdump_file_paths) == 1, \
                (f"Expected exactly one taxdump file, got {len(set_taxdump_file_paths)}\n"
                 + "\n".join(set_taxdump_file_paths))
            db_dir = os.path.dirname(taxdump_yaml_obj.get_yaml())
            taxdump_path = set_taxdump_file_paths.pop()
            assert os.path.samefile(os.path.dirname(taxdump_path), db_dir), \
                ("taxdump.tar.gz and taxdump.yaml must lie in same folder.\n"
                 + f"taxdump.tar.gz path: {taxdump_path}\n"
                 + f"taxdump.yaml path: {taxdump_yaml_obj.get_yaml()}")
            files_dict = taxdump_cfg.get("files", {})
            files_dict["taxdump"] = os.path.basename(taxdump_path)
            taxdump_cfg["files"] = files_dict
            if not dry_run:
                for f in taxdump_cfg["files"].values():
                    path = os.path.join(db_dir, f)
                    assert os.path.exists(path), f"Required file does not exist: {path}"
                taxdump_yaml_obj._config = taxdump_cfg
            adjusted_files.append(taxdump_yaml_obj.get_yaml())
        return adjusted_files

    def add_required_dbs_to_nr_and_uniprot(self, dry_run=False):
        adjusted_files = []
        lst_file_keys = ["acc2tax", "taxdump_map", "taxdump", "taxmap"]
        for db_type in ["ncbi_nr", "uniprot_complete", "uniprot_sp", "uniprot_tr"]:
            for y_obj in self.get_yamls_objs_matching_db_types([db_type]):
                if not dry_run:
                    y_cfg = y_obj.get_config_deepcopy()
                    lst_required_files = []
                    for key, value in y_cfg.items():
                        if key in lst_file_keys:
                            if key == "acc2tax":
                                lst_required_files += value
                            else:
                                lst_required_files.append(value)
                    set_required_folders = {os.path.dirname(f) for f in lst_required_files}
                    lst_helper_yaml_objs = []
                    for folder in set_required_folders:
                        lst_helper_yaml_objs += [
                            y for y in self.lst_yaml_objs
                            if os.path.samefile(folder, os.path.dirname(y.get_yaml()))
                        ]
                    if len(lst_helper_yaml_objs) != len(set_required_folders):
                        raise ValueError(f"Expected {len(set_required_folders)} helper DBs, got {len(lst_helper_yaml_objs)}\n"
                                         + f"Folders:\n\t" + "\n\t".join(set_required_folders) + "\n"
                                         + f"Required DB yamls:\n\t"
                                         + "\n\t".join([y.get_yaml() for y in lst_helper_yaml_objs]) + "\n")
                    required_dbs = {}
                    for helper_db_yaml_obj in lst_helper_yaml_objs:
                        helper_db_cfg = helper_db_yaml_obj.get_config_deepcopy()
                        helper_db_type = helper_db_cfg["type"]
                        required_dbs[helper_db_type] = {}
                        required_dbs[helper_db_type]["version"] = helper_db_cfg["version"]
                        y_obj._config["required_dbs"] = required_dbs
                adjusted_files.append(y_obj.get_yaml())
        return adjusted_files

    def add_files_section_to_eggnog(self, dry_run=False):
        y_obj = self.get_one_yaml_obj_matching_db_type("eggnog")
        y_path = y_obj.get_yaml()
        adjusted_files = [y_path]
        y_cfg = y_obj.get_config_deepcopy()
        eggnog_dir = os.path.dirname(y_path)
        if "files" in y_cfg:
            files_dict = y_cfg["files"]
        else:
            files_dict = {}
        for key in ["data_dir", "og_annotations", "fun_cats"]:
            if key == "data_dir":
                new_key = "emapper_data"
            else:
                new_key = key
            full_file_path = y_cfg[key]
            rel_file_path = os.path.relpath(full_file_path, eggnog_dir)
            files_dict[new_key] = rel_file_path
        y_cfg["files"] = files_dict
        if not dry_run:
            y_obj._config = y_cfg
        return adjusted_files

    def remove_non_relative_paths(self, dry_run=False):
        adjusted_files = []
        for db_type in ["eggnog", "ncbi_nr", "uniprot_complete", "uniprot_sp", "uniprot_tr"]:
            if db_type == "eggnog":
                keys_to_remove = ["data_dir", "og_annotations", "fun_cats"]
            else:
                keys_to_remove = ["acc2tax", "taxdump_map", "taxdump", "taxmap"]
            for y_obj in self.get_yamls_objs_matching_db_types([db_type]):
                y_cfg = y_obj.get_config_deepcopy()
                keys_in_y = [k for k in keys_to_remove if k in y_cfg]
                if keys_in_y:
                    adjusted_files.append(y_obj.get_yaml())
                    for k in keys_in_y:
                        del y_cfg[k]
                if not dry_run:
                    y_obj._config = y_cfg
        return adjusted_files

    def get_one_yaml_obj_matching_db_type(self, db_type):
        lst_y_objs = self.get_yamls_objs_matching_db_types([db_type])
        if len(lst_y_objs) != 1:
            raise ValueError(f"Expected exactly one yaml object with type '{db_type}'\n"
                             + f", got: {len(lst_y_objs)}\n"
                             + "\n".join(map(str, lst_y_objs)))
        return lst_y_objs[0]

    def get_one_yaml_obj_matching_db_type_and_version(self, db_type, db_version):
        lst_y_objs = self.get_yamls_objs_matching_db_types([db_type])
        matching_objs = [y for y in lst_y_objs if y._config["version"] == db_version]
        if len(matching_objs) > 1:
            raise ValueError(f"Expected exactly one yaml object with type '{db_type}'\n"
                             + f", got: {len(matching_objs)}\n"
                             + "\n".join(map(str, matching_objs)))
        elif len(matching_objs) < 1:
            raise ValueError(f"Expected exactly one yaml object with type '{db_type}' and version '{db_version}'\n"
                             + f", got: {len(matching_objs)}\n"
                             + "object with correct type but wrong version:\n"
                             + "\n".join(map(str, lst_y_objs)))
        return matching_objs[0]

    def rename_db_folder_to_yaml_attributes_and_adjust_yaml_path(self, dry_run=False):
        adjusted_files = []
        yaml_objs_to_migrate_wo_skip = [y for y in self.yaml_objs_to_migrate if not y.get_yaml().endswith("skip.yaml")]
        # group by db_type dir
        db_type_dct = defaultdict(list)
        for y in yaml_objs_to_migrate_wo_skip:
            db_type_dct[y.get_config_deepcopy()['type']].append(y)
        for db_type_str, lst_yaml_objs in db_type_dct.items():
            db_type_subdir_renamed = False
            for yaml_object in lst_yaml_objs:
                yaml_obj_config = yaml_object.get_config_deepcopy()
                db_version_str = str(yaml_obj_config['version'])
                yaml_path = yaml_object.get_yaml()
                db_version_subdir = self.get_version_from_yaml_path(yaml_path)
                if db_type_subdir_renamed:
                    db_type_subdir = db_type_str
                else:
                    db_type_subdir = self.get_type_from_yaml_path(yaml_path)
                db_base_dir = self.get_db_base_dir_from_yaml_path(yaml_path)
                lst_old_names_new_names = [
                    (os.path.join(db_base_dir, db_type_subdir, db_version_subdir), os.path.join(db_base_dir, db_type_subdir, db_version_str)),
                    (os.path.join(db_base_dir, db_type_subdir), os.path.join(db_base_dir, db_type_str))
                ]
                for old_name, new_name in lst_old_names_new_names:
                    if not old_name == new_name:
                        if not dry_run:
                            os.rename(old_name, new_name)
                        adjusted_files += [old_name, new_name]
                new_yaml_name = os.path.join(db_base_dir, db_type_str, db_version_str,
                                             os.path.basename(yaml_object.get_yaml()))
                db_type_subdir_renamed = True
                adjusted_files.append(new_yaml_name)
                if not dry_run:
                    yaml_object.yaml = new_yaml_name
        return adjusted_files

    def add_unprocessed_sequence_lib_to_yaml(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            db_type = yaml_obj_config['type']
            db_is_uniprot_or_ncbi_based = (db_type == "ncbi_nr"
                                           or db_type in ["uniprot_complete", "uniprot_sp", "uniprot_tr"])
            yaml_path_basename = os.path.basename(yaml_obj.get_yaml())
            yaml_path_basename_wo_ext = os.path.splitext(yaml_path_basename)[0]
            if "files" in yaml_obj_config:
                files_dict = yaml_obj_config["files"]
            else:
                files_dict = {}
            if db_type in ['tigrfams', 'pfams']:
                adjusted_files.append(yaml_obj.get_yaml())

                if db_type == 'tigrfams':
                    prog_db_string = yaml_path_basename_wo_ext + ".LIB"
                elif db_type == 'pfams':
                    prog_db_string = yaml_path_basename_wo_ext + ".hmm"
                files_dict['hmm_lib'] = prog_db_string
            elif db_is_uniprot_or_ncbi_based:
                adjusted_files.append(yaml_obj.get_yaml())
                prog_db_string = yaml_path_basename_wo_ext + ".gz"
                files_dict['fasta_gz'] = prog_db_string
            # eggnog data dir is already added in add_files_section_to_eggnog
            yaml_obj_config["files"] = files_dict
            for f in yaml_obj_config["files"].values():
                path = os.path.join(os.path.dirname(yaml_obj.get_yaml()), f)
                assert os.path.exists(path), f"Required file does not exist: {path}"
            if not dry_run:
                yaml_obj._config = yaml_obj_config
                for f in yaml_obj.get_config_deepcopy()["files"].values():
                    path = os.path.join(os.path.dirname(yaml_obj.get_yaml()), f)
                    assert os.path.exists(path), f"Required file does not exist: {path}"
        return adjusted_files

    def get_version_from_yaml_path(self, yaml_path):
        return os.path.basename(os.path.dirname(yaml_path))

    def get_type_from_yaml_path(self, yaml_path):
        return os.path.basename(os.path.dirname(os.path.dirname(yaml_path)))

    def get_db_base_dir_from_yaml_path(self, yaml_path):
        return os.path.dirname(os.path.dirname(os.path.dirname(yaml_path)))


class MigrateV03toV04(Migrator):
    """
    replace hmm libs file entries with the gzipped (*.gz) files (except for dbCAN, since downloaded HMMs are plain txt).
        If the gzipped files are not present, create them.
    Remove plain versions of hmm libs
    add files required for pfam map creation to respective files section
    add files required for tigrfam map creation to respective files section
    add uniprot complete requirements: uniprot_sp and uniprot_tr
    remove resfams_{core,full} tsv map-resource file
    add resfams_core files Resfams-core.xlsx to yaml files section
    add resfams_full files Resfams-full.xlsx to yaml files section
    add foam file FOAM-hmm_rel1a.tsv to yaml files section
    add dbcan file dbCAN-HMMdb-V8.fam-activities.txt to yaml files section
    remove dependency on taxdump and accession2taxid DBs from Uniprot DBs
    bump "db_schema_version: 4" in yaml
    """
    from_version = 3
    to_version = 4

    @property
    def list_of_ordered_migration_methods(self):
        return [
            self.replace_hmm_libs_with_gz_and_rename_processed_files,
            self.add_pfam_files_for_mapping,
            self.add_tigrfam_files_for_mapping,
            self.add_sp_and_tr_to_uniprot_complete_requirements,
            self.add_dbcan_files_for_mapping,
            self.add_foam_files_for_mapping,
            self.remove_resfam_tsv,
            self.add_resfam_files_for_mapping,
            self.remove_uniprot_complete_files,
            self.remove_redundant_uniprot_dependancies
        ]

    def get_yamls_objs_matching_db_types(self, lst_of_db_types):
        matching_yaml_objs = []
        for y in self.lst_yaml_objs:
            y_cfg = y.get_config_deepcopy()
            db_type = y_cfg['type']
            if db_type in lst_of_db_types:
                matching_yaml_objs.append(y)
        return matching_yaml_objs

    def get_one_yaml_obj_matching_db_type(self, db_type):
        lst_y_objs = self.get_yamls_objs_matching_db_types(db_type)
        if len(lst_y_objs) != 1:
            raise ValueError(f"Expected exactly one yaml object with type '{db_type}'\n"
                             + f", got: {len(lst_y_objs)}\n"
                             + "\n".join(lst_y_objs))
        return lst_y_objs[0]

    def replace_hmm_libs_with_gz_and_rename_processed_files(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type in ['tigrfams', 'pfams', 'foam', 'dbcan', 'resfams_full', 'resfams_core']:
                adjusted_files.append(yaml_obj.get_yaml())
                plain_lib_full_path = os.path.join(os.path.dirname(yaml_obj.get_yaml()), files_dict['hmm_lib'])
                assert os.path.exists(plain_lib_full_path), f"Required file does not exist: {plain_lib_full_path}"
                adjusted_files.extend([plain_lib_full_path])
                if db_type == 'dbcan':
                    files_dict['hmm_lib'] = files_dict['hmm_lib'] + ".unzipped"
                else:
                    gz_lib_full_path = plain_lib_full_path + ".gz"
                    adjusted_files.extend([gz_lib_full_path])
                    files_dict['hmm_lib'] = files_dict['hmm_lib'] + ".gz"
                files_to_rename = []
                for ext in [".h3f", ".h3i", ".h3m", ".h3p"]:
                    src_f = plain_lib_full_path + ext
                    if os.path.exists(src_f):
                        adjusted_files.append(src_f)
                        files_to_rename.append(src_f)
                        tgt_f = plain_lib_full_path + ".unzipped" + ext
                        adjusted_files.append(tgt_f)
                yaml_obj_config["files"] = files_dict

                if not dry_run:
                    yaml_obj._config = yaml_obj_config
                    if (db_type != 'dbcan') and (not os.path.exists(gz_lib_full_path)):
                        with open(plain_lib_full_path, "rb") as f_in:
                            with gzip.open(gz_lib_full_path, "wb") as f_out:
                                shutil.copyfileobj(f_in, f_out)
                    shutil.move(plain_lib_full_path, plain_lib_full_path + ".unzipped")
                    for f in files_to_rename:
                        path_wo_ext, ext = os.path.splitext(f)
                        tgt_f = path_wo_ext + ".unzipped" + ext
                        shutil.move(f, tgt_f)
        return adjusted_files

    def add_pfam_files_for_mapping(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type == 'pfams':
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_basename = os.path.basename(os.path.splitext(yaml_obj.get_yaml())[0])
                files_dict['clan_txt'] = yaml_basename + ".clan.txt.gz"
                files_dict['pfam_clans'] = yaml_basename + ".clans.tsv.gz"
            yaml_obj_config["files"] = files_dict

            if not dry_run:
                yaml_obj._config = yaml_obj_config
        return adjusted_files

    def add_tigrfam_files_for_mapping(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type == 'tigrfams':
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_basename = os.path.basename(os.path.splitext(yaml_obj.get_yaml())[0])
                info_file_basename = yaml_basename.rstrip("_HMM")
                files_dict['role_name'] = info_file_basename + "_ROLE_NAMES"
                files_dict['role_link'] = info_file_basename + "_ROLE_LINK"
                files_dict['info_tar'] = info_file_basename + "_INFO.tar.gz"
            yaml_obj_config["files"] = files_dict

            if not dry_run:
                yaml_obj._config = yaml_obj_config
        return adjusted_files

    def add_dbcan_files_for_mapping(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type == 'dbcan':
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_basename = os.path.basename(os.path.splitext(yaml_obj.get_yaml())[0])
                files_dict['map_res'] = yaml_basename + ".fam-activities.txt"
            yaml_obj_config["files"] = files_dict

            if not dry_run:
                yaml_obj._config = yaml_obj_config
        return adjusted_files

    def add_foam_files_for_mapping(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type == 'foam':
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_basename = os.path.basename(os.path.splitext(yaml_obj.get_yaml())[0])
                files_dict['map_res'] = yaml_basename + ".tsv"
            yaml_obj_config["files"] = files_dict

            if not dry_run:
                yaml_obj._config = yaml_obj_config
        return adjusted_files

    def remove_resfam_tsv(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']

            if db_type in ['resfams_core', 'resfams_full']:
                yaml_path_wo_ext = os.path.splitext(yaml_obj.get_yaml())[0]
                map_res_tsv = yaml_path_wo_ext + ".tsv"
                adjusted_files.append(map_res_tsv)

        if not dry_run:
            for f in adjusted_files:
                os.remove(f)
        return adjusted_files

    def add_resfam_files_for_mapping(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()

            db_type = yaml_obj_config['type']
            files_dict = yaml_obj_config["files"]

            if db_type in ['resfams_core', 'resfams_full']:
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_basename = os.path.basename(os.path.splitext(yaml_obj.get_yaml())[0])
                files_dict['map_res_xlsx'] = yaml_basename + ".xlsx"
            yaml_obj_config["files"] = files_dict

            if not dry_run:
                yaml_obj._config = yaml_obj_config
        return adjusted_files

    def add_sp_and_tr_to_uniprot_complete_requirements(self, dry_run=False):
        adjusted_files = []
        uniprot_complete_objs = self.get_yamls_objs_matching_db_types(["uniprot_complete"])
        uniprot_sp_objs = self.get_yamls_objs_matching_db_types(["uniprot_sp"])
        uniprot_tr_objs = self.get_yamls_objs_matching_db_types(["uniprot_tr"])
        for o in uniprot_complete_objs:
            adjusted_files.append(o.get_yaml())
            cfg = o.get_config_deepcopy()
            cfg["required_dbs"]["uniprot_sp"] = {}
            cfg["required_dbs"]["uniprot_tr"] = {}

            ver = cfg["version"]

            lst_matching_sp_o = [sp_o for sp_o in uniprot_sp_objs if sp_o.get_config_deepcopy()["version"] == ver]
            if len(lst_matching_sp_o) != 1:
                raise ValueError(f"Expected exactly ony Uniprot swissprot database yaml of version {ver}\n"
                                 f"got: {len(lst_matching_sp_o)}")
            cfg["required_dbs"]["uniprot_sp"]["version"] = str(ver)

            lst_matching_tr_o = [tr_o for tr_o in uniprot_tr_objs if tr_o.get_config_deepcopy()["version"] == ver]
            if len(lst_matching_tr_o) != 1:
                raise ValueError(f"Expected exactly ony Uniprot swissprot database yaml of version {ver}\n"
                                 f"got: {len(lst_matching_tr_o)}")
            cfg["required_dbs"]["uniprot_tr"]["version"] = str(ver)

            if not dry_run:
                o._config = cfg
        return adjusted_files

    def remove_uniprot_complete_files(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            db_type = yaml_obj_config['type']
            if db_type == 'uniprot_complete':
                yaml_path = yaml_obj.get_yaml()
                adjusted_files.append(yaml_path)
                rel_fasta_path = yaml_obj_config['files']['fasta_gz']
                abs_fasta_path = os.path.join(os.path.dirname(yaml_path), rel_fasta_path)
                adjusted_files.append(abs_fasta_path)
                del yaml_obj_config['files']['fasta_gz']
                if not dry_run:
                    os.remove(abs_fasta_path)
                    yaml_obj._config = yaml_obj_config
        return adjusted_files

    def remove_redundant_uniprot_dependancies(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            db_type = yaml_obj_config['type']
            if db_type in ['uniprot_complete', 'uniprot_sp', 'uniprot_tr']:
                yaml_path = yaml_obj.get_yaml()
                adjusted_files.append(yaml_path)
                del yaml_obj_config['required_dbs']['taxdump']
                del yaml_obj_config['required_dbs']['accession2taxid']
                if not dry_run:
                    yaml_obj._config = yaml_obj_config
        return adjusted_files


class MigrateV04toV05(Migrator):
    """
    move file to subdir, depending on their origin:
        files downloaded by url: direct_dl
        files downloaded by emapper: emapper_dl
        output files of processing steps: processed
    bump "db_schema_version: 5" in yaml
    """
    from_version = 4
    to_version = 5

    @property
    def list_of_ordered_migration_methods(self):
        return [
            self.remove_non_mandatory_files,
            self.move_dl_files_to_subdirs,
            self.rename_db_yaml,
            self.rename_ncbi_helper_dbs
        ]

    def remove_non_mandatory_files(self, dry_run=False):
        """
        keep only db_config.yaml and the files that are mentioned in the db_config.yaml, remove everything else in each
        db-folder

        :param dry_run:
        :return:
        """
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            paths_to_keep = set()
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            yaml_f_abs = yaml_obj.get_yaml()
            paths_to_keep.add(yaml_f_abs)
            db_dir = os.path.dirname(yaml_f_abs)

            paths_in_db_dir = {os.path.join(db_dir, p) for p in os.listdir(db_dir)}
            paths_to_keep.update({os.path.join(db_dir, f) for f in yaml_obj_config["files"].values()})

            paths_to_remove = paths_in_db_dir - paths_to_keep
            adjusted_files.extend(paths_to_remove)
            if not dry_run:
                for p in paths_to_remove:
                    try:
                        os.remove(p)
                    except OSError:
                        os.rmdir(p)
        return adjusted_files

    def move_dl_files_to_subdirs(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            adjusted_files.append(yaml_obj.get_yaml())

            db_type = yaml_obj_config['type']
            db_dir = os.path.dirname(yaml_obj.get_yaml())
            files_dict = yaml_obj_config["files"]
            lst_src_files = []
            lst_tgt_files = []
            new_files_dict = {}
            for file_key, rel_src_file in files_dict.items():
                if (db_type == "eggnog") and (file_key == "emapper_data"):
                    dl_sub_dir = "emapper_dl"
                else:
                    dl_sub_dir = "direct_dl"
                rel_tgt_file = os.path.join(dl_sub_dir, rel_src_file)
                abs_src_file = os.path.join(db_dir, rel_src_file)
                abs_tgt_file = os.path.join(db_dir, rel_tgt_file)
                lst_src_files.append(abs_src_file)
                lst_tgt_files.append(abs_tgt_file)
                new_files_dict[file_key] = rel_tgt_file
            yaml_obj_config["files"] = new_files_dict
            adjusted_files.extend(lst_src_files + lst_tgt_files)
            for path in lst_src_files:
                assert os.path.exists(path), f"Required file does not exist: {path}"
            if not dry_run:
                yaml_obj._config = yaml_obj_config
                for f_src, f_tgt in zip(lst_src_files, lst_tgt_files):
                    tgt_dir = os.path.dirname(f_tgt)
                    if not os.path.exists(tgt_dir):
                        os.mkdir(tgt_dir)
                    shutil.move(f_src, f_tgt)
        return adjusted_files

    def rename_db_yaml(self, dry_run=False):
        adjusted_files = []
        for yaml_obj in self.yaml_objs_to_migrate:
            yaml_path = yaml_obj.get_yaml()
            src = yaml_path
            yaml_dir = os.path.dirname(yaml_path)
            yaml_ext = os.path.splitext(yaml_path)[1]
            new_path = os.path.join(yaml_dir, "db_config" + yaml_ext)
            tgt = new_path
            adjusted_files.extend([src, tgt])
            if not dry_run:
                yaml_obj.yaml = new_path
                tgt_dir = os.path.dirname(tgt)
                if not os.path.exists(tgt_dir):
                    os.mkdir(tgt_dir)
                shutil.move(src, tgt)
        return adjusted_files

    def rename_ncbi_helper_dbs(self, dry_run=False):
        """
        rename accession2taxid to ncbi_accession2taxid
        rename taxdump to ncbi_taxdump
        replace types in cfg files
        replace required_dbs in ncbi_nr:
          ncbi_taxdump:
            version: "{version}"
          ncbi_accession2taxid:
            version: "{version}"
        """
        adjusted_files = []
        is_db_renamed = {
            'accession2taxid': False,
            'taxdump': False
        }
        for yaml_obj in self.yaml_objs_to_migrate:
            lst_src_tgt_dir_tuples = []
            yaml_obj_config = yaml_obj.get_config_deepcopy()
            db_type = yaml_obj_config['type']
            tgt_yaml = yaml_obj.yaml
            if db_type == 'ncbi_nr':
                adjusted_files.append(yaml_obj.get_yaml())
                yaml_obj_config['required_dbs'] = {
                    'ncbi_taxdump': yaml_obj_config['required_dbs']['taxdump'],
                    'ncbi_accession2taxid': yaml_obj_config['required_dbs']['accession2taxid']
                }
            if db_type in ['accession2taxid', 'taxdump']:
                new_db_type = 'ncbi_' + db_type
                yaml_obj_config['type'] = new_db_type
                src_yaml = yaml_obj.get_yaml()
                full_path_up_to_version, yaml_basename = os.path.split(src_yaml)
                src_dir, subdir_version = os.path.split(full_path_up_to_version)
                full_path_up_to_base_dir, _ = os.path.split(src_dir)
                tgt_dir = os.path.join(full_path_up_to_base_dir, new_db_type)
                tgt_yaml = os.path.join(tgt_dir, subdir_version, yaml_basename)
                adjusted_files.extend([src_dir, tgt_dir, src_yaml, tgt_yaml])
                if not is_db_renamed[db_type]:
                    lst_src_tgt_dir_tuples.append((src_dir, tgt_dir))
                    is_db_renamed[db_type] = True
            if not dry_run:
                yaml_obj._config = yaml_obj_config
                yaml_obj.yaml = tgt_yaml
                for d_src, d_tgt in lst_src_tgt_dir_tuples:
                    os.rename(d_src, d_tgt)
        return adjusted_files
