"""
Run IgBLAST and output the AIRR-formatted result table
"""

# The code in here was copied from IgDiscover 0.15, https://github.com/NBISweden/IgDiscover-legacy
# and streamlined a bit to work stand-alone. Relevant files:
# - src/igdiscover/cli/igblastwrap.py
# - src/igdiscover/igblast.py
# - src/igdiscover/species.py
# - src/igdiscover/utils.py
# - src/igdiscover/dna.py

# Inline dependencies. Makes this script runnable with, e.g., 'pipx run'
# /// script
# dependencies = ["dnaio"]
# ///

# Code has been fixed to populate CDR3 and FR4 fields in the output using regex logic that existed but was not used.
# Relevant code needed to augment IgBLAST output was adapted from:
# - src/igdiscover/cli/augment.py

import csv
import logging
import multiprocessing
import os
import re
import shlex
import shutil
import subprocess
import sys
import tempfile
import time
from argparse import ArgumentParser
from contextlib import ExitStack
from dataclasses import dataclass
from io import StringIO
from itertools import islice
from pathlib import Path
from typing import Dict, Optional

import dnaio

GENETIC_CODE = {
    "AAA": "K",
    "AAC": "N",
    "AAG": "K",
    "AAT": "N",
    "ACA": "T",
    "ACC": "T",
    "ACG": "T",
    "ACT": "T",
    "AGA": "R",
    "AGC": "S",
    "AGG": "R",
    "AGT": "S",
    "ATA": "I",
    "ATC": "I",
    "ATG": "M",
    "ATT": "I",
    "CAA": "Q",
    "CAC": "H",
    "CAG": "Q",
    "CAT": "H",
    "CCA": "P",
    "CCC": "P",
    "CCG": "P",
    "CCT": "P",
    "CGA": "R",
    "CGC": "R",
    "CGG": "R",
    "CGT": "R",
    "CTA": "L",
    "CTC": "L",
    "CTG": "L",
    "CTT": "L",
    "GAA": "E",
    "GAC": "D",
    "GAG": "E",
    "GAT": "D",
    "GCA": "A",
    "GCC": "A",
    "GCG": "A",
    "GCT": "A",
    "GGA": "G",
    "GGC": "G",
    "GGG": "G",
    "GGT": "G",
    "GTA": "V",
    "GTC": "V",
    "GTG": "V",
    "GTT": "V",
    # 'TAA': stop
    "TAC": "Y",
    # 'TAG': stop,
    "TAT": "Y",
    "TCA": "S",
    "TCC": "S",
    "TCG": "S",
    "TCT": "S",
    # 'TGA': stop
    "TGC": "C",
    "TGG": "W",
    "TGT": "C",
    "TTA": "L",
    "TTC": "F",
    "TTG": "L",
    "TTT": "F",
}

logger = logging.getLogger(__name__)


def nt_to_aa(s, _get=GENETIC_CODE.get):
    """Translate a nucleotide sequence to an amino acid sequence"""
    return "".join([_get(s[i : i + 3], "*") for i in range(0, len(s), 3)])


class SerialPool:
    """
    An alternative to multiprocessing.Pool that runs things in serial for
    easier debugging
    """

    def __init__(self, *args, **kwargs):
        pass

    def __enter__(self):
        return self

    def __exit__(self, *args):
        pass

    def imap(self, func, iterable, chunksize):
        for i in iterable:
            yield func(i)


def available_cpu_count():
    """
    Number of available virtual or physical CPUs on this system
    """
    if sys.platform != "linux":
        try:
            import multiprocessing

            return multiprocessing.cpu_count()
        except (ImportError, NotImplementedError):
            return 1
    return len(os.sched_getaffinity(0))


def escape_shell_command(command):
    return " ".join(shlex.quote(arg) for arg in command)


