#!/usr/bin/env python

import argparse
import os
import sys

from csvkit import CSVKitReader, CSVKitWriter
from csvkit.cli import init_common_parser, extract_csv_reader_kwargs, extract_csv_writer_kwargs, install_exception_handler

def main(args):
    """
    Stack up the rows from multiple CSV files, optionally adding a grouping value.
    """
    install_exception_handler(args.verbose)
   
    if len(args.files) < 2:
        sys.exit('You must specify at least two files to stack.')

    if args.group_by_filenames:
        groups = [os.path.split(f.name)[1] for f in args.files] 
    elif args.groups:
        groups = args.groups.split(',')

        if len(groups) != len(args.files):
            sys.exit('The number of grouping values must be equal to the number of CSV files being stacked.')
            
    group_name = args.group_name if args.group_name else 'group'

    output = CSVKitWriter(sys.stdout, **extract_csv_writer_kwargs(args))

    for i, f in enumerate(args.files):
        rows = CSVKitReader(f, **extract_csv_reader_kwargs(args))
        headers = rows.next()

        if i == 0:
            if args.groups:
                headers.insert(0, group_name)
            
            output.writerow(headers)

        for row in rows:
            if args.groups:
                row.insert(0, groups[i])

            output.writerow(row)

if __name__ == '__main__':
    parser = init_common_parser(description='Stack up the rows from multiple CSV files, optionally adding a grouping value.', omitflags='f')

    parser.add_argument('files', metavar="FILES", nargs='+', type=argparse.FileType('r'))
    parser.add_argument('-g', '--groups', dest='groups',
                        help='A comma-seperated list of values to add as "grouping factors", one for each CSV being stacked. These will be added to the stacked CSV as a new column. You may specify a name for the grouping column using the -n flag.')
    parser.add_argument('-n', '--group-name', dest='group_name',
                        help='A name for the grouping column, e.g. "year". Only used when also specifying -g.')
    parser.add_argument('--filenames', dest='group_by_filenames', action='store_true',
                        help='Use the filename of each input file as its grouping value. When specified, -g will be ignored.')

    main(parser.parse_args())
