#!/usr/bin/env python
"""
Generate deliverables

Usage:

    deliver.py <spider> [options]

For listing valid values for spider argument, just run deliver.py without arguments.

This script does the following:

* searches for untagged finished crawl_manager.py jobs that corresponds to the given spider
* read their logs and finds all consumer jobs scheduled by it
* gets all the items generated by them and uploads to customer s3 bucket in files of 1M (by default) items each one
* tags processed crawl_manager jobs so they are not processed again
* tags processed spider jobs so they are not reprocessed as independent (no crawl manager) job, if deliver is ran again
  and there is no new finished crawl manager

If no crawl_manager job was found that handles the selected spider:

* searches for untagged target spider jobs
* gets all the items generated by them and uploads to customer s3 bucket in files of 1M (by default) items each one
* tags processed jobs so they are not processed again

Settings
--------

S3_BUCKET_NAME - Target s3 bucket

S3_FILE_ROOT - Root prefix for s3 files. Default value is empty string.
S3_FILE_PREFIX - Defines the file folder template.
                 {name} is replaced by scraper name, {time} is replaced by time (in the format given by
                 FILETIME_FORMAT, {itemtype} is replaced by item type, which is extracted from `_type`
                 metafield on each item. {field:<fieldname>} is replaced by the value of the given item field.
                 Both {field} and {type} variables implies to generate separate files for items with different
                 _type and specified field value. {argument:<argumentname>} is replaced by the value of the given
                 spider argument for the job from where the item was read.
S3_FILE_PREFIX_FOR_SPIDER - A dict that maps from a scraper name to a file folder (defined in the same way as
                S3_FILE_PREFIX). If given, first searches scraper name in this dict for getting the file folder
                template. If scraper name is not found, fallbacks to S3_FILE_PREFIX.

Final s3 key name prefix is build by concatenation S3_FILE_ROOT, rendered S3_FILE_PREFIX.
An alternative for more complex custom handling of file prefix, is to subclass the DeliverScript class and override
the method get_fileprefix_template()

S3_FILE_TIMEFORMAT - Format of the timestamp that replaces {time} in S3_FILE_PREFIX
S3_FILE_SIZE - Max size of each file. Default is 1,000,000

Following settings are also required:

AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY

Other settings:

S3_SUCCESS_FILE - Boolean. If True, an empty file will be uploaded at the end of the script, inside each lowest
                  hierarchy folder generated. Default value is False.

"""
import json
import re
import os
import gzip
import datetime
import logging
import time
import tempfile
from argparse import ArgumentParser
from tempfile import mktemp

from s3fs import S3FileSystem
from sqlitedict import SqliteDict

from scrapy.utils.project import get_project_settings
from scrapy.utils.conf import init_env

from scrapinghub import ScrapinghubClient

from .utils import resolve_project_id, kumo_settings


_LOG = logging.getLogger(__name__)
_LOG.setLevel(logging.INFO)

init_env()


settings = get_project_settings()
settings.setdict(kumo_settings())


_FIELD_RE = re.compile(r'\{field:(.+?)\}')
_ARGUMENT_RE = re.compile(r'\{argument:(.+?)\}')


class SqliteDictDupesFilter(object):
    def __init__(self):
        """
        SqlteDict based dupes filter
        """
        self.dupes_db_file = tempfile.mktemp()
        self.__filter = None

    def __create_db(self):
        self.__filter = SqliteDict(self.dupes_db_file, flag='n', autocommit=True)

    def __contains__(self, element):
        if self.__filter is None:
            self.__create_db()
        return element in self.__filter

    def add(self, element):
        if self.__filter is None:
            self.__create_db()
        self.__filter[element] = '-'

    def close(self):
        if self.__filter is not None:
            try:
                self.__filter.close()
                os.remove(self.dupes_db_file)
            except:
                pass


def _upload_file_to_s3(bucketname, keyname, filename=None):
    aws_access_key_id = settings['AWS_ACCESS_KEY_ID']
    aws_secret_access_key = settings['AWS_SECRET_ACCESS_KEY']
    s3 = S3FileSystem(key=aws_access_key_id, secret=aws_secret_access_key)
    with s3.open(f's3://{bucketname}/{keyname}', 'wb') as f:
        if filename is None:
            f.write(b'')
        else:
            with open(filename, 'rb') as r:
                f.write(r.read())