def run_igblast(sequences, blastdb_dir, species, sequence_type, penalty=None) -> str:
    """
    Run the igblastn command-line program.

    sequences -- list of Sequence objects
    blastdb_dir -- directory that contains BLAST databases. Files in that
    directory must be databases created by the makeblastdb program and have
    names V, D, and J.

    Return IgBLAST’s raw output as a string.
    """
    if sequence_type not in ("Ig", "TCR"):
        raise ValueError('sequence_type must be "Ig" or "TCR"')
    variable_arguments = []
    for gene in "V", "D", "J":
        variable_arguments += [
            "-germline_db_{gene}".format(gene=gene),
            os.path.join(blastdb_dir, "{gene}".format(gene=gene)),
        ]

    with tempfile.TemporaryDirectory() as tmpdir:
        tmp_path = Path(tmpdir)

        env = None
        # Work around internal_data/ not allowed to contain spaces on Windows
        if sys.platform == "win32" and "IGDATA" not in os.environ:
            if copy_internal_data(tmp_path):
                env = {**os.environ, "IGDATA": tmpdir}

        # An empty .aux suppresses a warning from IgBLAST. /dev/null does not work.
        empty_aux_path = tmp_path / "empty.aux"
        empty_aux_path.write_text("\n")
        arguments = []
        if penalty is not None:
            arguments += ["-penalty", str(penalty)]
        if species is not None:
            arguments += ["-organism", species]

        arguments += [
            "-auxiliary_data",
            str(empty_aux_path),
            "-ig_seqtype",
            sequence_type,
            "-num_threads",
            "1",
            "-domain_system",
            "imgt",
            "-num_alignments_V",
            "1",
            "-num_alignments_D",
            "1",
            "-num_alignments_J",
            "1",
            "-outfmt",
            "19",  # AIRR format
            "-query",
            "-",
        ]
        fasta_str = "".join(">{}\n{}\n".format(r.name, r.sequence) for r in sequences)

        # For some reason, it has become unreliable to let IgBLAST 1.10 write its result
        # to standard output using "-out -". The data becomes corrupt. This does not occur
        # when calling igblastn in the shell and using the same syntax, only with
        # subprocess.check_output. As a workaround, we write the output to a temporary file.
        path = os.path.join(tmpdir, "igblast.txt")
        command = ["igblastn"] + variable_arguments + arguments + ["-out", path]
        logger.debug("Running %s", escape_shell_command(command))
        output = subprocess.check_output(
            command, input=fasta_str, universal_newlines=True, env=env
        )
        assert output == ""
        with open(path) as f:
            return f.read()


def copy_internal_data(tmp_path) -> bool:
    """
    Guess where igblastn’s internal_data/ directory is located and copy it
    recursively to the target directory.

    igblastn on Windows does not recognize internal_data/ directories whose
    absolute path contains spaces (even though the standard installation
    location is a subdirectory of C:\\Program Files\\).

    This works only if the target (tmp_path) does not contain spaces.

    Return True iff the internal_data/ directory was found and copied.
    """
    igblastn_path = Path(shutil.which("igblastn"))
    installation_path = igblastn_path.parent.parent
    internal_data_path = installation_path / "internal_data"
    if not internal_data_path.exists():
        return False

    shutil.copytree(internal_data_path, tmp_path / "internal_data")
    logger.debug("Copied %s to %s", internal_data_path, tmp_path / "internal_data")
    return True


def chunked(iterable, chunksize: int):
    """
    Group the iterable into lists of length chunksize
    >>> list(chunked('ABCDEFG', 3))
    [['A', 'B', 'C'], ['D', 'E', 'F'], ['G']]
    """
    chunk = []
    for it in iterable:
        if len(chunk) == chunksize:
            yield chunk
            chunk = []
        chunk.append(it)
    if chunk:
        yield chunk


class RawRunner:
    """
    This is the target of a multiprocessing pool. The target needs to
    be pickleable, and because nested functions cannot be pickled,
    we need this separate class.

    It runs IgBLAST and returns raw AIRR-formatted output
    """

    def __init__(self, blastdb_dir, species, sequence_type, penalty, database, cache):
        self.blastdb_dir = blastdb_dir
        self.species = species
        self.sequence_type = sequence_type
        self.penalty = penalty
        self.database = database
        self.cache = cache

    def __call__(self, sequences):
        """
        Return raw IgBLAST output
        """
        return run_igblast(
            sequences, self.blastdb_dir, self.species, self.sequence_type, self.penalty
        )


class MakeBlastDbError(subprocess.CalledProcessError):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __str__(self):
        return (
            f"Running '{escape_shell_command(self.cmd)}' failed with "
            f"exit code {self.returncode}. "
            f"Standard output:\n{self.output.decode()}\n"
            f"Standard error:\n{self.stderr.decode()}"
        )


