#!/usr/bin/env python

import os
import sqlite3
import binascii
import itertools
import logging
log = logging.getLogger('gglsbl3')


class StorageBase(object):
    @staticmethod
    def compress_ranges(nums):
        """Compress consequtive ranges in sequence of numbers

        E.g. [1,2,3,4,7] -> '1-4,7'
        """
        if not nums:
            return None
        nums.sort()
        buf = []
        buf.append(nums[0])
        for i in range(1, len(nums)):
            if nums[i-1] == nums[i]:
                pass
            elif nums[i] - nums[i-1] == 1:
                if buf[-1] is not None:
                    buf.append(None)
            else:
                if buf[-1] is None:
                    buf.append(nums[i-1])
                buf.append(nums[i])
        if buf[-1] is None:
            buf.append(nums[-1])
        output = ','.join([str(i) for i in buf])
        output = output.replace(',None,', '-')
        return output

    @staticmethod
    def expand_ranges(list_of_ranges):
        """Do the opposite of compress_ranges()

        E.g. ['1-4,7', '9-11'] -> [1,2,3,4,7,9,10,11]
        """
        nums = []
        for ranges in list_of_ranges:
            for chunk_range in ranges.strip().split(','):
                if type(chunk_range) is int:
                    nums.append(chunk_range)
                elif chunk_range.isdigit():
                    nums.append(int(chunk_range))
                else:
                    try:
                        r1, r2 = chunk_range.split('-')
                        r1 = int(r1)
                        r2 = int(r2) + 1
                        nums.extend(range(r1, r2))
                    except ValueError as e:
                        log.error('Failed to parse chunk range %s (%s)', chunk_range, e)
                        raise
        return nums


