#!/usr/bin/env python
# -*- coding: utf-8 -*-

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

import sys
import time
import heapq
import functools

from injection.utils.table_scan import TableScanner, TableScannerChunk
from injection.utils.watch_dog import WatchDog
from trans2es.type_info import TypeInfo, get_type_info_map
from search.utils.es import get_es


class Timer(object):

    def __init__(self, watchdog_timeout):
        self._heap = []
        self._counter = 0
        self._watch_dog = WatchDog(watchdog_timeout)

    @property
    def now(self):
        return time.time()

    def wait_and_run(self):
        if len(self._heap) == 0:
            return False
        exec_time, _, job = heapq.heappop(self._heap)
        now = self.now
        sleep_time = exec_time - now
        if sleep_time > 0:
            time.sleep(sleep_time)
        with self._watch_dog.guard():
            job()
        return True

    def put(self, exec_time, job):
        exec_time = float(exec_time)
        assert callable(job)

        now = self.now
        if exec_time < now:
            exec_time = now
        self._counter += 1
        heapq.heappush(self._heap, (exec_time, self._counter, job))

    def run_forever(self):
        while self.wait_and_run():
            pass

    def set_interval(self, period, callback):
        period = float(period)
        assert callable(callback)

        def delay_callback(exec_time):
            callback()
            exec_time += period
            self.put(exec_time, functools.partial(delay_callback, exec_time))
        first_exec_time = self.now + period
        self.put(first_exec_time, functools.partial(delay_callback, first_exec_time))


class DataSyncRunnable(object):
    __es = None

    def __init__(self, type_info, index_prefix):
        assert isinstance(type_info, TypeInfo)
        self._type_info = type_info
        self._index_prefix = index_prefix
        self._scanner = TableScanner(queryset=type_info.queryset)
        self._chunk_iter = self._scanner.chunks(chunk_size=self._type_info.round_insert_chunk_size)

    @classmethod
    def get_es(cls):
        if cls.__es is None:
            cls.__es = get_es()
        return cls.__es

    def start(self, timer):
        assert isinstance(timer, Timer)
        timer.set_interval(
            period=self._type_info.round_insert_period,
            callback=self.run,
        )

    def run(self):
        chunk = next(self._chunk_iter)
        assert isinstance(chunk, TableScannerChunk)
        result = self._type_info.insert_table_chunk(
            index_prefix=self._index_prefix,
            table_chunk=chunk,
            es=self.get_es(),
        )
        print(result)
        sys.stdout.flush()


class Command(BaseCommand):
    args = ''
    help = 'dump data to elasticsearch, roundly'

    from optparse import make_option

    option_list = BaseCommand.option_list + (
        make_option('-t', '--type', dest='type', help='type name to dump data to elasticsearch', metavar='TYPE'),
        make_option('-i', '--index-prefix', dest='index_prefix', help='index name to dump data to elasticsearch', metavar='INDEX_PREFIX'),
    )

    def handle(self, *args, **options):
        type_info_map = get_type_info_map()

        permitted_type_set = ['_all'] + type_info_map.keys()

        if not options['type']:
            raise CommandError('type name must be specified')
        elif options['type'] not in permitted_type_set:
            raise CommandError('unknown type name, type must be one of [{}]'.format(', '.join(permitted_type_set)))

        if options['type'] == '_all':
            type_name_list = list(type_info_map.keys())
        else:
            type_name_list = [options['type']]

        if options['index_prefix']:
            index_prefix = options['index_prefix']
        else:
            index_prefix = settings.ES_INDEX_PREFIX

        timer = Timer(watchdog_timeout=300)
        for type_name in type_name_list:
            type_info = type_info_map[type_name]
            data_sync_runnable = DataSyncRunnable(
                type_info=type_info,
                index_prefix=index_prefix,
            )
            data_sync_runnable.start(timer=timer)

        timer.run_forever()

