from scrapy.http import Headers
from scrapy.responsetypes import responsetypes

import datetime
import logging
import gzip
import pickle  # nosec
import pymongo


logger = logging.getLogger(__name__)


class MongoCacheStorage:
    def __init__(self, settings):
        self.expiration_secs = settings.getint('HTTPCACHE_EXPIRATION_SECS', 604800) # 1 week

        self.url = settings.get('HTTPCACHE_MONGO_URL') or settings.get('MONGO_URL')
        assert self.url, 'HTTPCACHE_MONGO_URL or MONGO_URL is required'

        self.db_name = settings.get('HTTPCACHE_MONGO_DATABASE', 'scraping')
        self.col_name = settings.get('HTTPCACHE_MONGO_COLLECTION', 'cache')

        self._client = None
        self._col = None
        self._fingerprinter = None

    def open_spider(self, spider):
        self._client = pymongo.MongoClient(self.url)
        self._col = self._client[self.db_name][self.col_name]

        # create index
        self._col.create_index([('created_at', pymongo.ASCENDING)], expireAfterSeconds=self.expiration_secs)

        logger.debug(
            'Using Mongo cache storage',
            extra={'spider': spider},
        )

        assert spider.crawler.request_fingerprinter
        self._fingerprinter = spider.crawler.request_fingerprinter

    def close_spider(self, spider):
        self._client.close()

    def retrieve_response(self, spider, request):
        id = self._fingerprinter.fingerprint(request).hex()
        item = self._col.find_one({'_id': id})
        if not item:
            return None

        data_raw_compressed = item['data']
        data_raw = gzip.decompress(data_raw_compressed)
        data = pickle.loads(data_raw)

        url = data['url']
        status = data['status']
        headers = Headers(data['headers'])
        body = data['body']

        respcls = responsetypes.from_args(headers=headers, url=url, body=body)
        response = respcls(url=url, headers=headers, status=status, body=body)

        return response

    def store_response(self, spider, request, response):
        id = self._fingerprinter.fingerprint(request).hex()
        data = {
            'status': response.status,
            'url': response.url,
            'headers': dict(response.headers),
            'body': response.body,
        }
        data_raw = pickle.dumps(data, protocol=4)
        data_raw_compressed = gzip.compress(data_raw)
        item = {
            '_id': id,
            'data': data_raw_compressed,
            'created_at': datetime.datetime.now(tz=datetime.timezone.utc),
        }
        self._col.update_one({
            '_id': id
        }, {
            '$set': item
        }, upsert=True)
