# coding: utf-8
from __future__ import unicode_literals

from datetime import datetime

import pymongo

from emission.core.storage import BaseStorage


class ClientMongoDBStorage(BaseStorage):
    """MongoDB storage

    Поля документа:
        data
        creation_time
        msg_id
        master_name
    """

    def __init__(self, collection, master_name):
        self.collection = collection
        self.master_name = master_name

    def get_query(self):
        return {'master_name': self.master_name}

    def insert(self, msg_id, data, action):
        entry = self.get_one(msg_id)

        if not entry:
            entry = {
                'msg_id': msg_id,
                'data': data,
                'action': action,
                'creation_time': datetime.now(),
                'master_name': self.master_name,
            }

            self.collection.insert_one(entry)
        else:
            entry['data'] = data
            self.collection.replace_one({'msg_id': msg_id}, entry)

        return entry

    def get_one(self, msg_id):
        query = self.get_query()
        query.update({'msg_id': msg_id})

        return self.collection.find_one(query)

    def get_slice(self, start, stop=None, max_rows=None):
        query = self.get_query()
        query.update({'msg_id': {'$gte': start}})

        if stop is not None:
            query['msg_id'].update({'$lte': stop})

        kwargs = {}
        if max_rows:
            kwargs['limit'] = max_rows

        return self.collection.find(
            query,
            sort=[('msg_id', pymongo.ASCENDING)],
            **kwargs
        )

    def delete_one(self, msg_id):
        query = self.get_query()
        query.update({'msg_id': msg_id})

        return self.collection.delete_one(query)
