
import pickle
import logging
import traceback


class PickleDatabase(object):
    """ This object is a wrapper to access and handle our
    local pickle file. We'll use this pickle file as a sort of
    database for storing schedules and other info for the agent."""
    def __init__(self, file_path):
        self.file_path = file_path
        try:
            self.data = pickle.load(open(file_path, 'rb'))
        except Exception:
            # If something went wrong, the data is either corrupted or missing.
            log = logging.getLogger(self.__class__.__name__)
            log.warn('Unable to open database {}, creating an empty one'.format(self.file_path))
            self.data = {}
            open(self.file_path, 'wb')  # Create the DB file if it doesn't exist
            log.info('Database created: %s' % self.file_path)

    def save(self):
        # Dump current contents of our dict into the pickle file
        pickle.dump(self.data, open(self.file_path, 'wb'))

    def __getitem__(self, key):
        return self.data[key]

    def __setitem__(self, key, value):
        self.data[key] = value
        # Save our newly updated dictionary
        self.save()

    def __repr__(self):
        return str(self.data)

    def __contains__(self, key):
        return key in self.data