#!/usr/bin/python
import argparse
import datetime
import os
import sys
import logging
import pprint
import base64
import time
import json

import triton
import triton.store

log = logging.getLogger('triton')

GET_COMMAND = 'get'
PUT_COMMAND = 'put'
CAT_COMMAND = 'cat'


def setup_logging(options):
    if len(options.verbose) > 1:
        level = logging.DEBUG
    elif options.verbose:
        level = logging.INFO
    else:
        level = logging.WARNING

    log_format = "%(asctime)s %(levelname)s:%(name)s: %(message)s"

    # We just log straight to stdout. Generally oxd is run by some process that
    # handles collecting logging for you, like upstart or supervisord. It would be
    # easy enough to add an option if it was needed by someone though.
    logging.basicConfig(level=level, format=log_format, stream=sys.stdout)


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--verbose',
                        '-v',
                        dest='verbose',
                        action='append_const',
                        const=True,
                        default=list())

    subparsers = parser.add_subparsers(
        dest='command',
        help='available commands')

    parser_put = subparsers.add_parser(
        PUT_COMMAND,
        help='read lines from stdin, put to stream')

    parser_put.add_argument('--stream',
                            '-s',
                            dest='stream',
                            action='store',
                            required=True)

    parser_get = subparsers.add_parser(
        GET_COMMAND,
        help='read records from stream')

    parser_get.add_argument('--stream',
                            '-s',
                            dest='stream',
                            action='store',
                            required=True)
    parser_get.add_argument('--shard', '-d', dest='shards', type=int, nargs='+')

    parser_cat = subparsers.add_parser(
        CAT_COMMAND,
        help='read records from stdin, file or S3 bucket')
    parser_cat.add_argument('--stream',
                            '-s',
                            dest='stream',
                            action='store',
                            required=False)
    parser_cat.add_argument('--bucket',
                            '-b',
                            dest='bucket',
                            action='store',
                            required=False)
    parser_cat.add_argument('--start-date',
                            '-t',
                            dest='start_date',
                            action='store',
                            required=False)
    parser_cat.add_argument('--end-date',
                            '-e',
                            dest='end_date',
                            action='store',
                            required=False)

    args = parser.parse_args()

    setup_logging(args)

    config = triton.load_config(os.environ.get('TRITON_CONFIG',
                                               '/etc/triton.yaml'))
    if not config:
        parser.error("Failed to load TRITON_CONFIG")

    if args.command == GET_COMMAND:
        cmd_stream = triton.get_stream(args.stream, config)

        running = True
        i = cmd_stream.build_iterator_from_latest(shard_nums=args.shards)
        while True:
            try:
                for rec in i:
                    print rec
                    pprint.pprint(rec.data)
                    print
            except KeyboardInterrupt:
                i.stop()
                continue
            else:
                break

    elif args.command == PUT_COMMAND:
        cmd_stream = triton.get_stream(args.stream, config)

        for line in sys.stdin:
            cmd_stream.put(msg=line, ts=time.time())

    elif args.command == CAT_COMMAND:
        if not sys.stdin.isatty():
            stream = triton.store.decoder(sys.stdin)
        elif args.bucket:
            region = os.environ.get('AWS_DEFAULT_REGION', 'us-west-1')

            if not (args.start_date and args.end_date):
                parser.error("Dates required")

            start_dt = datetime.datetime.strptime(args.start_date, '%Y%m%d')
            end_dt = datetime.datetime.strptime(args.end_date, '%Y%m%d')

            bucket = triton.store.open_bucket(args.bucket, region)
            stream = triton.stream_from_s3_store(bucket, config[args.stream],
                                                 start_dt, end_dt)
        else:
            parser.error("Nothing to do")

        for rec in stream:
            print json.dumps(rec)


if __name__ == '__main__':
    main()
