#!/usr/bin/env python

from Bio import SeqIO
import sys
import argparse

parser = argparse.ArgumentParser(
    description=(
        "This script facilitates renaming headers in a fasta file. "
        "By default, each sequence is renamed to <wanted_text>_<n> for n=1,2,… "
        "If you provide --prefix and/or --suffix, those strings will be "
        "prepended/appended to each original header. "
        "For version info, run `bit-version`."
    )
)

required = parser.add_argument_group('required arguments')

required.add_argument("-i", "--input-fasta", metavar = "<FILE>", help = "starting fasta file", action = "store", required = True)
parser.add_argument("-w", "--wanted-text", metavar = "<STR>", help = 'base name to give seqs when renaming to "<wanted_text>_<n>" (default: "Seq")', action = "store", default = "Seq")
parser.add_argument("-o", "--output-fasta", metavar = "<FILE>", help = 'output fasta file (default: "renamed.fasta").', default = "renamed.fasta")

parser.add_argument("--prefix", metavar = "<STR>", help = "prepend this text to the original header (include separator if wanted)", default = "")
parser.add_argument("--suffix", metavar = "<STR>", help = "append this text to the original header (include separator if wanted)", default = "")


if len(sys.argv)==1:
    parser.print_help(sys.stderr)
    sys.exit(0)

args = parser.parse_args()

do_numbering = (args.prefix == "" and args.suffix == "")

try:
    input_fasta = open(args.input_fasta, "r")
except FileNotFoundError:
    print(f"Error: Input file '{args.input_fasta}' not found.")
    sys.exit(1)

out_fasta = open(args.output_fasta, "w")

if do_numbering:
    counter = 0
    for seq_record in SeqIO.parse(input_fasta, "fasta"):
        counter += 1
        out_fasta.write(f">{args.wanted_text}_{counter}\n")
        out_fasta.write(str(seq_record.seq) + "\n")

else:

    for seq_record in SeqIO.parse(input_fasta, "fasta"):
        # using description so it preserves the full header
        original_header = seq_record.description

        new_header = f"{args.prefix}{original_header}{args.suffix}"
        out_fasta.write(f">{new_header}\n{seq_record.seq}\n")

input_fasta.close()
out_fasta.close()
