#!/bin/python
import argparse
import sys
import logging
import datetime
import functools

import elasticsearch
import elasticsearch.helpers

import blueox.client

DEFAULT_HOST = "127.0.0.1:3513"

log = logging.getLogger('blueox.oxstash')


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"
    logging.basicConfig(level=level, format=log_format, stream=sys.stdout)


def build_doc(options, evt):
    timestamp = datetime.datetime.utcfromtimestamp(evt['end'])
    evt_type = evt.pop('type')

    if options.index_name:
        index = options.index_name
    else:
        index = "%s-%s" % (options.index_base, timestamp.strftime('%Y.%m.%d'))

    evt['duration'] = evt['end'] - evt['start']

    # This appears to be the default timestamp key logstash/kibana looks for.
    evt['@timestamp'] = timestamp.isoformat() + 'Z'

    action = {'_index': index, '_type': evt_type, '_source': evt,}

    return elasticsearch.helpers.expand_action(action)


def main():
    parser = argparse.ArgumentParser(
        description="Send blueox log events to elasticsearch")
    parser.add_argument('--verbose', '-v',
                        dest='verbose',
                        action='append_const',
                        const=True,
                        default=list())

    parser.add_argument(
        '--type-name', '-n',
        dest='type_name',
        action='store',
        default=None,
        help="What event types to display. Can end with '*' for prefix matches.")
    parser.add_argument('--host', '-H',
                        dest='host',
                        action='store',
                        default=None)
    parser.add_argument('--log-path', '-l',
                        dest='log_path',
                        action='store',
                        default=None)
    parser.add_argument(
        '--elasticsearch', '-e',
        action='append',
        default=[],
        help="Elasticsearch nodes to connect to (default: localhost:9200)")
    parser.add_argument('--index_name', '-i',
                        action='store',
                        default=None,
                        help="Elasticsearch index to use for events")
    parser.add_argument('--index_base', '-p',
                        action='store',
                        default="blueox",
                        help="Elasticsearch index to use, based on date")
    parser.add_argument(
        '--chunk_size',
        action='store',
        type=int,
        default=500,
        help="How many events to do together in a bulk operation")
    parser.add_argument('--request_timeout',
                        action='store',
                        type=int,
                        default=10,
                        help="Read timeout for bulk api")

    options = parser.parse_args()

    if len(options.elasticsearch) == 0:
        options.elasticsearch.append("localhost:9200")

    setup_logging(options)

    es = elasticsearch.Elasticsearch(options.elasticsearch)

    if sys.stdin.isatty() or options.host:
        log.info("Loading stream from oxd")

        if options.host:
            host = options.host
        else:
            host = DEFAULT_HOST

        out_stream = blueox.client.subscribe_stream(host, options.type_name)
    else:
        if options.type_name is not None:
            parser.error("Can't specify a name from stdin")
            sys.exit(1)

        log.info("Loading stream from stdin")
        out_stream = blueox.client.stdin_stream()

    action_callback = functools.partial(build_doc, options)
    bulk_streamer = elasticsearch.helpers.streaming_bulk(
        es, out_stream,
        request_timeout=options.request_timeout,
        chunk_size=options.chunk_size,
        expand_action_callback=action_callback)

    for result in bulk_streamer:
        batch_ok, item = result
        if item['create'].get('ok') is not True:
            print item


if __name__ == '__main__':
    try:
        main()
    except (KeyboardInterrupt, SystemExit):
        pass