def makeblastdb(fasta, database_name, prefix=""):
    """
    prefix -- prefix to add to sequence ids
    """
    n = 0
    with dnaio.open(fasta) as fr, open(database_name + ".fasta", "w") as db:
        for record in fr:
            name = prefix + record.name.split(maxsplit=1)[0]
            db.write(">{}\n{}\n".format(name, record.sequence))
            n += 1
    if n == 0:
        raise ValueError("FASTA file {} is empty".format(fasta))

    command = [
        "makeblastdb",
        "-parse_seqids",
        "-dbtype",
        "nucl",
        "-in",
        database_name + ".fasta",
        "-out",
        database_name,
    ]
    logger.debug("Running %s", escape_shell_command(command))
    try:
        process_output = subprocess.check_output(command, stderr=subprocess.PIPE)
    except subprocess.CalledProcessError as e:
        raise MakeBlastDbError(
            e.returncode, command, output=e.output, stderr=e.stderr
        ) from None

    if b"Error: " in process_output:
        raise MakeBlastDbError(0, command, stderr=process_output) from None


class Database:
    def __init__(self, path, sequence_type):
        """path -- path to database directory with V.fasta, D.fasta, J.fasta"""
        self.path = path
        self.sequence_type = sequence_type
        self._v_records = self._read_fasta(os.path.join(path, "V.fasta"))
        self.v = self._records_to_dict(self._v_records)
        self._d_records = self._read_fasta(os.path.join(path, "D.fasta"))
        self.d = self._records_to_dict(self._d_records)
        self._j_records = self._read_fasta(os.path.join(path, "J.fasta"))
        self.j = self._records_to_dict(self._j_records)
        self._cdr3_starts = dict()
        self._cdr3_ends = dict()
        for locus in ("IGH", "IGK", "IGL", "TRA", "TRB", "TRG", "TRD"):
            self._cdr3_starts[locus] = {
                name: cdr3_start(s, locus) for name, s in self.v.items()
            }
            self._cdr3_ends[locus] = {
                name: cdr3_end(s, locus) for name, s in self.j.items()
            }
        self.v_regions_nt, self.v_regions_aa = self._find_v_regions()

    @staticmethod
    def _read_fasta(path):
        records = []
        with dnaio.open(path) as sr:
            for record in sr:
                record.name = record.name.split(maxsplit=1)[0]
                records.append(record)
        return records

    @staticmethod
    def _records_to_dict(records):
        return {record.name: record.sequence.upper() for record in records}

    def v_cdr3_start(self, gene, locus):
        return self._cdr3_starts[locus][gene]

    def j_cdr3_end(self, gene, locus):
        return self._cdr3_ends[locus][gene]

    def _find_v_regions(self):
        """
        Run IgBLAST on the V sequences to determine the nucleotide and amino-acid sequences of the
        FR1, CDR1, FR2, CDR2 and FR3 regions
        """
        v_regions_nt = dict()
        v_regions_aa = dict()
        for record in igblast_records(self.path, self._v_records, self.sequence_type):
            nt_regions = dict()
            aa_regions = dict()
            for region in ("FR1", "CDR1", "FR2", "CDR2", "FR3"):
                nt_seq = record.region_sequence(region)
                if nt_seq is None:
                    break
                if len(nt_seq) % 3 != 0:
                    logger.warning(
                        "Length %s of %s region in %r is not divisible by three; region "
                        "info for this gene will not be available",
                        len(nt_seq),
                        region,
                        record.query_name,
                    )
                    # not codon-aligned, skip entire record
                    break
                nt_regions[region] = nt_seq
                try:
                    aa_seq = nt_to_aa(nt_seq)
                except ValueError as e:
                    logger.warning(
                        "The %s region could not be converted to amino acids: %s",
                        region,
                        str(e),
                    )
                    break
                if "*" in aa_seq:
                    logger.warning(
                        "The %s region in %r contains a stop codon (%r); region info "
                        "for this gene will not be available",
                        region,
                        record.query_name,
                        aa_seq,
                    )
                    break
                aa_regions[region] = aa_seq
            else:
                v_regions_nt[record.query_name] = nt_regions
                v_regions_aa[record.query_name] = aa_regions

        return v_regions_nt, v_regions_aa


