from utils.exceptions import CustomMapError


def run_acc2annot_mapper(query_accessions, custom_acc2tax_map, result_path):
    """
    the custom_acc2tax_map file must be a  tab-separated table (tsv) file, with the first column holding accessions of
    the same format as those the query_accessions, all additional columns will be used for annotation

    Example custom_acc2tax_map content:
    acc	name_of_level1	name_of_level2	name_of_level3
    KZK33282.1	Bacteria	Firmicutes	Bacilli
    NP_268346.1	Bacteria	Firmicutes	Bacilli
    WP_003131952.1	Bacteria	Firmicutes	Bacilli
    PKC80086.1	Bacteria	Actinobacteria	Bifidobacteriales
    XP_963622.1	Eukaryota	Ascomycota	Sordariomycetes
    WP_000665196.1	Bacteria	Proteobacteria	Gammaproteobacteria
    WP_000665196.1	Bacteria	Proteobacteria	Gammaproteobacteria group

    :param query_accessions: iterable of accession-string to annotate
    :param custom_acc2tax_map: tsv-table
    :param result_path: path to output tsv
    :return: None
    """
    # build a tsv with columns:
    # [#acc, source, hit, annotation_info[0], annotation_info[1], ...]
    with open(custom_acc2tax_map, 'r') as custom_map_handle:
        column_names = custom_map_handle.readline().strip().split("\t")
        no_cols = len(column_names)
        if no_cols <= 1:
            raise CustomMapError("Expected at least one annotation column in custom map, found none. custom map: "
                                 f"{custom_acc2tax_map}")
        # drop name of accession column
        column_names = column_names[1:]
        headline = "\t".join(["#acc", "source", "hit"] + column_names) + "\n"
        with open(result_path, 'w') as outhandle:
            outhandle.write(headline)
            for line in custom_map_handle:
                fields = line.split("\t")
                if not len(fields) == no_cols:
                    raise CustomMapError(f"Malformed line in column map '{custom_acc2tax_map}'\n"
                                         + f"Expected: {no_cols} columns\n"
                                         + f"Got {len(fields)} columns")
                acc = fields[0].strip()
                if acc in query_accessions:
                    line = [acc, "custom_map"] + fields
                    outhandle.write("\t".join(line))
