#!/usr/bin/env python

import argparse
import re
import sys

from csvkit import CSVKitReader, CSVKitWriter
from csvkit.grep import FilteringCSVReader
from csvkit.cli import extract_csv_reader_kwargs, extract_csv_writer_kwargs, init_common_parser, parse_column_identifiers, print_column_names, install_exception_handler

def main(args):
    """
    Like the unix "grep" command, but for tabular data.
    """
    install_exception_handler(args.verbose)
   
    if args.names_only:
        print_column_names(args.file, sys.stdout, **extract_csv_reader_kwargs(args))
        sys.exit()

    if not args.pattern:
        sys.exit('A pattern must be specified unless using the -n option.')

    rows = CSVKitReader(args.file, **extract_csv_reader_kwargs(args))
    column_names = rows.next()

    column_ids = parse_column_identifiers(args.columns, column_names)
    
    if args.regex:
        pattern = re.compile(args.pattern)
    else:
        pattern = args.pattern
        
    patterns = dict((c,pattern) for c in column_ids)

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

    filter_reader = FilteringCSVReader(rows, header=False, patterns = patterns)

    for i, row in enumerate(filter_reader):
        output.writerow(row)
                
if __name__ == "__main__":
    """
    Process command line arguments.
    """
    parser = init_common_parser(description='Like the unix "grep" command, but for tabular data.', omitflags='f')
    parser.add_argument('-n', '--names', dest='names_only', action='store_true',
                        help='Display column names and indices from the input CSV and exit.')
    parser.add_argument('-c', '--columns', dest='columns',
                        help='A comma separated list of column indices or names to be searched.')
    parser.add_argument('-r', '--regex', dest='regex', action='store_true',
                        help='If specified, the search pattern will be treated as a regular expression.')
    parser.add_argument('pattern', metavar="PATTERN", nargs='?',
                        help='The pattern to search for.')
    parser.add_argument('file', metavar="FILE", nargs='?', type=argparse.FileType('rU'), default=sys.stdin,
                        help='The CSV file to operate on. If omitted, will accept input on STDIN.')

    main(parser.parse_args())