def igblast_records(
    database, sequences, sequence_type, species=None, penalty=None, use_cache=False
):
    """
    Run IgBLAST, parse results and yield RegionRecords objects.

    database -- Path to database directory with V./D./J.fasta files
    sequences -- an iterable of Sequence objects
    sequence_type -- 'Ig' or 'TCR'
    """
    # Create the BLAST databases in a temporary directory
    with tempfile.TemporaryDirectory() as blastdb_dir:
        make_vdj_blastdb(blastdb_dir, database)
        igblast_result = run_igblast(
            sequences, blastdb_dir, species, sequence_type, penalty
        )
        sio = StringIO(igblast_result)
        for record in parse_region_records(sio):
            yield record


def parse_region_records(file):
    csv.register_dialect(
        "airr",
        delimiter="\t",
        lineterminator="\n",
        strict=True,
    )
    reader = csv.DictReader(file, dialect="airr")
    for record in reader:
        yield RegionsRecord(query_name=record["sequence_id"], fields=record)


@dataclass
class RegionsRecord:
    query_name: str
    fields: Dict[str, str]

    COLUMNS_MAP = {
        "CDR1": "cdr1",
        "CDR2": "cdr2",
        "FR1": "fwr1",
        "FR2": "fwr2",
        "FR3": "fwr3",
    }

    def region_sequence(self, region: str) -> str:
        """
        Return the nucleotide sequence of a named region. Allowed names are:
        CDR1, CDR2, FR1, FR2, FR3. Sequences are extracted from the full read
        using begin and end coordinates from IgBLAST’s "alignment summary" table.
        """
        if region not in self.COLUMNS_MAP:
            raise KeyError(f"Region '{region}' not allowed")
        return self.fields[self.COLUMNS_MAP[region]]


def igblast_parallel_chunked(
    database,
    sequences,
    sequence_type,
    species=None,
    threads=1,
    penalty=None,
    cache=None,
):
    """
    Run IgBLAST on the input sequences and yield AIRR-formatted results.

    The input is split up into chunks of 1000 sequences and distributed to
    *threads* number of IgBLAST instances that run in parallel.

    database -- Path to database directory with V./D./J.fasta files
    sequences -- an iterable of Sequence objects
    sequence_type -- 'Ig' or 'TCR'
    threads -- number of threads.
    """
    with ExitStack() as stack:
        # Create the three BLAST databases in a temporary directory
        blastdb_dir = stack.enter_context(tempfile.TemporaryDirectory())
        make_vdj_blastdb(blastdb_dir, database)

        chunks = chunked(sequences, chunksize=1000)
        runner = RawRunner(
            blastdb_dir,
            species=species,
            sequence_type=sequence_type,
            penalty=penalty,
            database=database,
            cache=cache,
        )
        pool = stack.enter_context(
            multiprocessing.Pool(threads) if threads > 1 else SerialPool()
        )
        for igblast_output in pool.imap(runner, chunks, chunksize=1):
            yield igblast_output


def make_vdj_blastdb(blastdb_dir, database_dir):
    """Run makeblastdb for all {V,D,J}.fasta in the database_dir"""

    for gene in ["V", "D", "J"]:
        # Without adding the "%" prefix, IgBLAST reports record names that look like GenBank
        # ids as "gb|original_name|".
        makeblastdb(
            os.path.join(database_dir, gene + ".fasta"),
            os.path.join(blastdb_dir, gene),
            prefix="%",
        )

################################################################################
# The code below was not previously used, now it is called to populate CDR3 and FR4
# fields in the output. See new code below this section.
################################################################################
# Regular expressions for CDR3 detection
#
# The idea comes from D’Angelo et al.: The antibody mining toolbox.
# http://dx.doi.org/10.4161/mabs.27105
# The heavy-chain regex was taken directly from there, but the difference
# is that we express everything in terms of amino acids, not nucleotides.
# This simplifies the expressions and makes them more readable.
#
_CDR3_REGEX = {
    # Heavy chain
    "IGH": re.compile(
        """
        [FY] [FHVWY] C
        (?P<cdr3>
            [ADEGIKMNRSTV] .{3,31}
        )
        W[GAV]
        """,
        re.VERBOSE,
    ),
    # Light chain, kappa
    "IGK": re.compile(
        """
        [FSVY] [CFHNVY] [CDFGLSW]
        (?P<cdr3>
            .{4,15}
        )
        [FLV][GRV]
        """,
        re.VERBOSE,
    ),
    # Light chain, lambda
    "IGL": re.compile(
        """
        # the negative lookahead assertion ensures that the rightmost start is found
        [CDY](?![CDY][CFHSY][CFGW])[CFHSY][CFGW]
        (?P<cdr3>
            .{4,15}
        )
        [FS]G
        """,
        re.VERBOSE,
    ),
}

