#  Copyright 2015 Palo Alto Networks, Inc
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.

"""
Table implementation based on LevelDB (https://github.com/google/leveldb).
This is a sort of poor, lazy man implementation of IndexedDB schema.

**KEYS**

Numbers are 8-bit unsigned.

- Schema Version: (0)
- Index Last Global Id: (0,1, <indexnum>)
- Last Update Key: (0,2)
- Number of Indicators: (0,3)
- Table Last Global ID: (0,4)
- Custom Metadata: (0,5)
- Indicator Version: (1,0,<indicator>)
- Indicator: (1,1,<indicator>)

**INDICATORS**

Each indicators has 2 entries associated in the DB: a version and a value.

The version number is used to track indicator existance and versioning.
When an indicator value is updated, its version number is incremented.
The version number is a 64-bit LSB unsigned int.

The value of an indicator is a 64-bit unsigned int LSB followed by a dump of
a dictionary of attributes in JSON format.

To iterate over all the indicators versions iterate from key (1,0) to key
(1,1) excluded.

NULL indicators are not allowed.

**INDEXES**

Indicators are stored in alphabetical order. Indexes are secondary indexes
on indicators attributes.

Each index has an associated id in the range 0 - 255. The attribute associated
to the index is stored at (0,1,<index id>), if the key does not exist the
index does not exist.

There is also a Last Global Id per index, used to index indicators with the
same attribute value. Each time a new indicator is added to the index, the
Last Global Id is incremented. The Last Global Id of an index is stored at
(2,<index id>,0) as a 64-bit LSB unsigned int.

Each entry in the index is stored with a key
(2,<index id>,0xF0,<encoded value>,<last global id>) and value
(<version>,<indicator>). <encoded value> depends on the type of attribute.

When iterating over an index, the value of an index entry is loaded and if
the version does not match with current indicator version the index entry is
deleted. This permits a sort of lazy garbage collection.

To retrieve all the indicators with a specific attribute value just iterate
over the keys (2,<index id>,0xF0,<encoded value>) and
(2,<index id>,0xF0,<encoded value>,0xFF..FF)
"""

import os
import plyvel
import struct
import ujson
import time
import logging
import shutil
import gevent


SCHEMAVERSION_KEY = struct.pack("B", 0)
START_INDEX_KEY = struct.pack("BBB", 0, 1, 0)
END_INDEX_KEY = struct.pack("BBB", 0, 1, 0xFF)
LAST_UPDATE_KEY = struct.pack("BB", 0, 2)
NUM_INDICATORS_KEY = struct.pack("BB", 0, 3)
TABLE_LAST_GLOBAL_ID = struct.pack("BB", 0, 4)
CUSTOM_METADATA = struct.pack("BB", 0, 5)

LOG = logging.getLogger(__name__)


class InvalidTableException(Exception):
    pass


