from itemadapter import ItemAdapter
from scrapy import signals
from scrapy.exceptions import DropItem

import logging
import pymongo


logger = logging.getLogger(__name__)


class MongoPipeline:
    @classmethod
    def from_crawler(cls, crawler):
        s = cls(crawler.settings)
        crawler.signals.connect(s.spider_opened, signal=signals.spider_opened)
        crawler.signals.connect(s.spider_closed, signal=signals.spider_closed)
        return s

    def __init__(self, settings):
        self.url = settings.get('PIPELINE_MONGO_URL') or settings.get('MONGO_URL')
        assert self.url, 'PIPELINE_MONGO_URL or MONGO_URL is required'

        self.db_name = settings.get('PIPELINE_MONGO_DATABASE') or settings.get('MONGO_DATABASE')
        assert self.db_name, 'PIPELINE_MONGO_DATABASE or MONGO_DATABASE is required'

        self.buffer_max = settings.getint('PIPELINE_MONGO_BUFFER_MAX', 100)

        self._client = None
        self._database = None
        self._buffers_by_collections = dict()
        self._buffer_size = 0

    def spider_opened(self, spider):
        logger.debug("Connecting to MongoDB...")
        self._client = pymongo.MongoClient(self.url)
        self._database = self._client[self.db_name]

    def spider_closed(self, spider):
        logger.debug("Closing MongoDB connection...")

        self._flush(spider)
        self._client.close()

    def process_item(self, item, spider):
        item_json = ItemAdapter(item).asdict()

        # Update id field
        item_id = item_json.get('id')
        if item_id:
            item_json['_id'] = item_id
            del item_json['id']
        else:
            if '_id' not in item_json:
                raise DropItem("Item without an 'id' property")

        # Add in the correct collection buffer
        collection = item_json.get('collection')
        if not collection:
            raise DropItem("Item without a 'collection' property")

        item_order = pymongo.UpdateOne({
            '_id': item_id,
        }, {
            '$set': item_json
        }, upsert=True)

        buffers_by_collections = self._buffers_by_collections.get(collection)
        if buffers_by_collections:
            buffers_by_collections.append(item_order)
        else:
            self._buffers_by_collections[collection] = [item_order]

        self._buffer_size += 1

        if self._buffer_size >= self.buffer_max:
            self._flush(spider)

        return item

    def _flush(self, spider):
        if self._buffer_size <= 0:
            return

        logger.debug("Storing %d items to MongoDB..." % (self._buffer_size))

        # Bulk write items
        for collection_name, items in self._buffers_by_collections.items():
            self._database[collection_name].bulk_write(items)

        # Reset buffer
        self._buffers_by_collections.clear()
        self._buffer_size = 0
