#!/usr/bin/env python
"""
Prepare input sequence file from fasta file in a directory for BPP.
"""

import argparse
import sys
from pathlib import Path
from Bio import SeqIO


def get_fasta_files(input_dir: Path, ext=".fasta") -> list[Path]:
    """
    Get a list of fasta files in the specified directory with the given extension.
    Args:
        input_dir (Path): Path to the directory containing fasta files.
        ext (str): Extension of the fasta files.
    Returns:
        list[Path]: List of Paths to the fasta files.
    """
    files = list(input_dir.glob(f"*{ext}"))
    if not files:
        print(f"No fasta files with extension '{ext}' found in directory '{input_dir}'")
        sys.exit(1)
    return files


def transform_fasta_to_bpp(fasta_file: Path):
    """
    Transform a fasta file into BPP sequence format.
    Args:
        fasta_file (Path): Path to the input fasta file.
    Returns:
        list[str]: Lines formatted for BPP input.
    """
    num_ind = 0
    output_lines = []
    seq_len = []
    for record in SeqIO.parse(fasta_file, "fasta"):
        ind_id = record.id
        num_ind += 1
        output_lines.append(f"^{ind_id}    {str(record.seq)}\n")
        seq_len.append(len(record.seq))
    if len(set(seq_len)) != 1:
        print(f"Error: Sequences in file '{fasta_file}' have different lengths.")
        sys.exit(1)
    header = f"{num_ind} {seq_len[0]}\n"
    return [header, "\n", *output_lines, "\n"]


def make_bpp_seq(input_dir: Path, output_file: Path, ext=".fasta"):
    """
    Create a BPP formatted sequence file from fasta files in a directory.
    Args:
        input_dir (Path): Path to the directory containing fasta files.
        output_file (Path): Path to the output BPP formatted file.
        ext (str): Extension of the fasta files.
    """
    fasta_files = get_fasta_files(input_dir, ext)
    with open(output_file, "w") as out_f:
        for fasta_file in fasta_files:
            bpp_lines = transform_fasta_to_bpp(fasta_file)
            out_f.writelines(bpp_lines)
    print(f"BPP formatted sequence file written to '{output_file}'")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(
        description="Create a BPP formatted sequence file from fasta files in a directory."
    )
    parser.add_argument(
        "input_dir", type=Path, help="Directory containing fasta files."
    )
    parser.add_argument(
        "output_file", type=Path, help="Output BPP formatted sequence file."
    )
    parser.add_argument(
        "--ext",
        type=str,
        default=".fasta",
        help="Extension of the fasta files (default: .fasta).",
    )
    args = parser.parse_args()
    make_bpp_seq(args.input_dir, args.output_file, args.ext)