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

import datetime, time, sys, os, json
import logging, argparse

from grab.spider import Spider, Task

import dbdrivers, utils, display
import likes, lenta, lifenews

class NewsParser(Spider):

    def __init__(self, config, modules, display, **kwargs):
        Spider.__init__(self, **kwargs)
        self.modules = modules
        self.config = config
        self.display = display
        self.display.parser = self
        self.major_tasks = []  # It affects only how theese tasks will be displayed
        schema = []
        for module in self.modules:
            schema.extend(modules[module].schema)
            self.major_tasks.extend(getattr(modules[module], 'major_tasks', []))
        self.schema = schema
        if self.config['repair']:
            self.repair_tables = [
                table_name
                for table_name, _ in schema
                if self.find_modules_method('hook_' + table_name)[0]
            ]
        else:
            self.repair_tables = []
        self.hosts = { # Elements have following structure:
            # 'http://lenta.ru/' : {
            #      'last_request_time' : datetime.datetime(...),
            #      'last_url' : 'http://lenta.ru/news/world/2015/05/12/merkel'
            #      'success' : 10,
            #      'failure' : 0,
            #      'in_process' : False,
            # }
        }
        self.db_stats = {
            table : {
                # 'missed_rows' : 0,
                # 'init' : self.db.rows_count(table),
                'added' : 0,   # Number of rows added
            }
            for table, _ in self.schema
        }
        self.prepare_count = 0

    def prepare(self):
        self.display.prepare_begin()
        self.prepare_count += 1
        self.reached_repair_limit = False
        if 'csv' in self.config:
            self.db = dbdrivers.CSV(self.schema, self.db_write_callback, self.display, self.config['csv'])
        elif 'sqlite' in self.config:
            if self.config['repair']:
                cache_keys = zip(*self.schema)[0]
            else:
                cache_keys = []
            self.db = dbdrivers.SQLite(self.schema, self.db_write_callback, self.display, self.config['sqlite'], cache_keys=cache_keys)
        for module in self.modules:
            self.modules[module].db = self.db
        self.display.db = self.db
        self.modules_state = {
            module_name : {
                'generator' : getattr(module, 'generate_tasks')()
                              if hasattr(module, 'generate_tasks') and not self.config['notadd']
                              else None,
                'gen_queue_count' : 0,  # Number of tasks from this module in gen_queue
                'max_gen_queue_pos' : 0,
                'hook_queue_count' : 0,
                'took_count' : 0,
            }
            for module_name, module in self.modules.items()
        }
        self.gen_queue = []
        self.hooks_queue = []
        self.db.prepare(self.repair_tables)
        for table in self.db.schema:
            if 'init' not in self.db_stats[table]:
                self.db_stats[table]['init'] = self.db.prepare_info[table]['rows_read']
            if 'missed_rows' in self.db.prepare_info[table]:
                self.db_stats[table]['missed_rows'] = self.db.prepare_info[table]['missed_rows']
        self.display.prepare_end()

    def get_cur_curl_tasks(self):
        return [
            reg_obj['task'].url
            for reg_obj in self.transport.registry.values()
        ]

    def can_make_request(self, url, curl_urls):
        host = utils.get_host(url)
        # First, make sure there are no active requests to this host
        for active_task in curl_urls:
            task_host = utils.get_host(active_task)
            if task_host == host:
                return False
        # Than, check allowed request rate from config
        if host in self.config['hosts'] and host in self.hosts:
            timedelta = datetime.datetime.now() - self.hosts[host]['last_request_time']
            if timedelta.total_seconds() < self.config['hosts'][host]['min_delay']:
                return False
        return True

    def populate_gen_queue(self):
        max_size = self.config['module_max_queue_count']
        for i in range(max_size):    # Here will be unneccessary cycles, but algorithm is simplier
            for module, module_info in self.modules_state.items():
                if module_info['generator'] is not None:
                    if self.config['balance']:
                        # If we are in "balance" mode,
                        # there there should be no more than max_size tasks in total in gen_queue and hooks_queue
                        if module_info['gen_queue_count'] + module_info['hook_queue_count'] >= max_size:
                            continue
                    else:
                        # If we are not in "balance" mode,
                        # there should be no more than max_size tasks after last taken task from gen_queue
                        if module_info['gen_queue_count'] >= module_info['max_gen_queue_pos'] + max_size:
                            continue
                    try:
                        task = module_info['generator'].next()
                    except StopIteration:
                        module_info['generator'] = None
                    else:
                        if task is not None:
                            task.source_task = module
                            task.module = module
                            module_info['gen_queue_count'] += 1
                            self.gen_queue.append(task)

    def get_next_task(self):
        # first, look in hooks_queue
        curl_urls = self.get_cur_curl_tasks()
        for task_pos, task in enumerate(self.hooks_queue):
            if self.can_make_request(task.url, curl_urls):
                # Atomic, but I'm not sure there will be no problems with
                # async call of db_write_callback
                self.hooks_queue[task_pos:task_pos+1] = []
                self.modules_state[task.module]['hook_queue_count'] -= 1
                self.display.task_from_hook_queue(task)
                return task
        self.populate_gen_queue()
        # if not found, look in generators queue
        for task_pos, task in enumerate(self.gen_queue):
            if self.can_make_request(task.url, curl_urls):
                module_state = self.modules_state[task.module]
                module_state['gen_queue_count'] -= 1
                module_state['took_count'] += 1
                module_state['max_gen_queue_pos'] = max(module_state['max_gen_queue_pos'], task_pos)
                self.gen_queue[task_pos:task_pos+1] = []
                self.display.task_from_gen_queue(task)
                return task
        for module in self.modules_state.values():
            if module['generator'] is not None:
                return None
        return False

    def find_modules_method(self, method):
        for module_name, module_obj in self.modules.items():
            if hasattr(module_obj, method):
                return module_name, getattr(module_obj, method)
        return None, None

    def task_generator(self):
        while True:
            task = self.get_next_task()
            if task is False:
                if self.reached_repair_limit:
                    self.display.repair_restart()
                else:
                    self.display.all_tasks_complete()
                return
            if task is None:
                yield None
                continue
            host = utils.get_host(task.url)
            if host not in self.hosts:
                self.hosts[host] = {
                    'success' : 0,
                    'failure' : 0,
                }
            self.hosts[host]['last_url'] = task.url
            self.hosts[host]['in_process'] = True
            method_name = 'task_' + task.name
            module, handler = self.find_modules_method(method_name)
            if handler is None:
                raise Exception('Can\'t find method ' + method_name + ' in modules')
            def wrapper_wrapper(host, handler):
                def wrapper(grab, task):
                    self.display.task_started(task)
                    self.hosts[host]['success'] += 1
                    self.hosts[host]['in_process'] = False
                    handler(grab, task)
                return wrapper
            task.callback = wrapper_wrapper(host, handler)
            yield task

    def db_write_callback(self, table, row, repair_mode=False):
        if not repair_mode:
            self.db_stats[table]['added'] += 1
        hook_method_name = 'hook_' + table
        module, handler = self.find_modules_method(hook_method_name)
        if handler is not None:
            for task in handler(row):
                if repair_mode:
                    self.db.prepare_info[table]['missed_rows'] += 1
                task.source_task = hook_method_name
                if getattr(task, 'module', None) is None:
                    task.module = module
                if repair_mode and self.modules_state[task.module]['hook_queue_count'] > self.config['repair']:
                    self.reached_repair_limit = True
                    break
                self.modules_state[task.module]['hook_queue_count'] += 1
                self.hooks_queue.append(task)
                self.display.add_hook_callback_task(task)

    def log_failed_network_result(self, res):
        host = utils.get_host(res['task'].url)
        self.hosts[host]['failure'] += 1
        self.hosts[host]['in_process'] = False
        if res['ok']:
            msg = 'http-%s' % res['grab'].response.code
        else:
            msg = res['error_abbr']
        self.display.failed_network_result(res['task'].url, msg)
        super(NewsParser, self).log_failed_network_result(res)

    def process_handler_error(self, func_name, ex, task):
        self.display.failed_parser_result(task.url)
        super(NewsParser, self).process_handler_error(func_name, ex, task)

    def submit_task_to_transport(self, task, *args):
        host = utils.get_host(task.url)
        self.hosts[host]['last_request_time'] = datetime.datetime.now()
        super(NewsParser, self).submit_task_to_transport(task, *args)
        self.display.task_submitted_to_transport()

    def shutdown(self):
        self.display.finalize_begin()
        if hasattr(self, 'db'):
            self.db.finalize()
        self.display.finalize_end()

