"""
From https://github.com/shish/sqliteshelf (MIT License)

by default, things are stored in a "shelf" table

>>> d = SQLiteShelf("test.sdb")

you can put multiple shelves into a single SQLite database

>>> e = SQLiteShelf("test.sdb", "othertable")

both are empty to start with

>>> d
{}
>>> e
{}

adding stuff is as simple as a regular dict
>>> d['a'] = "moo"
>>> e['a'] = "moo"

regular dict actions work

>>> d['a']
'moo'
>>> e['a']
'moo'
>>> 'a' in d
True
>>> len(d)
1
>>> del d['a']
>>> 'a' in d
False
>>> len(d)
0
>>> del e['a']

objects can be stored in shelves

>> class Test:
..    def __init__(self):
..        self.foo = "bar"
..
>> t = Test()
>> d['t'] = t
>> print d['t'].foo
bar

errors are as normal for a dict

>>> d['x']
Traceback (most recent call last):
    ...
KeyError: 'x'
>>> del d['x']
Traceback (most recent call last):
    ...
KeyError: 'x'

Adding and fetching binary strings

>>> d[1] = "a\\x00b"
>>> d[1]
'a\\x00b'
"""

try:
    from UserDict import DictMixin
except ImportError:
    from collections.abc import MutableMapping as DictMixin

try:
    import cPickle as pickle
except ImportError:
    import pickle

import sys

valtype = str
if sys.version > "3":
    buffer = memoryview
    valtype = bytes

import sqlite3  # noqa: E402


class SQLiteDict(DictMixin):
    def __init__(
        self, filename=":memory:", table="shelf", flags="r", mode=None, valtype=valtype
    ):
        self.table = table
        self.valtype = valtype
        MAKE_SHELF = (
            "CREATE TABLE IF NOT EXISTS " + self.table + " (key TEXT, value BLOB)"
        )
        MAKE_INDEX = (
            "CREATE UNIQUE INDEX IF NOT EXISTS "
            + self.table
            + "_keyndx ON "
            + self.table
            + "(key)"
        )
        self.conn = sqlite3.connect(filename)
        self.conn.text_factory = str
        self.conn.execute(MAKE_SHELF)
        self.conn.execute(MAKE_INDEX)
        self.conn.commit()

    def __getitem__(self, key):
        GET_ITEM = "SELECT value FROM " + self.table + " WHERE key = ?"
        item = self.conn.execute(GET_ITEM, (key,)).fetchone()
        if item is None:
            raise KeyError(key)
        return self.valtype(item[0])

    def __setitem__(self, key, item):
        ADD_ITEM = "REPLACE INTO " + self.table + " (key, value) VALUES (?,?)"
        self.conn.execute(ADD_ITEM, (key, sqlite3.Binary(item)))
        self.conn.commit()

    def __delitem__(self, key):
        if key not in self:
            raise KeyError(key)
        DEL_ITEM = "DELETE FROM " + self.table + " WHERE key = ?"
        self.conn.execute(DEL_ITEM, (key,))
        self.conn.commit()

    def __iter__(self):
        c = self.conn.cursor()
        try:
            c.execute("SELECT key FROM " + self.table + " ORDER BY key")
            for row in c:
                yield row[0]
        finally:
            c.close()

    def keys(self):
        c = self.conn.cursor()
        try:
            c.execute("SELECT key FROM " + self.table + " ORDER BY key")
            return [row[0] for row in c]
        finally:
            c.close()

    ###################################################################
    # optional bits

    def __len__(self):
        GET_LEN = "SELECT COUNT(*) FROM " + self.table
        return self.conn.execute(GET_LEN).fetchone()[0]

    def close(self):
        if self.conn is not None:
            self.conn.commit()
            self.conn.close()
            self.conn = None

    def __del__(self):
        self.close()

    def __repr__(self):
        return repr(dict(self))


class SQLiteShelf(SQLiteDict):
    def __getitem__(self, key):
        return pickle.loads(SQLiteDict.__getitem__(self, key))

    def __setitem__(self, key, item):
        SQLiteDict.__setitem__(self, key, pickle.dumps(item))


if __name__ == "__main__":
    import doctest

    doctest.testmod()