class SqliteStorage(StorageBase):
    """Storage abstraction for local GSB cache"""
    def __init__(self, db_path):
        do_init_db = not os.path.isfile(db_path)
        log.info('Opening SQLite DB %s' % db_path)
        self.db = sqlite3.connect(db_path)
        self.dbc = self.db.cursor()
        if do_init_db:
            log.info('SQLite DB does not exist, initializing')
            self.init_db()

    def init_db(self):
        self.dbc.execute(
        """CREATE TABLE chunk (
            chunk_number integer NOT NULL,
            timestamp timestamp DEFAULT current_timestamp,
            list_name character varying(127) NOT NULL,
            chunk_type TEXT CHECK( chunk_type IN ('add','sub')) NOT NULL,
            PRIMARY KEY (chunk_number, list_name, chunk_type)
            )"""
        )
        self.dbc.execute(
        """CREATE TABLE full_hash (
            value BLOB NOT NULL,
            list_name character varying(127) NOT NULL,
            metadata TINYINT,
            downloaded_at timestamp DEFAULT current_timestamp,
            expires_at timestamp without time zone NOT NULL,
            PRIMARY KEY (value)
            )"""
        )

        self.dbc.execute(
        """CREATE TABLE hash_prefix (
            value BLOB NOT NULL,
            chunk_number integer NOT NULL,
            timestamp timestamp without time zone DEFAULT current_timestamp,
            list_name character varying(127) NOT NULL,
            chunk_type TEXT CHECK( chunk_type IN ('add','sub')) NOT NULL,
            full_hash_expires_at timestamp NOT NULL DEFAULT current_timestamp,
            PRIMARY KEY (value, chunk_number, list_name, chunk_type),
            FOREIGN KEY(chunk_number, list_name, chunk_type)
                REFERENCES chunk(chunk_number, list_name, chunk_type)
                ON DELETE CASCADE
            )"""
        )
        self.db.commit()

    def chunk_exists(self, chunk):
        "Check if given chunk records already exist in the database"
        q = 'SELECT COUNT(*) FROM chunk WHERE chunk_number=? AND \
            chunk_type=? AND list_name=?'
        self.dbc.execute(q, [chunk.chunk_number, chunk.chunk_type, chunk.list_name])
        if self.dbc.fetchall()[0][0] > 0:
            return True
        return False

    def store_chunk(self, chunk):
        "Store chunk in the database"
        log.debug('Storing %s chunk #%s for list name %s' % (chunk.chunk_type, chunk.chunk_number, chunk.list_name))
        self.insert_chunk(chunk)
        for hash_value in chunk.hashes:
            hash_prefix = {
                'list_name': chunk.list_name,
                'chunk_number': chunk.chunk_number,
                'chunk_type': chunk.chunk_type,
                'value': sqlite3.Binary(hash_value),
            }
            self.insert_hash_prefix(hash_prefix)
        self.db.commit()

    def insert_chunk(self, chunk):
        "Insert hash prefixes from the chunk to the database"
        q = 'INSERT INTO chunk (chunk_number, list_name, chunk_type) \
            VALUES (?, ?, ?)'
        self.dbc.execute(q, [chunk.chunk_number, chunk.list_name, chunk.chunk_type])

    def insert_hash_prefix(self, hash_prefix):
        "Insert individual hash prefix to the database"
        q = 'INSERT INTO hash_prefix (value, chunk_number, list_name, chunk_type) \
            VALUES (?, ?, ?, ?)'
        params = [hash_prefix[k] for k in
                        ('value', 'chunk_number', 'list_name', 'chunk_type')]
        try:
            self.dbc.execute(q, params)
        except sqlite3.IntegrityError as e:
            log.warning("Trying to insert existing hash prefix: '%s' (%s)", hash_prefix, e)


    def store_full_hashes(self, hash_prefix, hashes):
        "Store hashes found for the given hash prefix"
        self.cleanup_expired_hashes()
        cache_lifetime = hashes['cache_lifetime']
        for hash_info, metadata_info in itertools.zip_longest(list(hashes['hashes'].items()), list(hashes['metadata'].items())):
            list_name_bytes = hash_info[0]
            #  log.debug("list name bytes is: {list_name_bytes}".format(list_name_bytes=list_name_bytes))
            list_name = list_name_bytes.decode("ascii")  # convert the bytestring listname to a 'normal' python string
            #  log.debug("list_name is {list_name}".format(list_name=list_name))
            hash_values = hash_info[1]
            metadata_values = metadata_info[1]
            log.debug("list_name: {list_name}".format(list_name=list_name))
            for hash_value, metadata_value in itertools.zip_longest(hash_values, metadata_values):
                if metadata_value is None:
                    metadata_value = 0
                log.debug("metadata value is: {metadata_value}".format(metadata_value=metadata_value))
                log.debug('storing hash: "{hash}"'.format(hash=binascii.hexlify(hash_value).decode("ascii")))
                q = "INSERT INTO full_hash (value, list_name, metadata, downloaded_at, expires_at)\
                    VALUES (?, ?, ?, current_timestamp, datetime(current_timestamp, '+%d SECONDS'))"
                q = q % cache_lifetime
                log.debug(q)
                self.dbc.execute(q, [sqlite3.Binary(hash_value), list_name, metadata_value])
        q = "UPDATE hash_prefix SET full_hash_expires_at=datetime(current_timestamp, '+%d SECONDS') \
            WHERE chunk_type='add' AND value=?"
        self.dbc.execute(q % cache_lifetime, [sqlite3.Binary(hash_prefix)])
        self.db.commit()

    def full_hash_sync_required(self, hash_prefix):
        """Check if hashes for the given hash prefix have expired

        and that prefix needs to be re-queried
        """
        q = "SELECT COUNT(*) FROM hash_prefix WHERE \
            full_hash_expires_at > current_timestamp AND chunk_type='add' AND value=?"
        self.dbc.execute(q, [sqlite3.Binary(hash_prefix)])
        c = self.dbc.fetchall()[0][0]
        return c == 0

    def lookup_full_hash(self, hash_value):
        "Query DB to see if hash is blacklisted"
        q = 'SELECT list_name, metadata FROM full_hash WHERE value=?'
        self.dbc.execute(q, [sqlite3.Binary(hash_value)])
        fetched = self.dbc.fetchall()
        log.debug("fetched from database: {fetched}".format(fetched=fetched))
        return [{"list": h[0], "metadata": h[1]} for h in fetched]

    def lookup_hash_prefix(self, hash_prefix):
        """Check if hash prefix is in the list and does not have 'sub'
        status signifying that it should be evicted from the blacklist
        """
        q = 'SELECT list_name FROM hash_prefix WHERE chunk_type=? AND value=?'
        self.dbc.execute(q, ['add', sqlite3.Binary(hash_prefix)])
        lists_add = [r[0] for r in self.dbc.fetchall()]
        if len(lists_add) == 0:
            return False
        self.dbc.execute(q, ['sub', sqlite3.Binary(hash_prefix)])
        lists_sub = [r[0] for r in self.dbc.fetchall()]
        if len(lists_sub) == 0:
            return True
        if set(lists_add) - set(lists_sub):
            return True
        return False

    def cleanup_expired_hashes(self):
        "Delete all hashes that behind their expiration date"
        q = 'DELETE FROM full_hash WHERE expires_at < current_timestamp'
        self.dbc.execute(q)
        self.db.commit()

    def del_add_chunks(self, chunk_numbers):
        "Delete records associated with 'add' chunk"
        if not chunk_numbers:
            return
        log.info('Deleting "add" chunks %s' % repr(chunk_numbers))
        for cn in self.expand_ranges(chunk_numbers):
            q = 'DELETE FROM chunk WHERE chunk_type=? AND chunk_number=?'
            self.dbc.execute(q, ['add', cn])
        self.db.commit()

    def del_sub_chunks(self, chunk_numbers):
        "Delete records associated with 'sub' chunk"
        if not chunk_numbers:
            return
        log.info('Deleting "sub" chunks %s' % repr(chunk_numbers))
        for cn in self.expand_ranges(chunk_numbers):
            q = 'DELETE FROM chunk WHERE chunk_type=? AND chunk_number=?'
            self.dbc.execute(q, ['sub', cn])
        self.db.commit()

    def get_existing_chunks(self):
        "Get the list of chunks that are available in the local cache"
        output = {}
        for chunk_type in ('add', 'sub'):
            q = "SELECT list_name, group_concat(chunk_number) FROM chunk \
                WHERE chunk_type=? GROUP BY list_name"
            self.dbc.execute(q, [chunk_type])
            for list_name, chunks in self.dbc.fetchall():
                if list_name not in output:
                    output[list_name] = {}
                chunks = [int(c) for c in chunks.split(',')]
                output[list_name][chunk_type] = self.compress_ranges(chunks)
        return output

    # FiXME: add stats about number of chunks, number of fullhashes

    def total_cleanup(self):
        "Reset local cache"
        q = 'DROP TABLE hash_prefix'
        self.dbc.execute(q)
        q = 'DROP TABLE chunk'
        self.dbc.execute(q)
        q = 'DROP TABLE full_hash'
        self.dbc.execute(q)
        self.db.commit()
        self.init_db()

    def close(self):
        self.db.close()
