import threading
from itertools import imap
from multiprocessing import cpu_count
from threading import Thread, Lock

from django.conf import settings
from django.core.management import BaseCommand
from elasticsearch import helpers

from api.models.message import Conversation
from message.utils.es_abstract import get_migrate_esop, get_esop, table_conversation

console_lock = Lock()


def get_conversation_bulk_action(conv):
    return {
        '_index': 'gm_msg-conversation',
        '_type': 'conversation',
        '_id': conv['id'],
        '_source': conv
    }


def read_conversation_from_es(conversation_ids):
    if not conversation_ids:
        return []
    convs = get_esop().mget(
        table=table_conversation,
        body={'ids': [str(i) for i in conversation_ids]}
    )
    convs = convs['docs']
    return [m['_source'] for m in convs if m['found'] and 'id' in m['_source']]


def worker(start_id, stop_id, chunk_size=100):
    assert start_id <= stop_id
    with console_lock:
        print("[thread-{0}]conversations between {1} and {2} are starting".format(
            threading.currentThread().getName(), start_id, stop_id))
    for i in range(start_id, stop_id, chunk_size):
        right = min(stop_id, i + chunk_size - 1)

        # read from old and write to new one.
        print("range: {}-{}".format(i, right))
        conversations = read_conversation_from_es(range(i, right + 1))
        helpers.bulk(get_migrate_esop().client, imap(get_conversation_bulk_action, conversations))

    with console_lock:
        print("[thread-{0}]conversations between {1} and {2} are done".format(
            threading.currentThread().getName(), start_id, stop_id))


class Command(BaseCommand):
    def add_arguments(self, parser):
        parser.add_argument('-c', help="number of thread;defaults cpu's core", type=int)
        parser.add_argument('-s', help="sync start conversation id", type=int)
        parser.add_argument('-t', help="sync stop conversation id", type=int)

    def handle(self, *args, **options):
        num = cpu_count()
        start = 1
        stop = Conversation.objects.using(settings.SLAVE_DB_NAME).last().id
        if options['c']:
            num = options['c']
        if options['s']:
            start = options['s']
        if options['t']:
            stop = options['t']

        assert start <= stop
        print('max_conversation_id: {0}, thread_num: {1}'.format(stop, num))
        step = max((stop - start + 1) / num, 1)

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

        for t in threads:
            t.join()