class Table(object):
    def __init__(self, name, truncate=False, bloom_filter_bits=0):
        if truncate:
            try:
                shutil.rmtree(name)
            except:
                pass

        self.db = None
        self._compact_glet = None

        self.db = plyvel.DB(
            name,
            create_if_missing=True,
            bloom_filter_bits=bloom_filter_bits
        )
        self._read_metadata()

        self.compact_interval = int(os.environ.get('MM_TABLE_COMPACT_INTERVAL', 3600 * 6))
        self.compact_delay = int(os.environ.get('MM_TABLE_COMPACT_DELAY', 3600))
        self._compact_glet = gevent.spawn(self._compact_loop)

    def _init_db(self):
        self.last_update = 0
        self.indexes = {}
        self.num_indicators = 0
        self.last_global_id = 0

        batch = self.db.write_batch()
        batch.put(SCHEMAVERSION_KEY, struct.pack("B", 1))
        batch.put(LAST_UPDATE_KEY, struct.pack(">Q", self.last_update))
        batch.put(NUM_INDICATORS_KEY, struct.pack(">Q", self.num_indicators))
        batch.put(TABLE_LAST_GLOBAL_ID, struct.pack(">Q", self.last_global_id))
        batch.write()

    def _read_metadata(self):
        sv = self._get(SCHEMAVERSION_KEY)
        if sv is None:
            return self._init_db()
        sv = struct.unpack("B", sv)[0]
        if sv == 0:
            # add table last global id
            self._upgrade_from_s0()
        elif sv == 1:
            pass
        else:
            raise InvalidTableException("Schema version not supported")

        self.indexes = {}
        ri = self.db.iterator(
            start=START_INDEX_KEY,
            stop=END_INDEX_KEY
        )
        with ri:
            for k, v in ri:
                _, _, indexid = struct.unpack("BBB", k)
                if v in self.indexes:
                    raise InvalidTableException("2 indexes with the same name")
                self.indexes[v] = {
                    'id': indexid,
                    'last_global_id': 0
                }
        for i in self.indexes:
            lgi = self._get(self._last_global_id_key(self.indexes[i]['id']))
            if lgi is not None:
                self.indexes[i]['last_global_id'] = struct.unpack(">Q", lgi)[0]
            else:
                self.indexes[i]['last_global_id'] = -1

        t = self._get(LAST_UPDATE_KEY)
        if t is None:
            raise InvalidTableException("LAST_UPDATE_KEY not found")
        self.last_update = struct.unpack(">Q", t)[0]

        t = self._get(NUM_INDICATORS_KEY)
        if t is None:
            raise InvalidTableException("NUM_INDICATORS_KEY not found")
        self.num_indicators = struct.unpack(">Q", t)[0]

        t = self._get(TABLE_LAST_GLOBAL_ID)
        if t is None:
            raise InvalidTableException("TABLE_LAST_GLOBAL_ID not found")
        self.last_global_id = struct.unpack(">Q", t)[0]

    def _get(self, key):
        try:
            result = self.db.get(key)
        except KeyError:
            return None

        return result

    def __del__(self):
        self.close()

    def get_custom_metadata(self):
        cmetadata = self._get(CUSTOM_METADATA)
        if cmetadata is None:
            return None
        return ujson.loads(cmetadata)

    def set_custom_metadata(self, metadata=None):
        if metadata is None:
            self.db.delete(CUSTOM_METADATA)
            return

        cmetadata = ujson.dumps(metadata)
        self.db.put(CUSTOM_METADATA, cmetadata)

    def close(self):
        if self.db is not None:
            self.db.close()

        if self._compact_glet is not None:
            self._compact_glet.kill()

        self.db = None
        self._compact_glet = None

    def exists(self, key):
        if type(key) == unicode:
            key = key.encode('utf8')

        ikeyv = self._indicator_key_version(key)
        return (self._get(ikeyv) is not None)

    def get(self, key):
        if type(key) == unicode:
            key = key.encode('utf8')

        ikey = self._indicator_key(key)
        value = self._get(ikey)
        if value is None:
            return None

        # skip version
        return ujson.loads(value[8:])

    def delete(self, key):
        if type(key) == unicode:
            key = key.encode('utf8')

        ikey = self._indicator_key(key)
        ikeyv = self._indicator_key_version(key)

        if self._get(ikeyv) is None:
            return

        batch = self.db.write_batch()
        batch.delete(ikey)
        batch.delete(ikeyv)
        self.num_indicators -= 1
        batch.put(NUM_INDICATORS_KEY, struct.pack(">Q", self.num_indicators))
        batch.write()

    def _indicator_key(self, key):
        return struct.pack("BB", 1, 1) + key

    def _indicator_key_version(self, key):
        return struct.pack("BB", 1, 0) + key

    def _index_key(self, idxid, value, lastidxid=None):
        key = struct.pack("BBB", 2, idxid, 0xF0)

        if type(value) == unicode:
            value = value.encode('utf8')

        if type(value) == str:
            key += struct.pack(">BL", 0x0, len(value))+value
        elif type(value) == int or type(value) == long:
            key += struct.pack(">BQ", 0x1, value)
        else:
            raise ValueError("Unhandled value type: %s" % type(value))

        if lastidxid is not None:
            key += struct.pack(">Q", lastidxid)

        return key

    def _last_global_id_key(self, idxid):
        return struct.pack("BBB", 2, idxid, 0)

    def create_index(self, attribute):
        if attribute in self.indexes:
            return

        if len(self.indexes) == 0:
            idxid = 0
        else:
            idxid = max([i['id'] for i in self.indexes.values()])+1

        self.indexes[attribute] = {
            'id': idxid,
            'last_global_id': -1
        }

        batch = self.db.write_batch()
        batch.put(struct.pack("BBB", 0, 1, idxid), attribute)
        batch.write()

    def put(self, key, value):
        if type(key) == unicode:
            key = key.encode('utf8')

        if type(value) != dict:
            raise ValueError()

        ikey = self._indicator_key(key)
        ikeyv = self._indicator_key_version(key)

        exists = self._get(ikeyv)
        self.last_global_id += 1
        cversion = self.last_global_id

        now = time.time()
        self.last_update = now

        batch = self.db.write_batch()
        batch.put(ikey, struct.pack(">Q", cversion)+ujson.dumps(value))
        batch.put(ikeyv, struct.pack(">Q", cversion))
        batch.put(LAST_UPDATE_KEY, struct.pack(">Q", self.last_update))
        batch.put(TABLE_LAST_GLOBAL_ID, struct.pack(">Q", self.last_global_id))

        if exists is None:
            self.num_indicators += 1
            batch.put(
                NUM_INDICATORS_KEY,
                struct.pack(">Q", self.num_indicators)
            )

        for iattr, index in self.indexes.iteritems():
            v = value.get(iattr, None)
            if v is None:
                continue

            index['last_global_id'] += 1

            idxkey = self._index_key(index['id'], v, index['last_global_id'])
            batch.put(idxkey, struct.pack(">Q", cversion) + key)

            batch.put(
                self._last_global_id_key(index['id']),
                struct.pack(">Q", index['last_global_id'])
            )

        batch.write()

    def query(self, index=None, from_key=None, to_key=None,
              include_value=False, include_stop=True, include_start=True,
              reverse=False):
        if type(from_key) is unicode:
            from_key = from_key.encode('ascii', 'replace')
        if type(to_key) is unicode:
            to_key = to_key.encode('ascii', 'replace')

        if index is None:
            return self._query_by_indicator(
                from_key=from_key,
                to_key=to_key,
                include_value=include_value,
                include_stop=include_stop,
                include_start=include_start,
                reverse=reverse
            )
        return self._query_by_index(
            index,
            from_key=from_key,
            to_key=to_key,
            include_value=include_value,
            include_stop=include_stop,
            include_start=include_start,
            reverse=reverse
        )

    def _query_by_indicator(self, from_key=None, to_key=None,
                            include_value=False, include_stop=True,
                            include_start=True, reverse=False):
        if from_key is None:
            from_key = struct.pack("BB", 1, 1)
            include_stop = False
        else:
            from_key = self._indicator_key(from_key)

        if to_key is None:
            to_key = struct.pack("BB", 1, 2)
            include_start = False
        else:
            to_key = self._indicator_key(to_key)

        ri = self.db.iterator(
            start=from_key,
            stop=to_key,
            include_stop=include_stop,
            include_start=include_start,
            reverse=reverse,
            include_value=False
        )
        with ri:
            for ekey in ri:
                ekey = ekey[2:]
                if include_value:
                    yield ekey.decode('utf8', 'ignore'), self.get(ekey)
                else:
                    yield ekey.decode('utf8', 'ignore')

    def _query_by_index(self, index, from_key=None, to_key=None,
                        include_value=False, include_stop=True,
                        include_start=True, reverse=False):
        if index not in self.indexes:
            raise ValueError()

        idxid = self.indexes[index]['id']

        if from_key is None:
            from_key = struct.pack("BBB", 2, idxid, 0xF0)
            include_start = False
        else:
            from_key = self._index_key(idxid, from_key)

        if to_key is None:
            to_key = struct.pack("BBB", 2, idxid, 0xF1)
            include_stop = False
        else:
            to_key = self._index_key(
                idxid,
                to_key,
                lastidxid=0xFFFFFFFFFFFFFFFF
            )

        ldeleted = 0
        ri = self.db.iterator(
            start=from_key,
            stop=to_key,
            include_value=True,
            include_start=include_start,
            include_stop=include_stop,
            reverse=reverse
        )
        with ri:
            for ikey, ekey in ri:
                iversion = struct.unpack(">Q", ekey[:8])[0]
                ekey = ekey[8:]

                evalue = self._get(self._indicator_key_version(ekey))
                if evalue is None:
                    # LOG.debug("Key does not exist")
                    # key does not exist
                    self.db.delete(ikey)
                    ldeleted += 1
                    continue

                cversion = struct.unpack(">Q", evalue)[0]
                if iversion != cversion:
                    # index value is old
                    # LOG.debug("Version mismatch")
                    self.db.delete(ikey)
                    ldeleted += 1
                    continue

                if include_value:
                    yield ekey.decode('utf8', 'ignore'), self.get(ekey)
                else:
                    yield ekey.decode('utf8', 'ignore')

        LOG.info('Deleted in scan of {}: {}'.format(index, ldeleted))

    def _compact_loop(self):
        gevent.sleep(self.compact_delay)

        while True:
            try:
                gevent.idle()

                counter = 0
                for idx in self.indexes.keys():
                    for i in self.query(index=idx, include_value=False):
                        if counter % 512 == 0:
                            gevent.sleep(0.001)  # yield to other greenlets
                        counter += 1

            except gevent.GreenletExit:
                break
            except:
                LOG.exception('Exception in _compact_loop')

            try:
                gevent.sleep(self.compact_interval)

            except gevent.GreenletExit:
                break

    def _upgrade_from_s0(self):
        LOG.info('Upgrading from schema version 0 to schema version 1')

        LOG.info('Loading indexes...')
        indexes = {}
        ri = self.db.iterator(
            start=START_INDEX_KEY,
            stop=END_INDEX_KEY
        )
        with ri:
            for k, v in ri:
                _, _, indexid = struct.unpack("BBB", k)
                if v in indexes:
                    raise InvalidTableException("2 indexes with the same name")
                indexes[v] = {
                    'id': indexid,
                    'last_global_id': 0
                }
        for i in indexes:
            lgi = self._get(self._last_global_id_key(indexes[i]['id']))
            if lgi is not None:
                indexes[i]['last_global_id'] = struct.unpack(">Q", lgi)[0]
            else:
                indexes[i]['last_global_id'] = -1

        LOG.info('Scanning indexes...')
        last_global_id = 0
        for i, idata in indexes.iteritems():
            from_key = struct.pack("BBB", 2, idata['id'], 0xF0)
            include_start = False
            to_key = struct.pack("BBB", 2, idata['id'], 0xF1)
            include_stop = False

            ri = self.db.iterator(
                start=from_key,
                stop=to_key,
                include_value=True,
                include_start=include_start,
                include_stop=include_stop,
                reverse=False
            )
            with ri:
                for ikey, ekey in ri:
                    iversion = struct.unpack(">Q", ekey[:8])[0]
                    if iversion > last_global_id:
                        last_global_id = iversion+1

        LOG.info('Last global id: {}'.format(last_global_id))
        batch = self.db.write_batch()
        batch.put(SCHEMAVERSION_KEY, struct.pack("B", 1))
        batch.put(TABLE_LAST_GLOBAL_ID, struct.pack(">Q", last_global_id))
        batch.write()