_CDR3_VH_ALTERNATIVE_REGEX = re.compile(
    """
        C
        (?P<cdr3> . [RK] .{3,30})
        [WF]G.G
""",
    re.VERBOSE,
)


def find_cdr3(sequence, locus):
    """
    Find the CDR3 in the given sequence, assuming it comes from the given locus
    (chain type). If the locus is not one of 'IGH', 'IGK', 'IGL', return None.

    Return a tuple (start, stop) if found, None otherwise.
    """
    try:
        regex = _CDR3_REGEX[locus]
    except KeyError:
        return None
    matches = []
    for offset in 0, 1, 2:
        aa = nt_to_aa(sequence[offset:])
        match = regex.search(aa)
        if not match and locus == "IGH":
            match = _CDR3_VH_ALTERNATIVE_REGEX.search(aa)
        if match:
            start, stop = match.span("cdr3")
            matches.append((start * 3 + offset, stop * 3 + offset))
    return min(matches, default=None)


# The following code is used for detecting CDR3 start sites within V
# reference sequences and CDR3 end sites within J reference sequences.


# Matches the start of the CDR3 within the end of a VH sequence
_CDR3_START_VH_REGEX = re.compile(
    """
    [FY] [FHVWY] C
    (?P<cdr3_start>
        [ADEGIKMNRSTV*] | $
    )
    """,
    re.VERBOSE,
)


_CDR3_START_VH_ALTERNATIVE_REGEX = re.compile(
    """
    C
    (?P<cdr3_start> . [RK])
    """,
    re.VERBOSE,
)


_CDR3_START_REGEXES = {
    "IGK": re.compile("[FSVY][CFHNVY][CDFGLSW]"),
    "IGL": re.compile("[CDY](?![CDY][CFHSY][CFGW])[CFHSY][CFGW]"),
    "TRG": re.compile("[YFH]C"),  # TODO test whether this also works for alpha and beta
    "TRD": re.compile("[YFH]C"),
}


def _cdr3_start_heavy(aa):
    head, tail = aa[:-15], aa[-15:]
    match = _CDR3_START_VH_REGEX.search(tail)
    if not match:
        match = _CDR3_START_VH_ALTERNATIVE_REGEX.search(tail)
    if not match:
        return None
    return len(head) + match.start("cdr3_start")


def cdr3_start(nt, locus):
    """
    Find CDR3 start location within a V gene (Ig or TCR)

    nt -- nucleotide sequence of the gene
    locus -- one of "IGH", "IGK", "IGL", "TRA", "TRB", "TRG", "TRD"
    """
    aa = nt_to_aa(nt)
    if locus == "IGH":
        start = _cdr3_start_heavy(aa)
        if start is None:
            return None
        return 3 * start
    if locus in ("IGK", "IGL", "TRG", "TRD"):
        head, tail = aa[:-15], aa[-15:]
        match = _CDR3_START_REGEXES[locus].search(tail)
        if match:
            return 3 * (len(head) + match.end())
        else:
            return None
    elif locus in ("TRA", "TRB"):
        head, tail = aa[:-8], aa[-8:]
        pos = tail.find("C")
        if pos == -1:
            return None
        else:
            return 3 * (len(head) + pos + 1)


# Matches after the end of the CDR3 within a J sequence
_CDR3_END_REGEXES = {
    "IGH": re.compile("W[GAV]"),
    "IGK": re.compile("FG"),
    "IGL": re.compile("FG"),
    "TRA": re.compile("FG"),
    "TRB": re.compile("FG"),
    "TRG": re.compile("FG"),
    "TRD": re.compile("FG"),
}


