import threading
from multiprocessing import cpu_count

from django.conf import settings
from django.core.management import BaseCommand
from message.utils.es_abstract import get_migrate_esop, get_esop
from api.models.message import Message, Conversation

_map = {
    'conversation': {
        'index': 'gm_msg-conversation',
        'type': 'conversation'
    },
    'message': {
        'index': 'gm_msg-message',
        'type': 'message'
    }
}
_missing = []
_write_lock = threading.Lock()


def read_from_es(index, flag, seq):
    """
        flag indentify read from older or newer, 0:older, 1:newer
    """
    op = get_migrate_esop() if flag else get_esop()
    data = op.client.mget(
        body={'ids': [str(i) for i in seq]},
        index=_map[index]['index'],
        doc_type=_map[index]['type']
    )

    # filter not found object
    data = data['docs']
    return [d['_source'] for d in data if d['found']]


def worker(index, start, stop, chunk=100):
    for i in range(start, stop, chunk):
        right = min(stop, i + chunk - 1)
        old = read_from_es(index, 0, range(i, right))
        new = read_from_es(index, 1, range(i, right))
        if old != new:
            m1 = {item['id']: item for item in old}
            m2 = {item['id']: item for item in new}
            # find missing ids.
            for k, v in m1.iteritems():
                if k not in m2 or v != m2[k]:
                    with _write_lock:
                        _missing.append(k)


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('--index', help="[message|conversation]",
                            type=str)
        parser.add_argument('-c', help="number of thread.defaults cpu's core", type=int)

    def handle(self, *args, **options):
        num = cpu_count()
        if options['c']:
            num = options['c']

        assert options['index'] in ['message', 'conversation']
        model = Message if options['index'] == 'message' else Conversation
        stop = model.objects.using(settings.SLAVE_DB_NAME).last().id
        print("max_id: {}".format(stop))

        step = max(stop / num, 1)
        threads = []
        for idx, i in enumerate(range(1, stop, step)):
            right = min(stop, i + step - 1)
            t = threading.Thread(target=worker, args=(options['index'], i, right), name=idx + 1)
            threads.append(t)
            t.start()

        for t in threads:
            t.join()

        print("Missing {} ids:".format(options['index']))
        print(_missing)
