import argparse
import os
import sys
import tempfile
import operator
import csv


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def sort_and_write(args, writer, base_csv_args, data):
    """
    The arguments to this function are the scrip args, a csvwriter
    (either unicode or not) and the data (list of lists) to be written to a
    file
    """
    # Sort, write to a temp file
    try:
        # Attmempt to sort according to the keys we passed to the script
        # this should work ok apart from hen the column numbers of the
        # keys are not in all the lines, in which case an IndexError is
        # thrown.
        data.sort(key=operator.itemgetter(*args.key))
    except IndexError:
        # Catch the IndexError and give an equaly perplexing error message
        print('[fatal] key indexes longer than the line', file=sys.stderr)
        sys.exit(1)

    # Get a temp file name and open it and pass the hadle to the csv writer
    n, fn = tempfile.mkstemp()
    fh = open(fn, "w")
    w = writer(fh, **base_csv_args)

    # Write the buffer to file
    for i in data:
        w.writerow(i)
    fh.close()
    return fn


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def mergeiter(*iterables, **kwargs):
    """
    Taken from:
    http://stackoverflow.com/questions/14465154

    This is a generator that returns the contents of all the (sorted) iterables
    that have been passed to it in the correct sort order. Takes an optional
    `key` callable to compare values by.
    """
    # Turn a tuple of iterables into a list of iterables
    iterables = [iter(it) for it in iterables]

    # Create a dict data sctructure with the file number as a key and a list
    # of the current line in the file, file number and iterable as the values
    iterables = {i: [next(it), i, it] for i, it in enumerate(iterables)}

    # If we havn't supplied a key then we do a line sort on all the values in
    # the line and not just a specific column
    if 'key' not in kwargs:
        key = operator.itemgetter(0)
    else:
        # If we have supplied a key then it will be a callable (i.e. a function
        # that returns the values we want to sort on)
        key = kwargs['key']

    # This is where the magic happens
    while True:
        # The original function on stackover flow used min instead of sorted
        # I couldn't work out how to get min to use the keys I wanted so
        # I used sorted and selected the files element it gave back in the
        # sorted list. This will not be as efficient as min especially if we
        # have loads of files
        # print(pp.pformat(iterables.values()), file=sys.stderr)
        value, i, it = sorted(iterables.values(), key=key)[0]

        # 'return' the value
        yield value
        try:
            # For the iterable that had the value returned we want to move on
            # to the next line, if there are no more lines then we catch the
            # error and delete the iterable from our dict. If after doing that
            # out dict is empty then we raise which I presume is the stop
            # condition for the generator
            iterables[i][0] = next(it)
        except StopIteration:
            del iterables[i]
            if not iterables:
                break


# ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
def main():
    """
    Main command line entry point
    """
    # Read in the command line arguments
    parser = argparse.ArgumentParser(
        description='Takes tabular STDIN and makes the keys unique and '
        'collapses all other values'
    )
    parser.add_argument('-d', '--delimiter', type=str,
                        default="\t", help="The delimiter of STDIN")
    parser.add_argument('-g', '--grp-del', type=str, default=",",
                        help="The delimiter of the grouped columns")
    parser.add_argument(
        '-q', '--quote',
        type=str,
        default='"',
        help="The Quoting character (to protect dlimiting characters that "
        "are not real delimiters)"
    )
    parser.add_argument(
        '-e', '--encoding',
        type=str,
        default='UTF-8',
        help="The unicode encoding in the file"
    )
    parser.add_argument('-b', '--buffer', type=int, default=100000,
                        help="The number of lines to sub sort in memory")
    parser.add_argument(
        '-k', '--key',
        action="append",
        help="The column(s) number to group on, can be column name if "
        "--header works similar to unix sort -k1 -k4 will unique on "
        "columns 1 and 4. Note this is 1-based"
    )
    parser.add_argument('--header', action="store_true",
                        help="Does the file have a header")
    args = parser.parse_args()

    # Make sure key has a default value of the first column
    if args.key is None:
        args.key = [1]

    # Read in stdinput
    with sys.stdin as csvfile:
        # These are the arguments that we will give to whatever
        # csvreader/writer we use
        base_csv_args = {'delimiter': args.delimiter}

        # Make sure that the quotechar argument is only parsed if the quotechar
        # is enabled otherwise we will get:
        # TypeError: quotechar must be set if quoting enabled errors
        if args.quote != "":
            base_csv_args['quotechar'] = args.quote

        reader = csv.reader(csvfile, **base_csv_args)

        # If we have a file header then grab it
        if args.header:
            header = next(reader)

        # Now get the column numbers of the keys, they may already be numbers
        # 1-based in which case we will make them 0-based but they could also
        # be column names
        for c, k in enumerate(args.key):
            # First we try to convert the key to an integer, if that works
            # then we make it 0-based so we can subset correctly from lists
            # However, there is a chance that the user has given a rubbish key
            # that is still an integer, i.e. 0 or negative value, so if that is
            # the case we will error out
            try:
                args.key[c] = int(k)-1

                if args.key[c] < 0:
                    print('[fatal] key must be a positive integer not: %i' %
                          int(k), file=sys.stderr)
                    sys.exit(1)
            except ValueError:
                # If we end up here then the key must have been a column name
                # (or cocked up input ~ either way it is not an integer). So
                # we look for the key in the header. If it is not in the header
                # we error out. If we don't have a header we also error out
                # but try to be helpful
                try:
                    args.key[c] = header.index(k)
                except ValueError:
                    print('[fatal] can not find key in header: %s' %
                          str(k), file=sys.stderr)
                    sys.exit(1)
                except NameError:
                    print('[fatal] key ({0}) is not an integer and there is no'
                          ' header (forgot to add --header ?)'.format(
                              str(k)), file=sys.stderr)
                    sys.exit(1)

        # Will hold out --buffer number of lines in memory
        stdin_buffer = []

        # Will hold the file name of the temp file when sorted buffer lines
        # have been written
        tmp_files = []

        # Now loop through each of the rows
        for nr, row in enumerate(reader):
            #        print(len(row))
            # If we are have filled the buffer then sort the buffer and write
            # to a temp file before clearing the buffer.
            if nr % args.buffer == 0 and len(stdin_buffer) > 0:
                #            print(nr)
                #            print(args.buffer)
                #            print(len(stdin_buffer))
                #            print(filewriter)
                tmp_files.append(sort_and_write(
                    args, csv.writer, base_csv_args, stdin_buffer))
                stdin_buffer = []
            # Add the row the the buffer
            stdin_buffer.append(row)

        # When we get here we have finished going through the file and we may
        # still have some data in the buffer so do 1 more flush
        # Make sure that we still output if we have something in the buffer
        if len(stdin_buffer) > 0:
            tmp_files.append(sort_and_write(
                args, csv.writer, base_csv_args, stdin_buffer))
            stdin_buffer = []

        # When we get to here we should have created a load of tmp files with
        # the data in the in our desired sort order. These now have to be
        # merged in to the final sorted list and then we can write the grouped
        # data to STDOUT First job is to open all the tmp files
        iterables = set()
        for i in tmp_files:
            reader = csv.reader(open(i), **base_csv_args)
            iterables.add(reader)

        def getkeys(v, key=args.key):
            """
            This is a callable function that I am using to pass to mergeiter,it
            provides the correct columns for the sort command inside mergeiter
            """
            return [v[0][i] for i in key]

        # Holds the previous key values from the last iteration
        prevkey = []

        # Holds all the values relating to the current key, this will be
        # output and reset when the key changes
        forconcat = []

        # Firstly if we have a header then ouput it to STDOUT but we want to
        # use our csv writer for that
        write_out = csv.writer(sys.stdout, **base_csv_args)
        if args.header:
            write_out.writerow(header)

        # Now we can get our data back in the correct sort order for
        # out keys from the tmp files
        for i in mergeiter(*iterables, key=getkeys):
            # Get the values for the key columns this will be
            # compared to the previous key
            key = [i[j] for j in args.key]

            # If we have a previous key and our current
            # key is different then we do not want to group anything
            # else, rather we want to output what we have already grouped
            if len(prevkey) > 0 and prevkey != key:
                # Loop throug all the sets of data for the key and join them
                # this works as I am using sets() and the key columns are
                # only present once
                for c, v in enumerate(forconcat):
                    forconcat[c] = args.grp_del.join(v)

                # Now write the line according to our specs
                write_out.writerow(forconcat)

                # As we are moving onto the next key then empty the contents
                # of out record storage
                forconcat = []

            # Here we are looping through all the columns and adding the data
            # to our record storage sets in the order of the columns
            for c, v in enumerate(i):
                try:
                    forconcat[c].add(v)
                except IndexError:
                    forconcat.append(set([v]))

            # Before we move to the next record make sure the current key is
            # assigned to the previous one
            prevkey = key

        # This makes sure that the last record gets output
        for c, v in enumerate(forconcat):
            forconcat[c] = args.grp_del.join(v)

        # Now write the last line
        write_out.writerow(forconcat)
        forconcat = []

        # Finally remove all the temp files, this needs to me in some sort of
        # destroy block so that it happens even if the script exist early
        # see: http://stackoverflow.com/questions/12686160
        # http://stackoverflow.com/questions/1112343
        # https://docs.python.org/2/library/signal.html
        # http://stackoverflow.com/questions/3347775
        for c, i in enumerate(iterables):
            os.remove(tmp_files[c])