def cdr3_end(nt, locus):
    """
    Find the position of the CDR3 end within a J sequence

    nt -- nucleotide sequence of the J gene
    locus -- one of "IGH", "IGK", "IGL", "TRA", "TRB", "TRG", "TRD"
    """
    regex = _CDR3_END_REGEXES[locus]
    for frame in 0, 1, 2:
        aa = nt_to_aa(nt[frame:])
        match = regex.search(aa)
        if match:
            return match.start() * 3 + frame
    return None


# When searching for the CDR3, start this many bases to the left of the end of
# the V match.
CDR3_SEARCH_START = 30


################################################################################
# This new code uses the regular expressions above to detect CDR3s and FR4.
# This is then used to populate the CDR3 and FR4 fields in the output.
################################################################################
# 
#
COLUMN_TYPES = {
    "v_sequence_start": int,
    "v_sequence_end": int,
    "v_germline_start": int,
    "v_germline_end": int,
    "j_sequence_start": int,
    "j_sequence_end": int,
    "j_germline_start": int,
    "j_germline_end": int,
    "cdr3_start" : int,
    "cdr3_end" : int
}

ALLOWED_LOCI = {"IGH", "IGK", "IGL", "TRA", "TRB", "TRG", "TRD"}

def set_cdr3_columns(record, database):
    if (
        not record["v_call"]
        or not record["j_call"]
        or record["locus"] not in ALLOWED_LOCI
    ):
        return

    # CDR3 start
    cdr3_ref_start = database.v_cdr3_start(
        record["v_call"], record["locus"]
    )
    if cdr3_ref_start is None:
        return
    cdr3_query_start = query_position(record, "v", reference_position=cdr3_ref_start)
    if cdr3_query_start is None:
        # Alignment is not long enough to cover CDR3 start position; try to rescue it
        # by assuming that the alignment would continue without indels.
        cdr3_query_start = record["v_sequence_end"] + (
            cdr3_ref_start - record["v_germline_end"]
        )

    # CDR3 end
    cdr3_ref_end = database.j_cdr3_end(record["j_call"], record["locus"])
    if cdr3_ref_end is None:
        return

    cdr3_query_end = query_position(record, "j", reference_position=cdr3_ref_end)
    if cdr3_query_end is None:
        return
    cdr3_nt = record["sequence"][cdr3_query_start:cdr3_query_end]

    record["cdr3_start"] = cdr3_query_start + 1
    record["cdr3_end"] = cdr3_query_end
    record["cdr3"] = cdr3_nt
    record["cdr3_aa"] = nt_to_aa(cdr3_nt)


def set_fwr4_columns(record, database):
    j_call = record["j_call"]
    if not j_call or record["locus"] not in ALLOWED_LOCI:
        return

    cdr3_ref_end = database.j_cdr3_end(record["j_call"], record["locus"])
    cdr3_query_end = record["cdr3_end"]
    if cdr3_ref_end is None or not cdr3_query_end:
        return

    fwr4_nt = record["sequence"][cdr3_query_end : record["j_sequence_end"]]

    # This overwrites some existing columns
    record["fwr4_start"] = cdr3_query_end + 1
    record["fwr4_end"] = record["j_sequence_end"]
    record["fwr4"] = fwr4_nt
    record["fwr4_aa"] = nt_to_aa(fwr4_nt)


def query_position(record, gene: str, reference_position: int):
    """
    Given a position on the reference, return the same position but relative to
    the full query sequence.
    """
    ref_pos = record[gene + "_germline_start"] - 1
    query_pos = record[gene + "_sequence_start"] - 1

    # Iterate over alignment columns
    if ref_pos == reference_position:
        return query_pos

    sequence_alignment = record[gene + "_sequence_alignment"]
    germline_alignment = record[gene + "_germline_alignment"]

    for ref_c, query_c in zip(germline_alignment, sequence_alignment):
        if ref_c != "-":
            ref_pos += 1
        if query_c != "-":
            query_pos += 1
        if ref_pos == reference_position:
            return query_pos
    return None

def parse_record(record):
    for col, typ in COLUMN_TYPES.items():
        record[col] = typ(record[col]) if record[col] else None
    # Then remove leading %
    for col in "v_call", "d_call", "j_call":
        record[col] = record[col].lstrip("%")
    return record

