import sqlite3
import pickle

class RosterDB(object):
    """
    Implements a roster persistence object (expected by sleekxmpp roster)
    with a sqlite3 backend
    """
    def __init__(self, name, table):
        self.db = sqlite3.connect(name, check_same_thread=False)
        self.table = table

        try:
            self.db.execute('create table {} (o varchar, k varchar,  v blob)'.format(self.table))
        except sqlite3.OperationalError as e:
            pass

    def entries(self, key, dblist={}):
        if key is None:
            r = self.db.execute('select distinct(o) from {}'.format(self.table)).fetchall()
            r = map(lambda x : x[0], r)
            print("DB ENTRIES AT TOP LEVEL", r)
            return r

        r = self.db.execute('select k from {} where o = ?'.format(self.table), (key,)).fetchall()
        r = map(lambda x: x[0], r)
        print("DB ENTRIES FOR KEY ", key, r)
        return r

    def load(self, owner, key, state):
        print("TYPE OF OWNER", type(owner)," TYPE OF KEY ", type(key))
        r = self.db.execute('select v from {} where o = ? and k = ?'.format(self.table), (owner, key)).fetchone()
        print("DB LOAD ", owner, key, " = ", r);
        if r is None:
            return r
        return pickle.loads(r[0])

    def save(self, owner, key, item, state):
        print("DB SAVE ", owner, key, " = ", item)
        try:
            self.db.execute('delete from {} where o = ? and k = ?'.format(self.table), (owner, key))
        except sqlite3.OperationalError as e:
            pass

        p = pickle.dumps(item)
        self.db.execute('insert into {} values (?, ?, ?)'.format(self.table),
                        (owner, key, p))
        self.db.commit()