class OutputFile(object):

    def __init__(self, keyprefix, testmode=False):
        self.__key_prefix = keyprefix

        self.__filename = mktemp()
        self.__file = gzip.open(self.__filename, "wb")
        self.__filecount = 0
        self.__count = 0
        self.__filesize = settings.get('S3_FILE_SIZE', 1000000)
        self.__bucketname = settings['S3_BUCKET_NAME']
        self.__testmode = testmode

        spath = self.__key_prefix.split('/')[:-1]
        spath.append('_SUCCESS')
        self.success_file = os.path.join(*spath)

    def write(self, item):
        self.__file.write(json.dumps(item).encode('utf8') + b'\n')
        self.__count += 1
        if self.__count == self.__filesize:
            self.flush(new_file=True)

    def flush(self, new_file=False):
        if self.__count > 0:
            self.__file.close()
            self._upload_file()
            self.__filecount += 1
            self.__count = 0
            if new_file:
                self.__file = gzip.open(self.__filename, "wb")
        elif not new_file:
            self.__file.close()

    def _upload_file(self):
        keyname = self.__key_prefix + "%05d.jsonl.gz" % self.__filecount
        if not self.__testmode:
            _upload_file_to_s3(self.__bucketname, keyname, self.__filename)
        _LOG.info("Saved %d items in s3://%s/%s", self.__count, self.__bucketname, keyname)


class OutputFileDict(object):

    def __init__(self, testmode=False):
        self.__testmode = testmode
        self.__outputfiles = {}

    def __getitem__(self, key):
        if key not in self.__outputfiles:
            self.__outputfiles[key] = OutputFile(keyprefix=key, testmode=self.__testmode)
        return self.__outputfiles[key]

    def values(self):
        return self.__outputfiles.values()

    def keys(self):
        return self.__outputfiles.keys()

    def pop(self, key):
        return self.__outputfiles.pop(key)


