# -*- coding: utf-8 -*-
import logging

import pymongo
from pymongo.errors import ConnectionFailure

from .abstractdb import AbstractDB


class MongoDB(AbstractDB):

    def __init__(self, ip="127.0.0.1", port=27017, dbname="pastepwn", collectionname="pastes"):
        super().__init__()
        self.logger = logging.getLogger(__name__)
        self.logger.debug("Initializing MongoDB - {0}:{1}".format(ip, port))
        self.db = pymongo.MongoClient(ip, port, serverSelectionTimeoutMS=5000)

        try:
            self.db.admin.command("ismaster")
        except ConnectionFailure as e:
            self.logger.error(e)
            raise e

        self.logger.debug("Connected to database!")

        self.db = self.db[dbname]
        self.collection = self.db[collectionname]
        self.collection.create_index([('key', pymongo.ASCENDING)], unique=True)

    def _insert_data(self, data):
        self.collection.update_one({'key': data['key']}, {'$set': data}, upsert=True)

    def _get_data(self, key, value):
        return self.collection.find({key: value})

    # TODO def update_data(self, )

    # TODO def delete_data(self,)

    def count(self, key, value):
        return self.collection.find({key: value}).count()

    def count_all(self):
        return self.collection.count()

    def store(self, paste):
        self.logger.debug("Storing paste {0}".format(paste.key))

        try:
            self._insert_data(paste.to_dict())
        except pymongo.errors.DuplicateKeyError:
            self.logger.debug("Duplicate key '{0}' - Not storing paste".format(paste.key))

    def get(self, key):
        return self._get_data("key", key)