#!/usr/bin/env python3

"""
blast_sequence_retriever.py

DESCRIPTION:
    Performs BLAST searches using one or more query sequences against a folder
    of genome/transcriptome FASTA files. For each query, retrieves the top N hits
    per genome file, optionally includes flanking regions, orients sequences to match
    the query strand, and saves results to a FASTA file.

USAGE:
    blast_sequence_retriever.py [-h]
                                   [--top_hits TOP_HITS]
                                   [--flanking_bp FLANKING_BP]
                                   [--min_evalue MIN_EVALUE]
                                   [-o OUTPUT_FOLDER]
                                   [-f FINAL_FASTA]
                                   [-l LOG_FILE]
                                   [--keep_db]
                                   {blastn,blastp,blastx,tblastn,tblastx}
                                   query_gene
                                   genome_folder                       
"""


import os
import subprocess
import argparse
from Bio import SeqIO
from Bio.Seq import Seq

# Argument parsing
parser = argparse.ArgumentParser(description="Run BLAST for multiple FASTA files, extract top hits with flanking regions, and orient sequences like the query.")

parser.add_argument("blast_type", choices=["blastn", "blastp", "blastx", "tblastn", "tblastx"],
                    help="Type of BLAST to use.")
parser.add_argument("query_gene", help="FASTA file with the query gene/protein.")
parser.add_argument("genome_folder", help="Folder with genome/transcriptome files.")
parser.add_argument("--top_hits", type=int, default=1, help="Number of top hits to extract per file (default: 1).")
parser.add_argument("--flanking_bp", type=int, default=1000, help="Flanking bases to include upstream/downstream (default: 1000).")
parser.add_argument("--min_evalue", type=float, default=1e-5, help="Minimum e-value threshold (default: 1e-5).")
parser.add_argument("-o", "--output_folder", default="blast_results", help="Output folder (default: blast_results).")
parser.add_argument("-f", "--final_fasta", default="final_top_hits.fasta", help="FASTA file to store all extracted hits.")
parser.add_argument("-l", "--log_file", default="blast_log.txt", help="Log file for summary (default: blast_log.txt).")
parser.add_argument("--keep_db", action="store_true", help="Keep intermediate BLAST database files.")

args = parser.parse_args()
os.makedirs(args.output_folder, exist_ok=True)

# Logging
with open(args.log_file, "w") as log:
    log.write(f"BLAST type: {args.blast_type}\nQuery: {args.query_gene}\nOutput folder: {args.output_folder}\n\n")

# Function to extract sequence + flanks
def extract_flanking_sequence(seq_record, start, end, flank):
    seq_len = len(seq_record.seq)
    left = max(0, min(start, end) - flank)
    right = min(seq_len, max(start, end) + flank)
    return seq_record.seq[left:right]

# Iterate over files in genome folder
for file in os.listdir(args.genome_folder):
    if not file.lower().endswith((".fa", ".fasta", ".fna")):
        continue

    base_name = os.path.splitext(file)[0]
    genome_path = os.path.join(args.genome_folder, file)
    db_prefix = os.path.join(args.output_folder, base_name)
    blast_output_file = os.path.join(args.output_folder, f"{base_name}.blast.out")

    # Create BLAST database
    db_type = "nucl" if args.blast_type in ["blastn", "blastx", "tblastx"] else "prot"
    subprocess.run([
        "makeblastdb",
        "-in", genome_path,
        "-dbtype", db_type,
        "-out", db_prefix
    ], check=True)

    # Run BLAST
    blast_cmd = [
        args.blast_type,
        "-query", args.query_gene,
        "-db", db_prefix,
        "-out", blast_output_file,
        "-outfmt", "6 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore"
    ]
    subprocess.run(blast_cmd, check=True)

    # Parse hits
    with open(blast_output_file) as bf:
        hits = [line.strip().split("\t") for line in bf if float(line.strip().split("\t")[10]) <= args.min_evalue]

    if not hits:
        continue

    top_hits = sorted(hits, key=lambda x: float(x[10]))[:args.top_hits]
    genome_dict = SeqIO.to_dict(SeqIO.parse(genome_path, "fasta"))

    # Extract and write sequences
    with open(args.final_fasta, "a") as out_fasta:
        for hit in top_hits:
            sseqid = hit[1]
            sstart, send = int(hit[8]), int(hit[9])
            strand = "plus" if sstart < send else "minus"

            if sseqid not in genome_dict:
                continue

            seq = extract_flanking_sequence(genome_dict[sseqid], sstart, send, args.flanking_bp)
            if strand == "minus":
                seq = seq.reverse_complement()

            header = f">{base_name}|{sseqid}_{sstart}_{send}|strand:{strand}"
            out_fasta.write(f"{header}\n{seq}\n")

    with open(args.log_file, "a") as log:
        log.write(f"{file}: extracted {len(top_hits)} hit(s)\n")

    # Clean up BLAST DB files
    if not args.keep_db:
        for ext in [".nhr", ".nin", ".nsq", ".phr", ".pin", ".psq", ".pal", ".nal"]:
            db_file = db_prefix + ext
            if os.path.exists(db_file):
                os.remove(db_file)

print("✅ All BLAST searches completed and sequences extracted.")