class DeliverScript(object):

    def __init__(self):
        self.project_id = resolve_project_id()
        self.client = ScrapinghubClient()
        self.args = self.parse_args()

        for scrapername in self.args.scrapername:
            _LOG.info("Delivery folder template for spider %s: %s", scrapername,
                      self.get_fileprefix_template(scrapername))

        self.set_output_files()
        self.set_fileformat()

        self.itemcount = 0
        self.filecount = 0
        self.dupes_filter = {i: SqliteDictDupesFilter() for i in self.args.filter_dupes_by_field}

    def set_output_files(self):
        self.output_files = OutputFileDict(self.args.test_mode)
        
    def set_fileformat(self):
        filetime_format = settings.get('S3_FILE_TIMEFORMAT', '%Y-%m-%dT%H:%M:%S')
        self.__datetime = datetime.datetime.now().strftime(filetime_format)

    def parse_args(self):
        self.argparser = ArgumentParser(description=__doc__)
        try:
            _LOG.info("Project selected: %d (if this is not the correct project, set it via PROJECT_ID\
                      environment variable)", int(self.project_id))
        except:
            self.argparser.error("Wrong project %s" % self.project_id)

        self.add_argparser_options()
        args = self.argparser.parse_args()
        return args

    def add_argparser_options(self):

        project = self.client.get_project(self.project_id)
        all_spiders = [s['id'] for s in project.spiders.iter()]

        self.argparser.add_argument('scrapername', help='Indicate scraper name', choices=sorted(all_spiders), nargs='*')
        self.argparser.add_argument('--filter-dupes-by-field', default=[],
                                    help='Dedupe by any of the given item field. Can be given multiple times')
        self.argparser.add_argument('--one-file-per-job', action='store_true', help='Generate one file per job.')
        self.argparser.add_argument('--test-mode', action='store_true',
                                    help='Run in test mode (performs all processes, but doesn\'t\
                                          upload files nor tag jobs)')
        self.argparser.add_argument('--name',
                                    help='Use the given value as replacement text for {name}. By default \
                                    uses spider name.')

    def get_fileprefix_template(self, scrapername):
        root_prefix = settings.get('S3_FILE_ROOT', '')
        default_prefix = settings['S3_FILE_PREFIX']
        prefix = settings.getdict('S3_FILE_PREFIX_FOR_SPIDER').get(scrapername, default_prefix)
        return os.path.join(root_prefix, prefix)

    def gen_keyprefix(self, scrapername, job, item, fileprefix):

        # don't separate files per itemtype if {itemtype} is not given on FILEPREFIX
        itemtype = ""
        if '{itemtype}' in str(fileprefix):
            itemtype = item.pop("_type").lower()

        field = ''
        m = _FIELD_RE.search(fileprefix)
        if m:
            fieldname = m.groups()[0]
            if fieldname in item:
                field = item[fieldname]
                fileprefix = _FIELD_RE.sub("{field}", fileprefix)
            else:
                raise KeyError('field %s not found in item' % fieldname)

        argument = ''
        m = _ARGUMENT_RE.search(fileprefix)
        if m:
            argumentname = m.groups()[0]
            if argumentname in job.metadata['spider_args']:
                argument = job.metadata['spider_args'][argumentname]
                fileprefix = _ARGUMENT_RE.sub("{argument}", fileprefix)
            else:
                raise KeyError('argument %s not found in job' % argumentname)

        return fileprefix.format(
            name=self.args.name or scrapername,
            time=self.__datetime,
            itemtype=itemtype,
            field=field,
            argument=argument
        )

    def _process_job_items(self, scrapername, spider_job):
        first_keyprefix = None
        for item in spider_job.items.iter_values():
            seen = False
            for field in self.dupes_filter.keys():
                if field in item:
                    if item[field] in self.dupes_filter[field]:
                        seen = True
                    else:
                        self.dupes_filter[field].add(item[field])
            if seen:
                continue

            try:
                fileprefix = self.get_fileprefix_template(scrapername)
                keyprefix = self.gen_keyprefix(scrapername, spider_job, item, fileprefix)
                first_keyprefix = first_keyprefix or keyprefix
            except KeyError as e:
                if first_keyprefix is None:
                    _LOG.info("Skipped job: %s. %s", spider_job.key, str(e))
                    return
                keyprefix = first_keyprefix
            self.output_files[keyprefix].write(item)
            self.itemcount += 1
            if self.itemcount % 100000 == 0:
                _LOG.info("Processed %d items.", self.itemcount)
        _LOG.info("Processed spider job %s", spider_job.key)

    def run(self):
        hsc = self.client._hsclient
        hsp = hsc.get_project(self.project_id)

        success_files = set()
        all_jobs_to_tag = set()
        for scrapername in self.args.scrapername:
            _LOG.info("Processing spider %s", scrapername)
            jobs_to_tag = self.process_crawlmanager_jobs(scrapername, hsp)
            if not jobs_to_tag:
                jobs_to_tag = self.process_spider_jobs(scrapername, hsp)
            all_jobs_to_tag.update(jobs_to_tag)

            for ofile in self.output_files.values():
                ofile.flush()
                if settings.getbool('S3_SUCCESS_FILE'):
                    success_files.add(ofile.success_file)

            _LOG.info("Total Processed items for spider %s: %d", scrapername, self.itemcount)

        jcount = 0
        if self.args.test_mode:
            all_jobs_to_tag = set()
        for jkey in all_jobs_to_tag:
            j = hsp.get_job(jkey)
            tags = j.metadata.get("tags", [])
            if 'delivered' not in tags:
                tags.append("delivered")
                j.update_metadata(tags=tags)
            jcount += 1
            if jcount % 100 == 0:
                _LOG.info("Marked %d jobs as delivered", jcount)

        for success_file in success_files:
            if not self.args.test_mode:
                _upload_file_to_s3(settings['S3_BUCKET_NAME'], success_file)
            _LOG.info("Created s3://%s/%s", settings['S3_BUCKET_NAME'], success_file)

        self.close(success_files)

    def close(self, success_files):
        pass

    def process_spider_jobs(self, scrapername, hsp):
        jobs_to_tag = []
        for spider_job in hsp.jobq.list(spider=scrapername, count=100, state="finished", lacks_tag="delivered"):
            sj = hsp.get_job(spider_job['key'])
            self._process_job_items(scrapername, sj)
            jobs_to_tag.append(spider_job['key'])
        return jobs_to_tag

    def process_crawlmanager_jobs(self, scrapername, hsp):
        jobs_to_tag = []
        schedule_re = re.compile(r"%s: scheduled (\d+/\d+/\d+)" % scrapername)
        running = list(hsp.jobq.list(spider="py:crawl_manager.py", count=100, state="running"))
        finished = list(hsp.jobq.list(spider="py:crawl_manager.py", count=100, state="finished", lacks_tag="delivered"))
        for crawler_job in finished + running:
            cj = hsp.get_job(crawler_job['key'])
            for option in cj.metadata['job_cmd'][1:]:
                if re.match(r"--spider\s*=?\s*{}($|\s)".format(scrapername), option.strip()):
                    while cj.metadata['state'] == 'running':
                        _LOG.info("Crawl manager job %s is still running. Waiting 5 mins before continue...")
                        time.sleep(300)
                        cj = hsp.get_job(crawler_job['key'])
                    _LOG.info("Processing crawler job %s", crawler_job['key'])
                    for logline in cj.logs.iter_values():
                        if isinstance(logline, dict):
                            m = schedule_re.search(logline.get('message', ''))
                            if m:
                                sjkey = m.groups()[0]
                                sj = hsp.get_job(sjkey)
                                if sj.metadata['state'] == 'finished':
                                    self._process_job_items(scrapername, sj)
                                    # job will be tagged as delivered in order to avoid to be processed again as
                                    # standalone job, but if crawl manager is untagged, it will (and must) be
                                    # processed again. That is why we don't check in this method whether the job is
                                    # already tagged.
                                    jobs_to_tag.append(sjkey)
                                    if self.args.one_file_per_job:
                                        for val in self.output_files.values():
                                            val.flush(new_file=True)
                    jobs_to_tag.append(crawler_job['key'])
                    break
            else:
                _LOG.info("Skipped jobid %s", crawler_job['key'])
        return jobs_to_tag


if __name__ == '__main__':
    deliver = DeliverScript()
    deliver.run()
