import os
import yaml

# --------------------------------------------------------------------------------
# CONFIGURATION LOADING
# --------------------------------------------------------------------------------

# Load paths and parameters from the external config file
with open("config.yaml", "r") as file:
    config = yaml.safe_load(file)

# --------------------------------------------------------------------------------
# HELPER FUNCTIONS
# --------------------------------------------------------------------------------

# Function to scan the input 16S directory and return a list of genome base names
def get_genomes():
    genome_dir_16S = config["genome_dir_16S"]
    # List files ending in .fasta and remove the extension for wildcard usage
    return [f.replace(".fasta", "") for f in os.listdir(genome_dir_16S) if f.endswith(".fasta")]

# --------------------------------------------------------------------------------
# TARGET RULE
# --------------------------------------------------------------------------------

# The default rule that defines the final desired outputs of the workflow
rule all:
    input:
        # Generates a BLAST CSV result for every genome found in the input directory
        expand(os.path.join(config["blast_dir"], "{genome}.blast.csv"), genome=get_genomes())

# --------------------------------------------------------------------------------
# WORKFLOW RULES
# --------------------------------------------------------------------------------

# Rule to download and prepare the NCBI BLAST 16S database if it does not exist locally
rule setup_blastdb:
    output:
        # Tracks one of the specific BLAST index files to verify installation
        os.path.join(config["blast_db_dir"], config["blast_db_name"] + ".nhr") 
    conda:
        "./envs/blast.yaml" 
    shell:
        """
        mkdir -p {config[blast_db_dir]}
        export BLASTDB={config[blast_db_dir]}
        cd {config[blast_db_dir]}
        
        # Check if the database index file exists; if not, download and extract from NCBI
        if [ ! -f {config[blast_db_name]}.nhr ]; then
            wget ftp://ftp.ncbi.nlm.nih.gov/blast/db/{config[blast_db_name]}.tar.gz
            tar -xzf {config[blast_db_name]}.tar.gz
        else
            echo "BLAST database already exists, skipping download."
        fi
        """

# Rule to perform the BLASTn search for the query 16S sequences
rule blast:
    input:
        # Dynamically locate the input FASTA based on the genome wildcard
        fasta=lambda wildcards: os.path.join(config["genome_dir_16S"], f"{wildcards.genome}.fasta"),
        # Dependency on the database setup rule
        db_installed=os.path.join(config["blast_db_dir"], config["blast_db_name"] + ".nhr")
    output:
        # Individual CSV output for each genome
        os.path.join(config["blast_dir"], "{genome}.blast.csv")
    params:
        # Full path to the BLAST database prefix
        db=os.path.join(config["blast_db_dir"], config["blast_db_name"])
    conda:
        "./envs/blast.yaml"
    shell:
        """
        # Create output directory if it doesn't exist
        mkdir -p "{config[output_dir]}"
        
        # Write CSV headers to the output file
        echo "query_id,subject_id,percentage_identity,alignment_length,mismatches,gap_opens,q_start,q_end,s_start,s_end,evalue,bit_score" > "{output}"
        
        # Execute blastn with custom comma-separated output format (10)
        blastn -query "{input.fasta}" -db "{params.db}" -outfmt "10 qseqid sseqid pident length mismatch gapopen qstart qend sstart send evalue bitscore" >> "{output}"
        """
