#!/usr/bin/env python

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

def main(args):
    """
    Like the unix "grep" command, but for tabular data.
    """
    rows = CSVKitReader(args.file, **extract_csv_reader_kwargs(args))
    column_names = rows.next()

    if args.names_only:
        for i, c in enumerate(column_names):
            sys.stdout.write('%3i: %s\n' % (i + 1, c))

        sys.exit()

    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.')
    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",
                        help='The pattern to search for.')

    main(parser.parse_args())

