"""Pad a delimited flat file out to a defined number of columns. Rows which
are short of columns are given an empty value.
"""
# Importing the version number (in __init__.py) and the package name
from pyaddons import (
    __version__,
    __name__ as pkg_name,
    utils,
    log
)
from pyaddons.flat_files import header

from tqdm import tqdm
import argparse
import stdopen
import sys
import os
import csv
# import pprint as pp


_SCRIPT_NAME = 'pad'
"""The name of the script that is output when verbose is True (`str`)
"""
_DESC = __doc__
"""The program description (`str`)
"""
_DEFAULT_DELIMITER = "\t"
"""The default delimiter (`str`)
"""
_DEFAULT_COMMENT = "#"
"""The default comment character (`str`)
"""
_DEFAULT_ENCODING = "UTF-8"
"""The default encoding (`str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script. For API access see:
    ``pyaddons.flat_files.rejig.rejig``
    """
    # Initialise and parse the command line arguments
    parser = _init_cmd_args()
    args = _parse_cmd_args(parser)

    logger = log.init_logger(
        _SCRIPT_NAME, verbose=args.verbose
    )
    log.log_prog_name(logger, pkg_name, __version__)
    log.log_args(logger, args)

    prog_verbose = log.progress_verbose(args.verbose)
    try:
        pad(args.infile, outfile=args.outfile, empty=args.empty,
            tmpdir=args.tmpdir, verbose=prog_verbose, npad=args.npad,
            delimiter=args.delimiter, comment=args.comment,
            skiplines=args.skip_lines, encoding=args.encoding)
        log.log_end(logger)
    except (BrokenPipeError, KeyboardInterrupt):
        # Python flushes standard streams on exit; redirect remaining
        # output to devnull to avoid another BrokenPipeError at shutdown
        devnull = os.open(os.devnull, os.O_WRONLY)

        try:
            os.dup2(devnull, sys.stdout.fileno())
        except ValueError as e:
            if not e.args[0] == "I/O operation on closed file":
                raise

        log.log_interrupt(logger)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args():
    """Initialise the command line arguments and return the parser.

    Returns
    -------
    args : :obj:`argparse.ArgumentParser`
        The argparse parser object with arguments added
    """
    parser = argparse.ArgumentParser(
        description=_DESC
    )
    parser.add_argument(
        '--infile', type=str, nargs='?',
        help="The path to the input file (can be gzip compressed), if not "
        "given then input is assumed to be from STDIN."
    )
    parser.add_argument(
        '-o', '--outfile', type=str,
        help="The path to the output file (if not given will default to"
        " STDOUT), if it has a .gz extension output will be compressed"
    )
    parser.add_argument(
        '-N', '--npad', type=int,
        help="Pad to a certain number of columns. The default is to use the"
        " header length."
    )
    parser.add_argument(
        '-E', '--empty', type=str, default="",
        help="The empty value to use when padding."
    )
    parser.add_argument(
        '-d', '--delimiter', type=str, default="\t",
        help="The delimiter of the input (and output) file"
    )
    parser.add_argument(
        '-e', '--encoding', type=str, default="UTF8",
        help="The encoding of the input file"
    )
    parser.add_argument(
        '-s', '--skip-lines', type=int, default=0,
        help="The skip this many lines before processing the header"
    )
    parser.add_argument(
        '-c', '--comment', type=str, default="#",
        help="Skip lines that start with this string"
    )
    parser.add_argument(
        '-T', '--tmpdir', type=str, default="#",
        help="Alternative tmpdir to use, output is written to tmp before"
        " final output (unless STDOUT)"
    )
    parser.add_argument(
        '-v', '--verbose',  action="count", default=0,
        help="give more output, -vv will turn on progress monitoring"
    )
    parser.add_argument(
        '--keep', action="store_true",
        help="Retain columns not in the transform."
    )

    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _parse_cmd_args(parser):
    """Parse the command line arguments.

    Parameters
    ----------
    parser : :obj:`argparse.ArgumentParser`
        The argparse parser object with arguments added.

    Returns
    -------
    args : :obj:`argparse.Namespace`
        The argparse namespace object containing the arguments
    """
    args = parser.parse_args()
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def pad(infile, outfile=None, tmpdir=None, verbose=False,
        delimiter=_DEFAULT_DELIMITER, comment=_DEFAULT_COMMENT,
        skiplines=0, encoding=_DEFAULT_ENCODING, empty=None, npad=None):
    """Pad a file to ensure a consistent number of columns. The main high level
    API function that takes an input file and generates and output file.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    outfile : `str` or `NoneType`, optional, default: `NoneType`
        The output file. If ``NoneType`` then output will be to ``STDOUT``.
    tmpdir : `str` or `NoneType`, optional, default: `NoneType`
        If writing to a file, a temp file is used and copied over to the final
        location when finished. Use this to specify a different temp location.
        If `NoneType` then the system default temp location is used.
    verbose : `bool`, optional, default: `False`
        Report progress.
    delimiter : `str`, optional, default `\t`
        The delimiter of the input (and output) files.
    comment : `str`, optional, default: `#`
        The symbol for a comment row that will be skipped if they occur before
        the header.
    skiplines : `int`, optional, default: `0`
        A fixed number of rows to skip before looking for the header.
    encoding : `str`, optional, default: `UTF-8`
        The default encoding of the input (and output) files.
    empty : `str`, optional, default: `NoneType`
        The empty value to use for padding.
    npad : `int`, optional, default: `NoneType`
        Pad the columns to this value irrespective of the header length.
    """
    read_method = utils.get_open_method(infile, allow_none=True)
    write_method = utils.get_open_method(outfile, allow_none=True)

    with stdopen.open(infile, method=read_method, encoding=encoding) as infile:
        reader = csv.reader(infile, delimiter=delimiter)
        h, row_no = header.get_header_from_reader(
            reader, comment=comment, skiplines=skiplines
        )
        pad_len = npad or len(h)
        h = pad_header(h, pad_len)
        with stdopen.open(outfile, method=write_method, mode='wt',
                          use_tmp=True, tmpdir=tmpdir,
                          encoding=encoding) as outfile:
            writer = csv.writer(outfile, delimiter=delimiter,
                                lineterminator=os.linesep)
            writer.writerow(h)
            for row in tqdm(reader, desc="[progress] padding: ",
                            unit=" rows", disable=not verbose):
                writer.writerow(pad_row(row, pad_len, empty=empty))


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def pad_header(row, pad_len):
    """Pad the header with more columns if needed (or remove columns).

    Parameters
    ----------
    row : `list` of `str`
        The header row to process.
    pad_len : `int`
        The length we want it padded to.

    Returns
    -------
    padded_header : `list` of `str`
        The padded header row.
    """
    row_len = len(row)
    row = row[:pad_len+1]
    nadd = max(pad_len - len(row), 0)
    for i in range(nadd):
        row.append(f"col{row_len+i+1}")
    return row


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def pad_row(row, pad_len, empty=None):
    """Pad the row with more columns if needed (or remove columns).

    Parameters
    ----------
    row : `list` of `str`
        The row to process.
    pad_len : `int`
        The length we want it padded to.
    empty : `str`, optional, default: `NoneType`
        The value to pad with.

    Returns
    -------
    padded_row : `list` of `str`
        The padded row.
    """
    row = row[:pad_len+1]
    # pp.pprint(row)
    nadd = max(pad_len - len(row), 0)
    if nadd > 0:
        row.extend([empty] * nadd)
    return row


# @@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
if __name__ == '__main__':
    main()
