from logging import getLogger

from django.conf import settings
from django.core.management.base import BaseCommand
from django.db import InterfaceError, DatabaseError, connection

from intranet.search.core.utils import get_oauth_token
from startrek_client import Startrek

from intranet.search.core.models import Indexation
from intranet.search.core.sources.utils import get_persons
from intranet.search.core.swarm.indexer import STORAGES, Indexer
from intranet.search.core.swarm.tasks import reindex, parse_to_list
from time import sleep


logger = getLogger(__name__)


def get_at_keys(keys_gt):
    persons_filter = {'_sort': 'id', '_fields': 'login', '_limit': 1000}
    if keys_gt:
        persons_filter['_query'] = f'id>{keys_gt}'
    return [p['login'] for p in get_persons(persons_filter)]


def get_st_keys(keys_gt):
    endpoint = settings.ISEARCH['api']['st']['endpoint']
    client = Startrek(useragent=endpoint.get('headers').get('User-Agent'),
                      base_url=endpoint.url(),
                      token=get_oauth_token(endpoint.get('oauth_token')))
    all_queues = sorted(q.key for q in client.queues.get_all())
    if keys_gt:
        try:
            idx = all_queues.index(keys_gt)
        except ValueError:
            pass
        else:
            all_queues = all_queues[idx+1:]
    return all_queues


KEYS_MAP = {
    ('at', ''): get_at_keys,
    ('st', ''): get_st_keys,
}


class Command(BaseCommand):
    help = 'Аккуратная переиндексация больших источников'

    def add_arguments(self, parser):
        parser.add_argument('search', type=str)
        parser.add_argument('index', nargs='?', type=str)

        parser.add_argument('--concurrency', '-c', action='store', type=int, default=1,
                            help='max count of simultaneous new indexations')
        parser.add_argument('--keys', action='store', type=str, default='',
                            help='keys for reindex')
        parser.add_argument('--keys-count', '-k', action='store', type=int, default=10,
                            help='keys count in one indexation')
        parser.add_argument('--keys-gt', action='store', type=str, default='',
                            help='indexation keys greater filter')
        parser.add_argument('--comment', action='store', type=str, default='',
                            help='task comment')
        parser.add_argument('--max_iterations', '-m', action='store', type=int, default=1,
                            help='max iterations count')
        parser.add_argument('--countdown', '-d', action='store', type=int, default=30,
                            help='sleep countdown between iterations')
        parser.add_argument('--revision-id', '-r', action='store', default='',
                            help='revision id for indexation')
        parser.add_argument('--document_storages', '--storages', action='store',
                            default=','.join(Indexer.default_options['document_storages']),
                            help='comma separated types of storages, available types: %s' % ', '.join(STORAGES.keys()))
        parser.add_argument('--suffix', action='store', default='', help='worker suffix')
        parser.add_argument('--noflush', action='store_true', default=False)

    def get_reindex_kwargs(self, options):
        kwargs = {
            'revision_id': options['revision_id'],
            'document_storages': options['document_storages'],
            'suffix': options['suffix'],
            'nolock': True,
            'nort': False,
            'flush': not options['noflush'],
            'shadow_revision': True,
        }
        return kwargs

    def get_keys(self, search, index, options):
        keys_getter = KEYS_MAP.get((search, index))
        if keys_getter:
            keys = keys_getter(options['keys_gt'])
        elif options['keys']:
            keys = sorted(parse_to_list(options['keys']))
        else:
            keys = []
        return keys

    def handle(self, search, index='', **options):
        status_filter = {'search': search, 'index': index, 'status': Indexation.STATUS_NEW}
        reindex_kwargs = self.get_reindex_kwargs(options)
        keys = self.get_keys(search, index, options)
        i = 0
        start = 0

        while i < options['max_iterations'] and start < len(keys):
            try:
                if Indexation.objects.filter(**status_filter).count() < options['concurrency']:
                    reindex_kwargs['keys'] = ','.join(keys[start:start + options['keys_count']])
                    start += options['keys_count']
                    reindex_kwargs['comment'] = '{}: {}'.format(options['comment'],
                                                                reindex_kwargs['keys'])
                    logger.debug(f'[BATCH] reindex {search}.{index}, {reindex_kwargs}')
                    reindex(search, index=index, **reindex_kwargs)
                else:
                    logger.debug('[BATCH] wait')
            except (DatabaseError, InterfaceError) as e:
                logger.error(e)
                connection.close()
                connection.ensure_connection()
            except Exception as e:
                logger.error(e)
            sleep(options['countdown'])
            i += 1
