"""Perform a grep like search against specific columns in a flat file.
"""
# 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 re
import sys
# import pprint as pp


_SCRIPT_NAME = 'grep-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`)
"""


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
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:
        get_matching_rows(
            args.infile, args.search_term, columns=args.columns,
            outfile=args.outfile, delimiter=args.delimiter,
            verbose=prog_verbose, encoding=args.encoding,
            comment=args.comment, tmpdir=args.tmpdir,
            no_match=args.missing, ignore_case=args.ignore_case
        )
        log.log_end(logger)
    except (BrokenPipeError, IOError, 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(
        'search_term', type=str, nargs="?",
        help="A term to search for, search terms are treated as regexps unless"
        " --no-regexp is active. If --file is used then this shuld not be "
        "given"
    )
    parser.add_argument(
        '-i', '--infile', type=str,
        help="An input file to search, if not provied then input is assumed "
        "to be from STDINinput is 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(
        '-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 search terms (regexps) to find in the file. If"
        " not provided then a search term should be given as a positional "
        "argument"
    )
    parser.add_argument(
        '-C', '--columns', type=str, nargs="+",
        help="The column names to search within. If none are given then all"
        "columns are searched"
    )
    parser.add_argument(
        '--missing', action="store_true",
        help="Include rows that have no annotation in the output file."
    )
    parser.add_argument(
        '--no-regexp', action="store_true",
        help="Matched based on identity within the column not a regexp, i.e. "
        "using Pythons 'in' operator."
    )
    parser.add_argument(
        '-I', '--ignore-case', action="store_true",
        help="Perform case insensitive matching."
    )
    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()

    # check the search_term is valid
    if (args.search_term is None and args.file is None) or \
       (args.search_term is not None and args.file is not None):
        raise ValueError(
            "you must provide either a search term or a search file"
        )

    terms = []
    if args.file is not None:
        with open(args.file) as infile:
            terms = [i.rstrip('\r').rstrip('\n') for i in infile]
    else:
        terms = [args.search_term]

    if args.no_regexp is False and args.ignore_case is False:
        terms = [re.compile(i) for i in terms]
    if args.no_regexp is False and args.ignore_case is True:
        terms = [re.compile(i, re.IGNORECASE) for i in terms]

    # Put the parsed terms back into the arguments
    args.search_term = terms

    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def get_matching_rows(infile, search_terms, columns=None, outfile=None,
                      verbose=False, tmpdir=None, no_match=False,
                      delimiter=_DEFAULT_DELIMITER, ignore_case=False,
                      encoding=_DEFAULT_ENCODING,
                      comment=_DEFAULT_COMMENT):
    """The main high level API function that takes an input file and generates
    and output file of matching entries.

    Parameters
    ----------
    infile : `str`
        The input file with the ``gene_id`` column.
    search_terms : `list` or (`str` or `re.Pattern`)
        The terms to search for in ``columns``. If a term is a string then the
        Python ``in`` operator is used to perform the search, if a `re.Pattern`
        object then `re.search` is used.
    columns : `list`of `str` or `NoneType`, optional, default: `NoneType`
        Specific column names to search in. If `NoneType` then all the columns
        are used.
    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.
    no_match : `bool`, optional, default: `False`
        Report all the entries that have no matches in addition to those that
        do match.
    delimiter : `str`, optional, default `\t`
        The delimiter of the input (and output) files.
    ignore_case : `bool`, optional, default: `False`
        When search terms are `str`, then perform case insensitive matching.
        Note that this only impacts search strings and not regular expressions.
        For pre-compiled regular expressions, compile them with `re.IGNORECASE`
        if you want case insensitive matching.
    encoding : `str`, optional, default: `UTF-8`
        The default encoding 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.
    """
    write_method = utils.get_open_method(outfile, allow_none=True)

    # This will open an output file if outfile is not None or - otherwise 'out'
    # will be STDOUT
    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_matching_rows(infile, search_terms, columns=columns,
                                    no_match=no_match, yield_header=True,
                                    delimiter=delimiter,
                                    encoding=encoding,
                                    comment=comment),
                desc="[progress] fetching matching rows: ",
                unit=" returned rows",
                disable=not verbose
        ):
            writer.writerow(row)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def yield_matching_rows(infile, search_terms, columns=None, no_match=False,
                        yield_header=True, ignore_case=False,
                        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.
    search_terms : `list` or (`str` or `re.Pattern`)
        The terms to search for in ``columns``. If a term is a string then the
        Python ``in`` operator is used to perform the search, if a `re.Pattern`
        object then `re.search` is used.
    columns : `list`of `str` or `NoneType`, optional, default: `NoneType`
        Specific column names to search in. If `NoneType` then all the columns
        are used.
    no_match : `bool`, optional, default: `False`
        Report all the entries that have no matches in addition to those that
        do match.
    yield_header : `bool`, optional, default: `True`
        When processing the input file the first row yielded should be a header
        row.
    ignore_case : `bool`, optional, default: `False`
        When search terms are `str`, then perform case insensitive matching.
        Note that this only impacts search strings and not regular expressions.
        For pre-compiled regular expressions, compile them with `re.IGNORECASE`
        if you want case insensitive matching.
    delimiter : `str`, optional, default `\t`
        The delimiter of the input (and output) files.
    encoding : `str`, optional, default: `#`
        The default encoding 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.

    Yields
    ------
    row : `list`
        A row of data. If ``yield_header=True`` then the first row will be the
        header row.

    Raises
    ------
    ValueError
        If any specified columns can not be found in the header.
    TypeError
        If any search terms are not `str` or `re.Pattern` objects.
    """
    # Process the search terms and assign them to their correct search function
    # based on if they are text or a pre-compiled regexp
    searches = []
    for i in search_terms:
        if isinstance(i, re.Pattern):
            searches.append((i, _search_regexp))
        elif isinstance(i, str) and ignore_case is False:
            searches.append((i, _search_in))
        elif isinstance(i, str) and ignore_case is True:
            searches.append((i, _search_in_ignore_case))
        else:
            raise TypeError("unknown search type: {0}".format(type(i)))

    read_method = utils.get_open_method(infile, allow_none=True)

    with stdopen.open(infile, method=read_method, encoding=encoding) as incsv:
        reader = csv.reader(incsv, delimiter=delimiter)
        h, nrow = header.get_header_from_reader(reader, comment=comment)
        if columns is None:
            columns = h

        # Now make sure the columns exist
        column_map = [(h.index(i), i) for i in columns]

        if yield_header is True:
            yield h + \
                ['match_text', 'match_column', 'search_term']

        for idx, row in enumerate(reader, 1):
            has_match = False
            for term, func in searches:
                for col_idx, col_name in column_map:
                    match_start, match_text, search = func(term, row[col_idx])
                    if match_text is not None:
                        has_match = True
                        yield row + [match_text, col_name, search]
                        break
                if has_match is True:
                    break
            if has_match is False and no_match is True:
                yield row + [None, None, None]


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _search_regexp(search_term, value):
    """Perform a regular expression search on a string value.

    Parameters
    ----------
    search_term : `re.Pattern`
        A pre-compiled regular expression to query against the ``value``
    value : `str` or `NoneType`
        A value to determine if it matches against ``search_term``.
        ``NoneType`` values do not error out but return no match.

    Returns
    -------
    search_start : `int`
        The 0-based start position of any string match. If there is no match
        then this will be -1.
    match_str : `str` or `NoneType`
        The string that matches within value. This will be ``NoneType`` if
        there was no match
    search_term : `str` or `NoneType`
        The regular expression pattern that produced the match. This will be
        ``NoneType`` if there was no match.
    """
    try:
        match = search_term.search(value)
    except TypeError:
        # probably NoneType
        return -1, None, None

    if match:
        return match.start(), value[match.start():match.end()], \
            search_term.pattern
    return -1, None, None


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _search_in(search_term, value):
    """Perform a case sensitive string IN search on a string value.

    Parameters
    ----------
    search_term : `str`
        A string to query against the ``value``
    value : `str` or `NoneType`
        A value to determine if it matches against ``search_term``.
        ``NoneType`` values do not error out but return no match.

    Returns
    -------
    search_start : `int`
        The 0-based start position of any string match. If there is no match
        then this will be -1.
    match_str : `str` or `NoneType`
        The string that matches within value. This will be ``NoneType`` if
        there was no match
    search_term : `str` or `NoneType`
        The regular expression pattern that produced the match. This will be
        ``NoneType`` if there was no match.
    """
    try:
        # It is faster to use if than try, except when you expect most to fail
        if search_term in value:
            return value.index(search_term), search_term, search_term
    except TypeError:
        # Probably NoneType
        pass

    return -1, None, None


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _search_in_ignore_case(search_term, value):
    """Perform a case in-sensitive string IN search on a string value.

    Parameters
    ----------
    search_term : `str`
        A string to query against the ``value``
    value : `str` or `NoneType`
        A value to determine if it matches against ``search_term``.
        ``NoneType`` values do not error out but return no match.

    Returns
    -------
    search_start : `int`
        The 0-based start position of any string match. If there is no match
        then this will be -1.
    match_str : `str` or `NoneType`
        The string that matches within value. This will be ``NoneType`` if
        there was no match
    search_term : `str` or `NoneType`
        The regular expression pattern that produced the match. This will be
        ``NoneType`` if there was no match.
    """
    try:
        search_term = search_term.lower()
        value = value.lower()

        # It is faster to use if than try, except when you expect most to fail
        if search_term in value:
            return value.index(search_term), search_term, search_term
    except (AttributeError, TypeError):
        # Probably NoneType
        pass

    return -1, None, None


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