def print_log(message):
    print(message.replace('\n', ''))
    sys.stdout.flush()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='newsgrab version 3.0.')
    db_type = parser.add_mutually_exclusive_group(required=True)
    db_type.add_argument('-csv', type=str, help='directory for csv tables', metavar='CSV_DIR')
    db_type.add_argument('-sqlite', type=str, help='file with sqlite3 database', metavar='FILE.DB')
    parser.add_argument('-config', type=str, help='configuration file in JSON format', default=None, metavar='CONFIG.JSON')
    parser.add_argument('-report', type=str, help='append short report to log file', default='newsgrab_report.log', metavar='FILE.LOG')
    # parser.add_argument('-curses', help='show realtime information to console', action="store_true")
    parser.add_argument('-cron', help='print report to stdout', action="store_true")
    parser.add_argument('-log', type=str, help='spider request log file', default=None, metavar='FILE.LOG')
    # parser.add_argument('-threads', type=int, help='set number of threads', default=5)
    # parser.add_argument('-repair', type=int, help='check database for unfinished hooks and retry them', default=0)
    # parser.add_argument('-notadd', help='disable generators', action="store_true")
    # parser.add_argument('-balance', help='do not limit queue size', action="store_true")
    args = parser.parse_args()
    logger = logging.getLogger('grab')
    logger.setLevel(logging.DEBUG)
    err = logging.StreamHandler(sys.stderr)
    err.setLevel(logging.ERROR)
    logger.addHandler(err)
    if args.config:
        config = json.loads(open(args.config).read())
        config['config'] = args.config
    else:
        config = {'threads' : 10, 'repair' : 0, 'hosts': {}, 'likes' : None, 'lenta' : None, 'lifenews' : None}
    if args.log is not None:
        config['log'] = args.log
        open(config['log'], 'w').close()
        logger.addHandler(logging.FileHandler(config['log']))
    display_obj = display.Default(sys.stdout, open(args.report, 'a'), cron=args.cron)
    if args.csv:
        config['csv'] = args.csv
    if args.sqlite:
        config['sqlite'] = args.sqlite
    if args.threads:
        config['threads'] = args.threads
    config['notadd'] = False
    config['balance'] = True
    config['module_max_queue_count'] = 10
    module_classes = {
        'likes' : likes.Module,
        'lenta' : lenta.Module,
        'lifenews' : lifenews.Module,
    }
    modules = {}
    for module_name, module_class in module_classes.items():
        if module_name in config:
            modules[module_name] = module_class(config[module_name])
    try:
        while True:
            spider = NewsParser(config, modules, display_obj, thread_number=config['threads'], max_task_generator_chunk=1)
            spider.run()
            if not spider.reached_repair_limit:
                break
    except KeyboardInterrupt:
        print('\nKeyboardInterrupt received')
        display_obj.keyboard_int()
    display_obj.finalize()