def dict_to_tsv(row_dict, fieldnames):
    """
    Turn a row dictionary into a tab-separated string in the same column order as 'fieldnames'.
    If a column is missing, use empty string.
    """
    return "\t".join(str(row_dict.get(f, "") or "") for f in fieldnames)

################################################################################

def main():
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    parser = ArgumentParser(description=__doc__)
    arg = parser.add_argument
    arg(
        "--threads",
        "-t",
        "-j",
        type=int,
        default=1,
        help="Number of threads. Default: 1. Use 0 for no. of available CPUs.",
    )
    arg(
        "--penalty",
        type=int,
        choices=(-1, -2, -3, -4),
        default=None,
        help="BLAST mismatch penalty (default: -1)",
    )
    arg(
        "--species",
        default="human",
        help="Tell IgBLAST which species to use. Note that this setting does "
        "not seem to have any effect since we provide our own database to "
        "IgBLAST. Default: Use IgBLAST’s default",
    )
    arg(
        "--sequence-type",
        default="Ig",
        choices=("Ig", "TCR"),
        help="Sequence type. Default: %(default)s",
    )
    arg("--limit", type=int, metavar="N", help="Limit processing to first N records")

    arg("--database", help="Database directory with V.fasta, D.fasta, J.fasta.")
    arg("--fasta", help="File with original reads")

    args = parser.parse_args()
    run_igblastwrap(**vars(args))


def run_igblastwrap(
    threads: int,
    penalty: Optional[int],
    species: str,
    sequence_type: str,
    limit: Optional[int],
    database: str,
    fasta: str,
):
    if threads == 0:
        threads = available_cpu_count()
    start_time = time.time()
    last_status_update = 0

    # Main function edited to perform inline augmentation of IgBLAST output
    with ExitStack() as stack:
        sequences = stack.enter_context(dnaio.open(fasta))
        sequences = islice(sequences, 0, limit)

        # 1) Create the Database object here, so we can pass it to set_cdr3_columns:
        db = Database(database, sequence_type)

        csv.register_dialect(
                "airr",
                delimiter="\t",
                lineterminator="\n",
                strict=True,
        )
        n = 0  # number of data lines printed so far
        chunk_count = 0
        # 2) The IgBLAST results come in "chunks", each chunk is a chunk of text from IgBLAST
        for chunk_output in igblast_parallel_chunked(
            database,
            sequences,
            sequence_type=sequence_type,
            species=species,
            threads=threads,
            penalty=penalty,
        ):
            chunk_count += 1

            # chunk_output is the raw multi-line AIRR table from IgBLAST for this chunk
            lines = chunk_output.splitlines()
            if not lines:
                continue

            # The first line in each chunk is typically the header if it's
            # the standard AIRR format. But we only want to print the header once.
            if chunk_count == 1:
                # the very first chunk: print the header line
                header_line = lines[0]
                print(header_line)
                data_lines = lines
            else:
                # subsequent chunks also have a header as lines[0], skip it
                data_lines = lines[1:]

            # Now parse the data lines into a CSV DictReader for augmentation
            if data_lines:
                text_data = "\n".join(data_lines)
                f = StringIO(text_data)
                reader = csv.DictReader(f, dialect="airr")
                fieldnames = reader.fieldnames

                # For each row in this chunk:
                for row_dict in reader:
                    # parse ints/floats if needed
                    row_dict = parse_record(row_dict)

                    # Now augment: set CDR3 and FR4
                    set_cdr3_columns(row_dict, db)
                    set_fwr4_columns(row_dict, db)

                    # Re-print the augmented row as a tab-separated line
                    print(dict_to_tsv(row_dict, fieldnames))
                    n += 1

            # Just for logging progress
            if n > 0 and n % 1000 == 0:
                elapsed = time.time() - start_time
                if elapsed >= last_status_update + 60:
                    logger.info(
                        "Processed {:10,d} sequences at {:.3f} ms/sequence".format(
                            n, elapsed / n * 1e3
                        )
                    )
                    last_status_update = elapsed

    # Done reading all chunks
    elapsed = time.time() - start_time
    if n > 0:
        logger.info(
            "Processed {:10,d} sequences at {:.1f} ms/sequence".format(
                n, elapsed / n * 1e3
            )
        )
    else:
        logger.info("Processed 0 sequences in {:.1f} s".format(elapsed))


if __name__ == "__main__":
    main()