"""Perform a column aware concatenation of files. All files must have the same
 delimiter/encoding and have a header.

This performs like regular cat (except slower) but performs a union of columns
in all the files. So it will align common columns between input file and add
blank columns in files that have missing columns.
"""

# 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 stdopen
import argparse
import csv
import os
import sys
# import pprint as pp


_SCRIPT_NAME = 'cat-column'
"""The name of the script that is implemented in this module (`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`)
"""
_STDIN_ID = '-'
"""Indicator that an input is from STDIN (`str`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script.
    """
    # 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:
        column_cat(
            args.infiles, outfile=args.outfile, delimiter=args.delimiter,
            verbose=prog_verbose, encoding=args.encoding,
            comment=args.comment, tmpdir=args.tmpdir
        )
        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(
        'infiles', type=str, nargs="*",
        help="Input files to concatenate, if input is from STDIN use a - "
        "(dash). Files listed here are also combined with any files given "
        "in --file"
    )
    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(
        '-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(
        '-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(
        '-f', '--file', type=str,
        help="A file containing input file paths. Must be absolute file paths"
        ", one per line"
    )
    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()

    if args.file is not None:
        with open(args.file) as infile:
            for i in infile:
                i = i.strip()
                if not i.startswith('#'):
                    args.infiles.append(i)

    if len(args.infiles) == 0:
        raise ValueError(
            "no input files to process, please use - (dash) for input from"
            " STDIN"
        )

    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def column_cat(infiles, outfile=None, verbose=False, tmpdir=None,
               delimiter=_DEFAULT_DELIMITER,
               encoding=_DEFAULT_ENCODING,
               comment=_DEFAULT_COMMENT):
    """The main high level API function that takes an input file and generates
    and output file of concatenated entries.

    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``.
    verbose : `bool`, optional, default: `False`
        Report progress.
    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.
    delimiter : `str`, optional, default `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        The symbol for a comment row that will be skipped if they occur before
        the header.
    """
    write_method = utils.get_open_method(outfile, allow_none=True)

    with stdopen.open(outfile, mode='wt', method=write_method, use_tmp=True,
                      tmpdir=tmpdir, encoding=encoding) as out:
        writer = csv.writer(out, delimiter=delimiter,
                            lineterminator=os.linesep)
        for row in tqdm(
                yield_column_cat(infiles, yield_header=True,
                                 delimiter=delimiter,
                                 encoding=encoding,
                                 comment=comment),
                desc="[progress] concatenating: ",
                unit=" rows",
                disable=not verbose
        ):
            writer.writerow(row)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_column_cat(infiles, yield_header=True,
                     delimiter=_DEFAULT_DELIMITER,
                     encoding=_DEFAULT_ENCODING,
                     comment=_DEFAULT_COMMENT):
    """Yield rows of matching entries.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    yield_header : `bool`, optional, default: `True`
        When processing the input file the first row yielded should be a header
        row.
    delimiter : `str`, optional, default `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        The symbol for a comment row that will be skipped if they occur before
        the header.

    Yields
    ------
    row : `list`
        A row of data. If ``yield_header=True`` then the first row will be the
        header row.
    """
    kwargs = dict(
        delimiter=delimiter,
        encoding=encoding,
        comment=comment
    )

    column_pool, stdin_input, stdin_header = get_column_pool(infiles, **kwargs)

    if yield_header is True:
        yield column_pool

    for i in infiles:
        fobj = None
        reader = stdin_input
        h = stdin_header
        if i != _STDIN_ID:
            read_method = utils.get_open_method(i)
            fobj = read_method(i, mode='rt', encoding=encoding)
            reader = csv.reader(fobj, delimiter=delimiter)
            h, nrow = header.get_header_from_reader(
                reader, comment=comment
            )
        try:
            seen = set()
            out_idx = []
            for c in h:
                # First deal with any duplicated column names in the same
                # header, we will append a numerical suffix on them
                idx = 1
                while c in seen:
                    c = "{0}_{1}".format(c, idx)
                    idx += 1

                out_idx.append(column_pool.index(c))
                seen.add(c)

            for row in reader:
                outrow = [None] * len(column_pool)
                for idx, c in enumerate(row):
                    outrow[out_idx[idx]] = c
                yield outrow
        finally:
            if i != _STDIN_ID:
                fobj.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_column_pool(infiles, delimiter=_DEFAULT_DELIMITER,
                    encoding=_DEFAULT_ENCODING, comment=_DEFAULT_COMMENT):
    """Get all the unique columns present in all of the input files.

    Parameters
    ----------
    infiles : `list` of `str` or `NoneType`
        The input file paths, if one is '-' then input is from STDIN
    delimiter : `str`, optional, default `con.DEFAULT_DELIMITER`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `con.DEFAULT_ENCODING`
        The default encoding of the input (and output) files.
    comment : `str`, optional, default: `con.DEFAULT_COMMENT`
        The symbol for a comment row that will be skipped if they occur before
        the header.

    Returns
    -------
    column_pool : `list` of `str`
        A unique set of columns in all files.
    stdin_input : `csv.Reader` or `NoneType`
        If any of the inputs are from STDIN then this will be a reader to
        stdin. If NoneType, then there are no inputs from STDIN.
    stdin_header : `list` of `str`
        If any inputs are from STDIN, this will be the header for them. If
        there are none then this will be an empty list.
    """
    stdin_input = None
    stdin_header = []
    column_pool = []

    for i in infiles:
        h = None
        if i != _STDIN_ID:
            read_method = utils.get_open_method(i)
            with read_method(i, mode='rt', encoding=encoding) as infile:
                reader = csv.reader(infile, delimiter=delimiter)
                h, nrow = header.get_header_from_reader(
                    reader, comment=comment
                )
        elif i == _STDIN_ID and stdin_input is None:
            stdin_input = csv.reader(sys.stdin, delimiter=delimiter)
            h, nrow = header.get_header_from_reader(
                stdin_input, comment=comment
            )
            stdin_header = h
        elif i == _STDIN_ID and stdin_input is None:
            raise TypeError("only a single STDIN - allowed")

        seen = set()
        for c in h:
            # First deal with any duplicated column names in the same
            # header, we will append a numerical suffix on them
            idx = 1
            while c in seen:
                c = "{0}_{1}".format(c, idx)
                idx += 1

            if c not in column_pool:
                column_pool.append(c)
            seen.add(c)

    return column_pool, stdin_input, stdin_header


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