from hpc_tools import (
    __version__,
    __name__ as PKG_NAME,
    file_helper,
    common
)
# import pprint as pp
import argparse
import sys
import csv
import os
import re
import warnings

# Make sure CSV can handle really long lines
csv.field_size_limit(sys.maxsize)

# The delimiter for the job array files
DELIMITER = "\t"

# The header for job array files
HEADER = ['rowidx', 'infile', 'outfile']


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def infile_readable(infiles):
    """
    Make sure each file in the list of input files is readable. This is done by
    attempting to open and close the file.

    Parameters
    ----------
    infiles : list of str
        A list of input file names to open and close

    Raises
    ------
    FileNotFoundError
        If an input file does not exist
    PermissionError
        If you do not have read permissions on the input file
    """
    for i in infiles:
        # This will generate a FileNotFound error is not readable
        open(i).close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def generate_output_files(infiles, outdir=None, prefix=None, suffix=None,
                          strip_extension=False, levels=0, regexp=None,
                          outfile_errors='error'):
    """
    Using the input files and the keyword arguments generate an output
    file. If the output file exists then an error is raised, also, if the
    output file can't be written in the output directory then an error is
    raised.

    Parameters
    ----------
    infiles : list of str
        A list of input file names to check and generate output file names
        from
    outdir : str or NoneType, optional
        A specific output directory to write all the output files, if it does
        not exist then it will be created. If `NoneType`, then the output
        directory will be gathered from the directory location of the input
        file (default=`NoneType`)
    prefix : str or NoneType, optional
        A prefix to add onto the output file name (default=`NoneType`)
    suffix : str or NoneType, optional
        A suffix to add on to the the end of the output file name, before
        the file extension (default=`NoneType`)
    strip_extension : bool, optional
        If `True` then the input file name extension is removed when generating
        the output file name (default=`False`)
    levels : int, optional
        Add number of directory `levels` up from the input file to below the
        output file directory. Only active if `outdir` is used (default=0)
    regexp : :obj:`re.Pattern`, optional
        A compiled regular expression that is applied over the whole input file
        path. This uses named capture groups to extract the parts that should
        be used in the output file name. the capture groups should be named
        in uppercase alphabet characters and the order in the alphabet is the
        order they appear in the output file name. i.e. `r'(?P<A>first_bit)'`.
        (default=`NoneType`)
    outfile_errors : str, optional
        What to do in the case that the output file names are not unique. The
        keyword 'error' will raise an exception, 'ignore' will ignore the
        non-unique output file names but will still issue a warning. This maybe
        intended for example if the step size if the job array is not 1. The
        'flatten' keyword, will make a pipe `|` concatinated string of the
        input files to make the output file names unique. i.e. we group on the
        outfiles. In addition, a warning is issued if `flatten` is used and
        each grouping of input files does not have the same number of files
        within it. (default=`'error'`)

    Returns
    -------
    all_files : list of tuple
        Each tuple has a `rowidx` (`int`) at element `[0]`, the full path to
        the input file name (`str`) at element `[1]`, and the full path to the
        output file name at element `[2]`.

    Raises
    ------
    FileExistsError
        If any of the output file names generated over all the input files are
        not unique, so this makes sure that output files will not be
        overwritten by other tasks when the job is running
    PermissionError
        If the output directory is not writable, or an output directory can't
        be created
    IsADirectoryError
        If the output file name exists and is a directory
    FileNotFoundError
        If the outdir can't be created because the path is wrong
    """

    try:
        # If the outdir is not None then create it and make sure we ignore any
        # errors
        outdir = os.path.abspath(os.path.expanduser(outdir))
        os.mkdir(outdir)
    except (TypeError, FileExistsError):
        # outdir is None or already exists
        pass

    # Make sure the prefix and suffix are represented as empty strings if they
    # are None, this makes the string formatting code easier
    if prefix is None:
        prefix = ""

    if suffix is None:
        suffix = ""

    # Will hold the tuples for returning
    all_files = []

    # Will hold the generated input files so I can quickly test for uniqueness
    # later
    outfiles = set()

    # Holds the number of regexp matches when a regexp is applied to the full
    # path of the input file. We are expecting that the regexp should produce
    # the same number of matches for each input file
    regexp_matches = set()
    for idx, infile in enumerate(infiles):
        # Process the input file into it's component parts. Note that this
        # treats .txt.gz, tab.gz as a single extension, this normally would
        # not happen with os.path.splitext
        indir, infile_name, extension = file_helper.get_file_components(
            infile)

        # Generate the output directory name. This will either be the indir
        # or outdir+levels of the indir
        cur_outdir = file_helper.get_output_directory(indir, outdir,
                                                      levels=levels)

        try:
            # Now create the output directory, this makes them recursively
            # We do not want to error out if they exist as there maybe multiple
            # runs of the job array, particularly if a regexp is used
            os.makedirs(cur_outdir)
        except FileExistsError:
            # The output directory already exists
            pass

        # Initialise the outfile name. So infile_name is the basename of the
        # input file without any extensions
        outfile_name = infile_name

        # If we are applying a regexp to the infile to generate an output file
        # name
        if regexp is not None:
            # ncomponents is the number of groups that have been captured by
            # the regular expression. We  use this to test for the same number
            # of regexp groups being captured for each input file
            outfile_name, ncomponents = file_helper.outfile_from_regexp(
                infile, regexp)
            regexp_matches.add(ncomponents)

        # Add any prefix and suffix to the output file name
        outfile_name = "{0}{1}{2}".format(prefix, outfile_name, suffix)

        # If we are not striping the extension then we will add it back onto
        # the the end of the output file
        if strip_extension is False:
            outfile_name = "{0}{1}".format(outfile_name, extension)

        # Generate the full output file path
        outfile = os.path.join(cur_outdir, outfile_name)

        try:
            # Now we have an output file name we test to see if it exists or
            # not
            open(outfile).close()
            raise FileExistsError("output file '{0}' already exists, is it "
                                  "the same as the input file "
                                  "path?".format(outfile))
        except FileNotFoundError:
            # file does not exist so all good so far
            pass

        # Now test if we can actually write the file,this may generate a
        # PermissionError or a FileNotFoundError
        open(outfile, 'w').close()
        os.unlink(outfile)

        # If we get here then everything checks out so far, now we need to
        # store the output file name along with the input file name. Note that
        # this is stored as a list as if 'flatten' is used then a list of
        # output files will be used (see below)
        all_files.append((idx+1, [infile], outfile))
        outfiles.add(outfile)

    # The final check is that each of the output files is unique, this is to
    # make sure that the output file from one job will not be overwritten by
    # the output file from another job
    if len(outfiles) != len(all_files):
        error_msg = "not all output files are unique, this will cause a" +\
            " conflict between tasks (unless the step size > 1, then it might"\
            " not)"

        if outfile_errors == 'error':
            raise FileExistsError(error_msg)
        elif outfile_errors == 'flatten':
            # make the replicated output files unique and use a list of
            # input files
            flat_files = {}

            # Loop through each tuple
            for idx, i, o in all_files:
                try:
                    # The outfile is the key and the list of matching input
                    # files is the value that is being extended with another
                    # input file
                    flat_files[o].extend(i)
                except KeyError:
                    # Initialise
                    flat_files[o] = i

            # Will be re-initialised to hold the final list of input files
            # that will be returned
            all_files = []
            # Store the numbers of input files for each unique output file.
            # This is so we can warn the user if there are different sized
            # groups of input files as this may be an error
            grp_counts = set()
            idx = 1
            for k, v in flat_files.items():
                grp_counts.add(len(v))
                all_files.append((idx, v, k))
                idx += 1

            # If we have different size groups, issue the warning
            if len(grp_counts) > 1:
                warnings.warn("not every infile grouping has the same number"
                              " of items: {0}".format(",".join(grp_counts)))
        elif outfile_errors == 'ignore':
            # if we are ignoring the replicated input files then we will issue
            # a warning anyway for good measure
            warnings.warn(error_msg)

    return all_files


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def write_job_array_file(all_files, header=HEADER, outfile=None):
    """
    Write the job array file to disk or STDOUT. This does not do anything
    special and will just loop through all_files writing the contents using
    `csv.writer`.

    Parameters
    ----------
    all_files : list of tuple
        Each tuple has a `rowidx` (`int`) at element `[0]`, the full path to
        the input file name (`str`) at element `[1]`, and the full path to the
        output file name at element `[2]`. However, there is no actual
        requirement for just this, as long, each tuple in `all_files` is the
        same length as `header` and element [1] is a list (as it will
        be joined with a pipe `|`). The other elements in the tuple could be
        anything.
    header : list of `str`, optional
        The header for the job array file, it defaults to the module level
        HEADER variable
    outfile : str or NoneType, optional
        The output file name for the job array file. If `NoneType` then output
        is to `STDOUT` (default=`NoneType`)

    Raises
    ------
    IndexError
        If any of the elements of `all_files` are not the same length as the
        `header`
    """
    if outfile is None:
        outfile = sys.stdout
    else:
        # We make sure to apply expanduser to copy with ~/
        outfile = os.path.expanduser(outfile)
        outfile = open(outfile, 'w')

    with outfile as csvout:
        writer = csv.writer(csvout, delimiter=DELIMITER,
                            lineterminator=os.linesep)
        writer.writerow(header)
        for idx, i in enumerate(all_files):
            i = list(i)
            i[1] = '|'.join(i[1])

            if len(i) != len(header):
                raise IndexError("line '{0}' does not equal the header "
                                 "length '{1}'".format(idx, len(header)))
            writer.writerow(i)

    outfile.close()


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def process_files(infiles, indir=None, inglob=None, outdir=None, prefix=None,
                  suffix=None, strip_extension=False, levels=0, regexp=None,
                  outfile_errors='error'):
    """
    Process the input files and the output files and add an idx parameter

    Parameters
    ----------
    infiles : list of str
        A list of input files, this will be merged with indir contents
    indir : str or NoneType, optional
        An input directory in addition to any `infiles` (default=`NoneType`)
    inglob : str or NoneType,  optional
        Used in conjunction with indir. If `NoneType`, then all files in the
        root of `indir` are gathered, otherwsie a glob of `indir` and `inglob`
        is used to gather files (default=`NoneType`)
    outdir : str or NoneType, optional
        A specific output directory to write all the output files, if it does
        not exist then it will be created. If `NoneType`, then the output
        directory will be gathered from the directory location of the input
        file (default=`NoneType`)
    prefix : str or NoneType, optional
        A prefix to add onto the output file name (default=`NoneType`)
    suffix : str or NoneType, optional
        A suffix to add on to the the end of the output file name, before
        the file extension (default=`NoneType`)
    strip_extension : bool, optional
        If `True` then the input file name extension is removed when generating
        the output file name (default=`False`)
    levels : int, optional
        Add number of directory `levels` up from the input file to below the
        output file directory. Only active if `outdir` is used (default=0)
    regexp : :obj:`re.Pattern`, optional
        A compiled regular expression that is applied over the whole input file
        path. This uses named capture groups to extract the parts that should
        be used in the output file name. the capture groups should be named
        in uppercase alphabet characters and the order in the alphabet is the
        order they appear in the output file name. i.e. `r'(?P<A>first_bit)'`.
        (default=`NoneType`)
    outfile_errors : str, optional
        What to do in the case that the output file names are not unique. The
        keyword 'error' will raise an exception, 'ignore' will ignore the
        non-unique output file names but will still issue a warning. This maybe
        intended for example if the step size if the job array is not 1. The
        'flatten' keyword, will make a pipe `|` concatenated string of the
        input files to make the output file names unique. i.e. we group on the
        outfiles. In addition, a warning is issued if `flatten` is used and
        each grouping of input files does not have the same number of files
        within it. (default=`'error'`)

    Returns
    -------
    all_files : list of tuple
        Each tuple has a `rowidx` (`int`) at element `[0]`, the full path to
        the input file name (`str`) at element `[1]`, and the full path to the
        output file name at element `[2]`.

    Raises
    ------
    ValueError
        If no files can be gathered
    FileExistsError
        If any of the output file names generated over all the input files are
        not unique, so this makes sure that output files will not be
        overwritten by other tasks when the job is running
    PermissionError
        If the output directory is not writable, or an output directory can't
        be created
    IsADirectoryError
        If the output file name exists and is a directory
    FileNotFoundError
        If the outdir can't be created because the path is wrong or if any of
        the input files do not exist
    """
    infiles = file_helper.gather_input_files(infiles=infiles, indir=indir,
                                             inglob=inglob)

    infile_readable(infiles)

    return generate_output_files(infiles, outdir=outdir,
                                 prefix=prefix,
                                 suffix=suffix,
                                 strip_extension=strip_extension,
                                 levels=levels,
                                 regexp=regexp,
                                 outfile_errors=outfile_errors)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def _init_cmd_args(description=None):
    """
    Use `argparse` to initialise the expected commandline arguments. Note that
    this does not parse the arguments. It is designed so that it returns the
    parser object for the user to use to add additional arguments to. The
    short arguments that are used are '-h', `-f`, `-i`, `-g`, `-o`, `-l`, `-e`,
    `-p`, `-r`, `-s`, `x`, `-v` - so these should not be used in any additional
    arguments added to the parser.

    Parameters
    ----------
    description : str or NoneType, optional
        An alternative program description for `argparse`, if not defined then
        the default one will be used (default=`NoneType`)

    Returns
    -------
    args : :obj:`argparse.ArgumentParser`
        An "unparsed" argparse parser
    """
    parser = None
    if description is None:
        description = ("Generate a skeleton for a job array file with input"
                       " file, output file columns and an index column. All"
                       " input files are checked to make sure they can be "
                       "accessed and it is a fatal error if they can not, "
                       "similarly, each output location is checked to see "
                       "if it is writable. If the output file already exists "
                       "then it is a fatal error")

    parser = argparse.ArgumentParser(description=description)

    parser.add_argument('infiles', type=str, nargs='*',
                        help="The input files to for the job array, if none "
                        "are supplied then --indir should be set and all the"
                        " files in indir are used as input")
    parser.add_argument('-f', '--outfile', type=str,
                        help="the output file, if not provided, output is to "
                        "STDOUT")
    parser.add_argument('-i', '--indir', type=str,
                        help="the input directory, this should be set if there"
                        "are no input files, it can also be used in "
                        "conjunction with input files")
    parser.add_argument('-g', '--inglob', type=str,
                        help="An input directory glob that is used in "
                        "conjunction with the --indir, the default is no "
                        "pattern")
    parser.add_argument('-o', '--outdir', type=str,
                        help="An output directory, if this is not set then "
                        "it is assumed that the output directory is the same"
                        " directory as the directory that each input file "
                        "resides in. If the outdir does not exist then it will"
                        " be created")
    parser.add_argument('-l', '--levels', type=int, default=0,
                        help="active when --outdir is used, this uses number "
                        "of levels directory structure down from infile and is"
                        " appended to the outdir. e.g. "
                        "outdir=/home/bob/my_data , "
                        "input file='/data/to/process/file.txt . If --levels"
                        " = 0 then the outfile will be "
                        "'/home/bob/my_data/file.txt', if --levels 1 is used,"
                        " then '/home/bob/my_data/process/file.txt' and if"
                        "levels 2 is used, then "
                        "'/home/bob/my_data/to/process/file.txt'")
    parser.add_argument('-e', '--outfile_errors', type=str, default="error",
                        choices=["error", "ignore", "flatten"],
                        help="action to take if there are errors with "
                        "non-unique output files. 'flatten', will make the "
                        "output file unique and group all the input files "
                        "into a single column delimited with a pipe '|'")
    parser.add_argument('-p', '--prefix', type=str,
                        help="A prefix to add to each output file, the default"
                        " is ''")
    parser.add_argument('-s', '--suffix', type=str,
                        help="A suffix to add to each output file, the default"
                        " is '_output'")
    parser.add_argument('-r', '--regexp', type=str,
                        help="A regular expression used to generate the "
                        "output file name from the full input file path. The"
                        " regexp will be applied using re.search. All groups"
                        " will be captured, although the groups required in "
                        "the output file name should be labels '1', '2', '3'"
                        " using the named capure groups `?P<1>` etc... . "
                        "Remember to also use non-capturing groups `(?:)` "
                        "where necessary.")
    parser.add_argument('-x', '--strip_ext', action='store_true',
                        help="shall the extension be stripped from the output"
                        "file name")
    parser.add_argument('-v', '--verbose', action='store_true',
                        help='shall I print progress to STDERR?')
    return parser


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def parse_cmd_line_args(parser):
    """
    Use argparse to parse the commandline arguments into a Namespace

    Parameters
    ----------
    args : :obj:`argparse.ArgumentParser`
        An "unparsed" argparse parser to parse

    Returns
    -------
    args : :obj:`Namespace`
        The arguments from parsing the cmd line args with `parser.parse_args()`
    """
    args = parser.parse_args()
    return args


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def set_regexp(regexp):
    """
    compile the regexp that has been passed on the command line

    Parameters
    ----------
    regexp : str or NoneType
        A regular expression to compile, if `NoneType`, then the it is just
        returned and no error is generated

    Returns
    -------
    compiled_regexp : :obj:`re.Pattern`
        A compiled regexp, note that no flags are added, this is a bog
        standard compile
    """
    if regexp is not None:
        regexp = r'{0}'.format(regexp)
        regexp = re.compile(regexp)

    return regexp


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def process_files_from_args(args):
    """
    Process files from arguments given on the command line. This is a
    convenience function that just calls `process_files`. The argument
    `Namespace` needs to contain, `infiles`, `indir`, `inglob`, `outdir`,
    `prefix`, `suffix`, `strip_ext`, `levels`, `regexp`, outfile_errors`.

    Parameters
    ----------
    args : :obj:`Namespace`
        The arguments from parsing the cmd line args with `parser.parse_args()`

    Returns
    -------
    all_files : list of tuple
        Each tuple has a `rowidx` (`int`) at element `[0]`, the full path to
        the input file name (`str`) at element `[1]`, and the full path to the
        output file name at element `[2]`.
    """
    args.regexp = set_regexp(args.regexp)

    return process_files(args.infiles, indir=args.indir,
                         inglob=args.inglob, outdir=args.outdir,
                         prefix=args.prefix, suffix=args.suffix,
                         strip_extension=args.strip_ext,
                         levels=args.levels, regexp=args.regexp,
                         outfile_errors=args.outfile_errors)


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """The main entry point for the script
    """
    parser = _init_cmd_args()
    args = parse_cmd_line_args(parser)
    msg = common.Msg(verbose=args.verbose, file=sys.stderr)
    msg.msg('=== make-jaf ({0} v{1}) ==='.format(
        PKG_NAME, __version__
    ))
    msg.prefix = '[info]'
    msg.msg_atts(args)

    try:
        all_files = process_files_from_args(args)
    except ValueError as e:
        print('[fatal] problem with input files: {0}'.format(str(e)),
              file=sys.stderr)
        sys.exit(1)
    except FileNotFoundError as e:
        print("[fatal] problem finding file or directory: '{0}'".format(
            str(e)), file=sys.stderr)
        sys.exit(1)
    except (FileExistsError, PermissionError, IsADirectoryError) as e:
        print("[fatal] problem generating output files: '{0}'".format(str(e)),
              file=sys.stderr)
        sys.exit(1)

    msg.msg('{0} input found'.format(len(all_files)))
    write_job_array_file(all_files, outfile=args.outfile)

    msg.msg("*** END ***", prefix="")


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