# coding=utf8
from __future__ import unicode_literals, absolute_import, print_function

from django.core.management.base import BaseCommand, CommandError

from injection.data_sync.impl.process.multi_generation import run


class Command(BaseCommand):
    help = 'run worker for data_sync'

    def add_arguments(self, parser):
        parser.add_argument('--configuration', default='pk_data_source')
        parser.add_argument('--connection-type', default='es', dest='connection_type')
        parser.add_argument('-t', '--type', dest='type_list', action='append'),
        parser.add_argument('--trace-kafka-brokers', dest='trace_kafka_brokers')
        parser.add_argument('--trace-kafka-topic', dest='trace_kafka_topic')
        parser.add_argument('--use-pubsub', default=False, type=bool, dest='use_pubsub')
        parser.add_argument('--max-queue-size', default=None, type=int, dest='max_queue_size')

    def handle(self, *args, **options):
        type_list = options['type_list'] or None
        run(
            configuration=options['configuration'],
            connection_type=options['connection_type'],
            data_source_name_list=type_list,
            trace_kafka={
                'brokers': options['trace_kafka_brokers'],
                'topic': options['trace_kafka_topic'],
            },
            use_pubsub=options['use_pubsub'],
            max_queue_size=options['max_queue_size'],
        )
