#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Copyright (C) 2012-2016 Xyne
#
# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# (version 2) as published by the Free Software Foundation.
#
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301, USA.

import datetime
import logging
import os.path
import sqlite3
import sys
import types



# TODO
# Find a way to get this from sqlite3 instead of hard-coding it.
SQLITE_MAX_EXPR_DEPTH = 1000-1



# At the time of writing there seems to be no better way to detect this error.
def sqlite_no_such_table_error(e):
  '''
  Return True if the error is the sqlite3.OperationError "no such table".
  '''
  return (str(e)[:15] == "no such table: ")



def stabilize(xs): # HARHAR
  '''
  Ensure that the given iterable is stable for iteration, i.e. that it is not
  a generator (single iteration) and that each iteration traverses the elements
  in the same order.
  '''
#   if isinstance(xs, types.GeneratorType):
  if not (isinstance(xs, list) or isinstance(xs, tuple)):
    return list(xs)
  else:
    return xs



class MDBError(Exception):
  '''
  Exception raised by fatal MemoizeDB errors.

  conn: The database connection.

  msg: The error message.

  error: The associated error, usually a sqlite3 Exception.
  '''

  def __init__(self, conn, msg, error=None):
    self.conn = conn
    self.msg = msg
    self.error = error

  def __str__(self):
    cursor = self.conn.cursor()
    cursor.execute("PRAGMA database_list;")
    db_names = ' or '.join('"{}"'.format(r[2]) for r in cursor.fetchall() if r[2])
    if db_names:
      display_msg = '{} ({})'.format(self.__class__.__name__, db_names)
    if self.error is not None:
      display_msg += ' [{}]'.format(self.error)
    return display_msg



class MemoizeDB():
  '''
  Store results in a DB for future requests.

  Results are stored in a database for future requests to increase speed and
  reduce calls to remote services. Results may be kept permanently or purged
  after a given interval.
  '''

  KEY_COLUMN = ('_key', 'TEXT')
  TIMESTAMP_COLUMN = ('_created', 'timestamp')

  def __init__(self, conn, glue, null_ttl=None):
    '''
    conn:
      An sqlite3 database connection. The connection should be created with the
      parameters "detect_types=sqlite3.PARSE_DECLTYPES|sqlite3.PARSE_COLNAMES",
      and "isolation_level=None", e.g.

          conn = sqlite3.connect(
            path,
            detect_types=sqlite3.PARSE_DECLTYPES|sqlite3.PARSE_COLNAMES,
            isolation_level=None
          )

      This may be extended in the future to support all connections implementing
      Python's DB-API.

    glue:
      A dictionary of tuples that provide the glue for retrieving the data and
      storing it in the database:

        { <table name> : (<retrieval function>, <type tuple>, <time to live>) }

      <table name>:
        The name of the database table, e.g. 'foo').

      <retrieval function>:
        The function to initially retrieve the requested data, e.g. It must
        accept an iterable of keys as its only argument and return an iterable
        over pairs of keys and row data. The row data must itself be an iterable
        that corresponds to the table data fields, or None if no data could be
        retrieved. Example:

            def retrieve_ex1(keys):
              for key in keys:
                try:
                  single_value = do_something(key)
                except Whatever:
                  yield key, None
                else:
                  # Return a tuple
                  yield key, (single_value,)

            def retrieve_ex2(keys):
              for key in keys:
                try:
                  some_iterable = do_something(key)
                except Whatever:
                  yield key, None
                else:
                  yield key, some_iterable

      <type tuple>:
        The database field names and types to use when storing the tuple
        retrieved by <retrieval function>. These will be directly used to form
        the SQL string, so make sure that they are sanitized before passing
        them. Example:

            (('foo', 'TEXT'), ('bar', 'TEXT'))

      <time to live>:
        The interval, in seconds, after which to purge the database entry. If
        None, then database entries in that table do not expire. Setting this to
        a very low value defeats the purpose of using this module.

    null_ttl:
      A time which, if given, will limit the lifetime of empty entries. This is
      useful when the data source may temporarily return empty values. The empty
      value will be cached temporarily to avoid multiple queries in a short
      interval, but will not be stored as long as other data (e.g. permanently).


    '''
    self.conn = conn
    self.glue = glue
    self.cursor = conn.cursor()
    self.null_ttl = null_ttl



  def raise_mdb_error(self, *args, **kwargs):
    '''
    Raise a MDBError.
    '''
    raise MDBError(self.conn, *args, **kwargs)



  ##############################################################################
  # Database Methods
  ##############################################################################

  def db_initialize(self):
    '''
    (Re-)initialize the database by creating missing tables and cleaning old
    entries.
    '''
    for t in self.glue:
      self.db_create(t)
    self.db_clean()



  def db_create(self, table):
    '''
    Create a database table.
    '''
    c = self.cursor
    try:
      func, types, ttl = self.glue[table]
      cols = ','.join('"{}" {}'.format(*tup) for tup in types + (self.TIMESTAMP_COLUMN,))
      c.execute('CREATE TABLE IF NOT EXISTS "{}" ("{}" {} PRIMARY KEY, {})'.format(table, self.KEY_COLUMN[0], self.KEY_COLUMN[1], cols))
      logging.debug('created table "{}"'.format(table))
    except sqlite3.OperationalError as e:
      self.raise_mdb_error('table creation failed', error=e)
    except KeyError:
      self.raise_mdb_error('invalid table "{}"'.format(table))
    self.conn.commit()


  def db_insert(self, table, key, values):
    '''
    Insert a row for the given key into the database table.
    '''
    if values is None:
      if self.null_ttl is None:
