#!/usr/bin/env python

import argparse
import re
import sys

from Bio import SeqIO

parser = argparse.ArgumentParser(
    formatter_class=argparse.ArgumentDefaultsHelpFormatter,
    description=(
        "Read a GenBank flat file and print GenBank or FASTA for matching records."
    ),
)

parser.add_argument(
    "genbankFile", metavar="FILE.gb", nargs="+", help="The GenBank file(s) to examine."
)

parser.add_argument(
    "--idRegex",
    metavar="REGEX",
    default=".",
    help=(
        "Only print records whose id (or description, if --matchDescription "
        "is used) matches this regular expression."
    ),
)

parser.add_argument(
    "--match",
    choices=("fullmatch", "match", "search"),
    default="search",
    help="How to match the regular expression. See pydoc re for an explanation.",
)

parser.add_argument(
    "--matchDescription",
    action="store_true",
    help="Match on the record description as well as the id.",
)

parser.add_argument(
    "--invert",
    "-v",
    action="store_true",
    help="Invert the match. Only print records that do not match.",
)

parser.add_argument(
    "--ignoreCase", "-i", action="store_true", help="Use case-insensitive matching."
)

parser.add_argument(
    "--format", choices=("fasta", "gb"), default="gb", help="The output format."
)

parser.add_argument(
    "--expectedCount",
    type=int,
    help=(
        "The expected number of records to find. If this number is not "
        "found, print a message and exit non-zero."
    ),
)

parser.add_argument(
    "--verbose",
    action="store_true",
    help="Print a summary of records read and printed.",
)

args = parser.parse_args()

recordCount = matchCount = 0
match = getattr(
    re.compile(args.idRegex, flags=(re.I if args.ignoreCase else 0)), args.match
)

matchDescription = args.matchDescription
invert = args.invert


def filterFunc(records):
    global matchCount, recordCount
    for record in records:
        recordCount += 1
        id_ = record.id + (f" {record.description}" if record.description else "")
        if match(id_ if matchDescription else record.id):
            emit = not invert
        else:
            emit = invert

        if emit:
            matchCount += 1
            yield record


for filename in args.genbankFile:
    with open(filename) as fp:
        SeqIO.write(filterFunc(SeqIO.parse(fp, "gb")), sys.stdout, args.format)

if args.verbose:
    print(
        f"Read {recordCount} record{'s' if recordCount > 1 else ''}, "
        f"printed {matchCount}",
        file=sys.stderr,
    )

if args.expectedCount is not None and matchCount != args.expectedCount:
    print(
        f"Expected {args.expectedCount} records but found {matchCount}. Exiting.",
        file=sys.stderr,
    )
    sys.exit(1)