#         self.db_delete(self, table, key)
        return
      else:
        values = (None,) * n
    c = self.cursor
    n = len(self.glue[table][1])
    # One for each value, plus key and timestamp.
    cols = '?, ' * n + '?, ?'
    query = 'REPLACE INTO "{}" VALUES ({})'.format(table, cols)
    now = datetime.datetime.utcnow()
    args = (key,) + tuple(values) + (now,)
    try:
      c.execute(query, args)
    except sqlite3.OperationalError as e:
      if sqlite_no_such_table_error(e):
        self.db_create(table)
        try:
          c.execute(query, args)
        except sqlite3.OperationalError as e:
          self.raise_mdb_error('row insertion(s) failed', error=e)
      else:
        self.raise_mdb_error('row insertion(s) failed', error=e)
    self.conn.commit()


  def where_clause(self, keys):
    '''
    Return a WHERE clause for the given keys.
    '''
    return ' OR '.join('{}=?'.format(self.KEY_COLUMN[0]) for k in keys)


  def db_delete(self, table, key):
    '''
    Delete a row matching the given key from the database.
    '''
    self.db_delete_many(table, (key,))


  def db_delete_many(self, table, keys):
    '''
    Delete rows matching the given keys from the database.
    '''
    keys = stabilize(keys)
    if not keys:
      return

    # Break up long queries to avoid exceeding the maximum depth of the
    # expression tree.
    # https://www.sqlite.org/limits.html
    elif len(keys) > SQLITE_MAX_EXPR_DEPTH:
      i = 0
      if not isinstance(keys, list):
        keys = list(keys)
      ks = keys[i:i+SQLITE_MAX_EXPR_DEPTH]
      while ks:
        self.db_delete_many(table, ks)
        i += SQLITE_MAX_EXPR_DEPTH
        ks = keys[i:i+SQLITE_MAX_EXPR_DEPTH]
      return

    where_clause = self.where_clause(keys)
    query = 'DELETE FROM "{}" WHERE {}'.format(table, where_clause)
    c = self.cursor
    try:
      c.execute(query, keys)
    except sqlite3.OperationalError as e:
      if sqlite_no_such_table_error(e):
        return
      else:
        self.raise_mdb_error('deletion failed', error=e)
    self.conn.commit()


  def db_select(self, table, key, cols=None):
    '''
    Select the values associated with the key in the database table.

    Return the associated row.
    '''
    for row in self.db_select_many(table, (key,), cols=cols):
      return row


  def db_select_many(self, table, keys, cols=None):
    '''
    Select the values associated with the keys in the database table.

    Return the associated rows. The rows are returned in order. If no row was
    found, None will be returned for that key.
    '''
    keys = stabilize(keys)
    if not keys:
      return

    # Break up long queries to avoid exceeding the maximum depth of the
    # expression tree.
    # https://www.sqlite.org/limits.html
    elif len(keys) > SQLITE_MAX_EXPR_DEPTH:
      i = 0
      if not isinstance(keys, list):
        keys = list(keys)
      ks = keys[i:i+SQLITE_MAX_EXPR_DEPTH]
      while ks:
        for x in self.db_select_many(table, ks, cols=cols):
          yield x
        i += SQLITE_MAX_EXPR_DEPTH
        ks = keys[i:i+SQLITE_MAX_EXPR_DEPTH]
      return

    where_clause = self.where_clause(keys)

    added_key_column = False
    key_column = 0
    if cols is None:
      query = 'SELECT * FROM "{}" WHERE {}'.format(table, where_clause)
    else:
      cols = tuple(cols)
      try:
        key_column = cols.index(self.KEY_COLUMN[0])
      except ValueError:
        cols = (self.KEY_COLUMN[0],) + cols
        added_key_column = True
      cs = ','.join('"{}"'.format(c.replace('"', '""')) for c in cols)
      query = 'SELECT {} FROM "{}" WHERE {}'.format(cs, table, where_clause)

    c = self.cursor
    try:
      backlog = dict()
      unread_results = True
      results = c.execute(query, keys)
      for key in keys:
        try:
          # Check if the matching row has already been read.
          yield backlog[key]
        except KeyError:
          if unread_results:
            # Read rows until we find the matching one.
            for rawrow in results:
              if added_key_column:
                row = rawrow[1:]
              else:
                row = rawrow
              # Yield the associated row.
              if rawrow[key_column] == key:
                yield row
                break
              # Add the row to the backlog if it does not match this key.
              else:
                backlog[rawrow[key_column]] = row
            else:
              unread_results = False
              # No entry for this key.
              yield None
          else:
            # No entry for this key.
            yield None
        else:
          # Delete the yielded row from the backlog.
          del backlog[key]
          continue

    except sqlite3.OperationalError as e:
      self.raise_mdb_error('query failed', error=e)



  # Clean up old tables in the database.
  def db_clean(self, wipe=False):
    '''
    Clean up the database.

    This will drop tables that no longer match the glue. It will also purge all
    outdated entries.
    '''
    try:
      c = self.cursor
      c.execute('SELECT tbl_name FROM "sqlite_master" WHERE type = "table"')
      for row in c.fetchall():
        table, = row
        # Drop if it it's no longer in the glue.
        if not table in self.glue:
          c.execute('DROP TABLE "{}"'.format(table))
          logging.debug('dropped table "{}"'.format(table))
        # Otherwise check that the columns match.
        else:
          expected_cols = (self.KEY_COLUMN,) + self.glue[table][1] + (self.TIMESTAMP_COLUMN,)
          cols = c.execute('PRAGMA table_info("{}")'.format(table)).fetchall()
          if len(expected_cols) == len(cols):
            for (name, typ), meta in zip(expected_cols, cols):
              if name != meta[1] or typ != meta[2]:
                c.execute('DROP TABLE "{}"'.format(table))
                logging.debug('dropped table "{}"'.format(table))
                break
            # Finally, delete expired entries.
            else:
              now = datetime.datetime.utcnow()
              if wipe or self.glue[table][2]:
                if wipe:
                  query = 'DELETE FROM "{}"'.format(table)
                  deleted = c.execute(query).rowcount
                else:
                  max_ttl = datetime.timedelta(seconds=self.glue[table][2])
                  arg = now - max_ttl
                  query = 'DELETE FROM "{}" WHERE "{}"<?'.format(table, self.TIMESTAMP_COLUMN[0])
                  deleted = c.execute(query, (arg,)).rowcount
                if deleted > 0:
                  logging.debug('deleted {:d} row(s) from table "{}"'.format(deleted, table))

              # Delete empty entries that have timed out.
              if not wipe and self.null_ttl is not None:
                cond = ' AND '.join('"{}" IS NULL'.format(f[0]) for f in self.glue[table][1])
                query = 'DELETE FROM "{}" WHERE {} AND "{}"<?'.format(
                  table, cond, self.TIMESTAMP_COLUMN[0]
                )
                max_ttl = datetime.timedelta(seconds=self.null_ttl)
                arg = now - max_ttl
                deleted = c.execute(query, (arg,)).rowcount
                if deleted > 0:
                  logging.debug('deleted {:d} row(s) from table"{}"'.format(deleted, table))
          else:
            c.execute('DROP TABLE "{}"'.format(table))
            logging.debug('dropped table "{}"'.format(table))
      c.execute('VACUUM')
      self.conn.commit()
    except sqlite3.OperationalError as e:
      self.raise_mdb_error('cleaning failed', error=e)




  ##############################################################################
  # Accessibility Methods
  ##############################################################################

  def get_cached_data(self, table, key):
    '''
    If the key is cached and still valid, return the values, otherwise None
    '''
    for values in self.get_cached_data_many(table, (key,)):
      return values


  def get_cached_data_many(self, table, keys, list_keys=False, cols=None):
    '''
    The same as get_cached_data for each key, in order. If list_keys is True
    then the key itself is returned instead of the values. This can be used to
    check which cached queries are still valid.
    '''
    keys = stabilize(keys)
    timestamp_index = -1
    appended_timestamp = True
    if list_keys:
      cols = (self.TIMESTAMP_COLUMN[0],)
    elif cols:
      cols = stabilize(cols)
      try:
        timestamp_index = cols.index(self.TIMESTAMP_COLUMN[0])
        appended_timestamp = False
      except ValueError:
        cols = tuple(cols) + (self.TIMESTAMP_COLUMN[0],)
    else:
      cols = None

    max_ttl = self.glue[table][2]
    # Values will be flanked by the key and timestamp fields.
    for key, values in zip(keys, self.db_select_many(table, keys, cols=cols)):
      if values is not None:
        # Check that it's still valid.
        if max_ttl is not None \
        or self.null_ttl is not None:
          timestamp = values[timestamp_index]
          now = datetime.datetime.utcnow()
          age = (now - timestamp).total_seconds()
          if (
            max_ttl is not None \
            and age > max_ttl
          ):
            yield None
            continue

        if list_keys:
          yield key
        elif cols is None:
          yield values[1:-1]
        elif appended_timestamp:
          yield values[:-1]
        else:
          yield values

      else:
        yield None


  def get(self, table, key):
    '''
    Get the values associated with the given key.
    '''
    values = self.get_cached_data(table, key)
    if values is None:
      func = self.glue[table][0]
      for key, values in func((key,)):
        break
      if values is not None:
        # In case it's a generator.
        values = tuple(values)
      self.db_insert(table, key, values)
    return values


  def get_many(self, table, keys):
    '''
    Iterate over the values corresponding to the given keys.
    '''
    keys = stabilize(keys)
    key_set = set(keys)
    cached_key_set = set(
      key for key in self.get_cached_data_many(
        table, keys, list_keys=True
      ) if key is not None
    )

    uncached_keys = list(key_set - cached_key_set)

    if uncached_keys:
      func = self.glue[table][0]
      for key, values in func(uncached_keys):
        self.db_insert(table, key, values)
    cols = tuple(x[0] for x in self.glue[table][1])
    for vs in self.db_select_many(table, keys, cols=cols):
      yield vs


  def get_one(self, table, key, n=0):
    '''
    Deprecated function. Use get_nth_field instead.
    '''
    return self.get_nth_field(table, key, n=n)


  def get_nth_field(self, table, key, n=0):
    '''
    Retrieve the nth item of the returned tuple. This is useful when the
    retrieval function only returns a single value.
    '''
    vs = self.get(table, key)
    if vs is None:
      return None
    else:
      return vs[n]


  def get_nth_field_many(self, table, keys, n=0):
    '''
    Variant of get_nth_field for multiple keys.
    '''
    for vs in self.get_many(table, keys):
      if vs is None:
        yield None
      else:
        yield vs[n]
