'''
TODO:
- Add auto-deletion of stale indexes (computed fields)
- Add logging
- Add documentation, particularly around compound indexes, also nested routes
'''


import pickle
import pymongo
import re
import string
import threading

from bson.codec_options import CodecOptions
from collections import defaultdict
from contextlib import contextmanager, suppress
from datetime import date, datetime, timezone
from time import time
from uuid import UUID, uuid4

from ..abstract_base_db import AbstractBaseDB
from ..exceptions import *
from ..index import Index
from ..util.fn_serialize import build_fn_key, dumps_fn, loads_fn


# The actual structure of our documents looks like the following:
# {
#   '_id': uuid
#   'data': serialized data (see _encode / _decode functions)
#   'hints': hints for deserializing data (see _encode / _decode functions)
#   'version': if relevant, stores the version ID for client-side indexes
#   <encoded_index_id_1>: stores data for the corresponding client-side index
#   <encoded_index_id_2>: the key is determined by the _encode_field_id function
#   ...
# }
# MongoDB stores the field names with every document, this we want them to be as small as possible. However, we want to
# make sure to avoid collisions with encoded index IDs. Thus, make sure that these characters you use for these fields
# do not overlap with those used by the '_encode_field_id' function.
_FIELD_DATA = '+'
_FIELD_HINTS = '-'
_FIELD_VERSION = '='


# For mapping the symbols in a 'Condition' filter to MongoDB filters
_CONDITION_FILTERS = {
  '<': '$lt',
  '<=': '$lte',
  '>': '$gt',
  '>=': '$gte',
  '==': '$eq',
  '!=': '$ne'
}


# A helper function that makes it easier to replace special characters in a string
# 'to_sub' is a list of (regex_pattern, string) tuples
def _sub_all(s, to_sub):
  for pattern, sub in to_sub:
    s = pattern.sub(sub, s)
  return s


# The following functions are concerned with encoding / decoding strings to fit within the limitations of MongoDB.

# Encoding collection names (data types):
# MongoDB imposes restrictions on the names of collections, and we have some reserved names for metadata collections.
# We want to allow any name rather than having esoteric exceptions.
# Example:
# 'my$data#type#D' -> 'my#Ddata##type##D'

_DATA_TYPE_SUB = [
  (re.compile('#'), '##'),
  (re.compile('\\$'), '#D'),
  (re.compile('\0'), '#0'),
  (re.compile('^system.'), '#system.'),
  (re.compile('^_data_type_index$'), '#_data_type_index'),
  (re.compile('^_data_type_index_seq$'), '#_data_type_index_seq')
]
def _encode_data_type(data_type):
  if not data_type or type(data_type) != str:
    raise Exception('Data type must be a non-empty string')
  return _sub_all(data_type, _DATA_TYPE_SUB)


# Encoding documents:
# We want to store MongoDB documents using the native BSON serialization to take advantage of native index
# functionality. BSON is limited - any types that can't be used for indexes (like objects) should be serialized using
# pickle. Raw dates (and potentially other index types we want to support in the future) are also not supported.
# MongoDB also has limitations on the names of fields, so we have to encode / decode those as well. For indexable
# types, like dates, it is important to preserve the "path" to the value within the document object, which means we
# cannot alter the field name to indicate that we should apply special treatment when deserializing the object. We also
# cannot alter the value, since MongoDB needs the raw value to build the index. This is why we use an additional
# 'hints' structure in these cases. Note that we have special handling for dictionary fields to allow any picklable
# type to be used.
# Example:
# { 'int': 1, 'a.b': 2, 'a?b': 3, 'object': object(), 'date': <python_date> } ->
# (
#   { 'int': 1, 'a?_b': 2, 'a??b': 3, 'object': { '?': object() }, 'nested': { date': <python_date> } },
#   { 'nested': { 'date': 'D' } } # Hints
# )

_FIELD_ENCODE_SUB = [
  (re.compile('\\?'), '?-'),
  (re.compile('\\.'), '?_'),
  (re.compile('\0'), '?0'),
  (re.compile('^\\$'), '?D'),
  (re.compile('^#'), '?#')
]
def _encode_field(field):
  if (type(field) is tuple):
    return ('.').join(_encode_field(x) for x in field)
  if (type(field) is str) and field:
    return _sub_all(field, _FIELD_ENCODE_SUB)
  return '#' + _encode_field(pickle.dumps(field).decode('utf-8', 'surrogateescape'))

_FIELD_DECODE_SUB = [
  (re.compile('^\\?#'), '#'),
  (re.compile('^\\?D'), '$'),
  (re.compile('\\?0'), '\0'),
  (re.compile('\\?_'), '.'),
  (re.compile('\\?-'), '?')
]
def _decode_field(field):
  if field[0] == '#':
    return pickle.loads(_decode_field(field[1:]).encode('utf-8', 'surrogateescape'))
  return _sub_all(field, _FIELD_DECODE_SUB)

_SUPPORTED_TYPES = {
  bool: None,
  str: None,
  int: None,
  float: None,
  UUID: None,
  date: {
    'hint': 'D',
    'to': lambda x: datetime(x.year, x.month, x.day, 0, 0)
  },
  datetime: None,
  bytes: None
}

_TRANSLATE_BY_HINT = {
  'D': lambda x: x.date()
}

def _encode_helper(data, path, hints):
  if data is None:
    return None
  data_type = type(data)
  if data_type is dict:
    return { _encode_field(key): _encode_helper(value, path + (key,), hints) for key, value in data.items() }
  if data_type in _SUPPORTED_TYPES:
    translate = _SUPPORTED_TYPES[data_type]
    if translate:
      hints[path] = translate['hint']
      return translate['to'](data)
    return data
  return { '?': pickle.dumps(data) }

def _encode(data):
  hints = {}
  encoded = _encode_helper(data, (), hints)
  # Note that we can't store a plain dictonary with non-sting keys using BSON (hence the pickling)
  return encoded, pickle.dumps(tuple(hints.items()))

def _decode_helper(data, path, hints):
  if type(data) is dict:
    if (len(data) == 1) and ('?' in data):
      return pickle.loads(data['?'])
    return { _decode_field(key): _decode_helper(value, path + (key,), hints) for key, value in data.items() }
  hint = hints.get(path)
  if hint:
    return _TRANSLATE_BY_HINT[hint](data)
  return data

def _decode(data, hints):
  return _decode_helper(data, (), dict(pickle.loads(hints)))


# Encoding index IDs:
# We generate numeric IDs to identify the indexes succinctly. This is desirable as these index IDs are stored as string
# fields in each document. We can shorten them further by encoding the integer ID int a field-safe UTF-8 representation
# of its bytes. Note that we never actually need to decode these IDs.

_CHAR_SET = string.ascii_letters + string.digits
def _encode_field_id(index_id):
  # Note: make sure that the output of this function does not overlap with any top-level reserved fields
  encoded = ''
  while index_id > 0:
    index_id, r = divmod(index_id, len(_CHAR_SET))
    encoded = _CHAR_SET[r - 1] + encoded
  return encoded


def _build_version_id(field_ids):
  return ('.').join(sorted(field_ids))


# A decorator that handles expected document validation exceptions (thrown if we try to create / update a value with
# missing indexes) and retries the decorated function after reloading the indexes. Realistically, one retry attempt
# should suffice, but having multiple retries is probably prudent to accommodate extreme race conditions.
_MAX_INDEX_REFRESH_ATTEMPTS = 3
def _retry_load_computed_fields(fn):
  def decorator(self, data_type, *args, **kwargs):
    for i in range(_MAX_INDEX_REFRESH_ATTEMPTS):
      try:
        return fn(self, data_type, *args, **kwargs)
      except pymongo.errors.WriteError as e:
        if str(e) != 'Document failed validation':
          raise
      self._load_computed_fields(data_type)
    raise Exception('Incompatible index version')
  return decorator


# Checks / transforms a DB filter condition
def _check_condition(value):
  for filter_key, filter_value in value.items():
    if filter_value is None:
      if (filter_key == '$eq') and (filter_value is None):
        return None
      if filter_key != '$ne':
        raise Exception("Can only use null values with '==' and '!=' queries")
  return value


def _build_db_filter_sort(index_fields, condition_filter, sort_direction):
  sort = [ *( (x, sort_direction) for x in index_fields ), ('_id', sort_direction) ]
  db_filter = defaultdict(dict)
  n_values = None
  for filter_key, values in condition_filter.items():
    n_values = len(values) if n_values is None else n_values
    if n_values != len(values):
      raise Exception('Inconsistent number of values in chained compound index condition')

    for i, value in enumerate(values):
      index_field = index_fields[i]
      if i == n_values - 1:
        db_filter[index_field][_CONDITION_FILTERS[filter_key]] = value
      else:
        if ('$eq' in db_filter[index_field]) and (db_filter[index_field]['$eq'] != value):
          raise Exception('Unequal prefix values in chained compound index condition')
        else:
          db_filter[index_field]['$eq'] = value
  for field, filter_value in db_filter.items():
    db_filter[field] = _check_condition(filter_value)
  return db_filter, sort


class DB(AbstractBaseDB):


  _DB_CODEC_OPTIONS = CodecOptions(tz_aware=True, tzinfo=timezone.utc)
  def __init__(self, connection_string, database):
    self._client = pymongo.MongoClient(connection_string)
    self._db = pymongo.database.Database(self._client, database, codec_options=self._DB_CODEC_OPTIONS)

    # This structure stores metadata about our data types:
    # {
    #   <data_type>: {
    #     'collection': <collection>, # The DB colleciton object (for convenience)
    #     'index_fields': {
    #       # This structure only contains indexes that we explicitly initialized
    #       <index_name>: <field>,
    #       ...
    #     },
    #     'indexes': {
    #       # This structure includes all client-side indexes (ones for which we have to generate and store a
    #       # specialized index value)
    #       <index_id>: <key_fn>,
    #       ...
    #     },
    #     # The version can be derived from the index IDs, but easier to compute it once when those are determined
    #     'version': <string>
    #   },
    #   ...
    # }
    self._data_types = {}

    self._thread_local = threading.local()

    # https://api.mongodb.com/python/current/api/pymongo/collection.html#pymongo.collection.Collection.create_index
    #
    # Create a collection that contains metadata about our computed fields - its documents are structured as follows:
    # {
    #   _id: <auto-generated>,
    #   data_type: string,
    #   fn: serialized getter function,
    #   fn_key: key that uniquely identifies the getter function,
    #   heartbeat_at: datetime that tells if this field is still in use
    # }
    collection = pymongo.collection.Collection(self._db, '_computed_field')
    collection.create_index([ ('data_type', pymongo.ASCENDING), ('fn_key', pymongo.ASCENDING) ], unique=True)

    # Contains metadata about our indexes:
    # {
    #   _id: <auto-generated>,
    #   data_type: string,
    #   field: string,
    #   heartbeat_at: datetime that tells if this index is still in use
    # }
    # collection = pymongo.collection.Collection(self._db, '_index')
    # collection.create_index([ ('data_type', pymongo.ASCENDING), ('field', pymongo.ASCENDING) ], unique=True)


  def describe(self):
    # TODO: fix this, it's broken
    output = {}
    for data_type, data_type_schema in self._data_types.items():
      output[data_type] = list(data_type_schema['index_fields'].keys())
    return output


  def _get_or_create_computed_field(self, data_type, computed_field):
    collection = pymongo.collection.Collection(self._db, '_computed_field')
    # TODO: test the function serialization / deserialization here
    key = { 'data_type': data_type, 'fn_key': build_fn_key(computed_field.fn) }
    doc = { '$set': { **key, 'fn': dumps_fn(computed_field.fn), 'field': computed_field.field, 'heartbeat_at': time() } }
    result = collection.find_one_and_update(key, doc, projection=[], upsert=True, return_document=pymongo.ReturnDocument.AFTER)
    # TODO: these IDs could take up lots of space - you could optimize this by mapping them down to much smaller names
    return str(result['_id'])


  def _load_computed_fields(self, data_type, last_heartbeat_threshold_s=7200):
    with self._client.start_session() as session:
      # Load the current active indexes from the DB and update the version
      collection = pymongo.collection.Collection(self._db, '_computed_field')
      oldest_heartbeat_at = time() - last_heartbeat_threshold_s
      computed_fields = collection.find({ 'data_type': data_type, 'heartbeat_at': { '$gt': oldest_heartbeat_at } })
      field_ids = []
      for computed_field in computed_fields:
        field_id = str(computed_field['_id'])
        self._data_types[data_type]['computed'].append((field_id, loads_fn(computed_field['fn'])))
        field_ids.append(field_id)
      version_id = _build_version_id(field_ids)
      self._data_types[data_type]['version'] = version_id

      # Set up the DB validation rules
      if field_ids:
        schema = {
          'required': [ _FIELD_VERSION ],
          'properties': { _FIELD_VERSION: { 'bsonType': 'string', 'enum': [ version_id ] } }
        }
      else:
        schema = {}
      self._db.command({
        'collMod': _encode_data_type(data_type),
        'validator': { '$jsonSchema': schema },
        'validationLevel': 'moderate'
      })


  def _backfill(self, data_type):
    collection = self._get_collection(data_type)
    version_id = self._data_types[data_type]['version']
    for doc in collection.find({ _FIELD_VERSION: { '$ne': version_id } }):
      data_id = doc['_id']
      data = self._build_data(doc)
      self.update(data_type, data_id, data)


  def init(self, data_type, computed=[], indexes=[]):
    if data_type in self._data_types:
      raise DuplicateDataTypeException("Data type '%s' already initialized" % data_type)

    encoded_data_type = _encode_data_type(data_type)
    try:
      data_type_collection = pymongo.collection.Collection(self._db, encoded_data_type, create=True)
    except pymongo.errors.OperationFailure:
      data_type_collection = pymongo.collection.Collection(self._db, encoded_data_type)

    self._data_types[data_type] = { 'collection': data_type_collection, 'computed': [], 'indexes': {} }

    computed_field_ids = {}
    for computed_field in computed:
      field_id = self._get_or_create_computed_field(data_type, computed_field)
      computed_field_ids[computed_field.field] = field_id
      # TODO: start the heartbeat for this computed field to report its heartbeat_at time

    for index in indexes:
      if (type(index) == str) or (type(index) == tuple):
        index = Index(index)
      encoded = tuple( computed_field_ids.get(x, _FIELD_DATA + '.' + _encode_field(x)) for x in index.fields )
      self._data_types[data_type]['indexes'][index.fields] = encoded
      data_type_collection.create_index([ (x, pymongo.ASCENDING) for x in encoded ])

    # Importantly, this call also creates DB validation that ensures all create / update calls from any client will
    # fail until they reload the computed fields
    self._load_computed_fields(data_type)

    if computed:
      self._backfill(data_type)


  def _get_collection(self, data_type):
    with suppress(KeyError):
      return self._data_types[data_type]['collection']
    raise MissingDataTypeException("Data type '%s' has not been initialized" % data_type)


  @contextmanager
  def transaction(self):
    with self._client.start_session() as session:
      with session.start_transaction():
        if not getattr(self._thread_local, 'transactions', False):
          self._thread_local.transactions = []
        self._thread_local.transactions.append(session)
        try:
          yield
        finally:
          self._thread_local.transactions.pop()


  def _get_session(self):
    sessions = getattr(self._thread_local, 'transactions', [])
    return sessions[-1] if sessions else None


  def _build_serializable_document(self, data_type, data_id, data):
    encoded, hints = _encode(data)
    serializable_document = { '_id': data_id, _FIELD_DATA: encoded, _FIELD_HINTS: hints }
    if self._data_types[data_type]['version']:
      serializable_document[_FIELD_VERSION] = self._data_types[data_type]['version']
      for encoded_field_id, fn in self._data_types[data_type]['computed']:
        try:
          serializable_document[encoded_field_id] = fn(data)
        except Exception as e:
          # TODO: it would be very nice to include the actual name of the field in this output
          msg = "Error computing field '%s' for data type '%s'" % (encoded_field_id, data_type)
          raise ComputedFieldException(msg) from e
    return serializable_document


  def _build_data(self, serializable_document):
    return _decode(serializable_document[_FIELD_DATA], serializable_document[_FIELD_HINTS])


  def _build_verified_db_filter_sort(self, data_type, conditions, sort_direction=pymongo.ASCENDING):
    db_filter = {}
    sort = []
    for condition in conditions:
      if condition.fields not in self._data_types[data_type]['indexes']:
        msg = "Index '%s' has not been initialized for data type '%s'" % (str(condition.fields), data_type)
        raise MissingIndexException(msg)
      encoded_index_fields = self._data_types[data_type]['indexes'][condition.fields]
      cur_db_filter, cur_sort = _build_db_filter_sort(encoded_index_fields, condition.filter, sort_direction)
      db_filter.update(cur_db_filter)
      sort = sort + cur_sort
    if not sort:
      sort.append(('_id', sort_direction))
    return db_filter, sort


  @_retry_load_computed_fields
  def create(self, data_type, data):
    collection = self._get_collection(data_type)
    data_id = uuid4()
    doc = self._build_serializable_document(data_type, data_id, data)
    collection.insert_one(doc, session=self._get_session())
    return data_id


  @_retry_load_computed_fields
  def update(self, data_type, data_id, data, *conditions):
    collection = self._get_collection(data_type)
    doc = self._build_serializable_document(data_type, data_id, data)
    # TODO: I don't think it is strictly necessary that these conditions are indexed since we're already querying by
    # the data ID, but the way we actually store the data underneath, its individual fields are not available for
    # queries UNLESS they are indexed.
    db_filter, _ = self._build_verified_db_filter_sort(data_type, conditions)
    db_filter['_id'] = data_id
    result = collection.replace_one(db_filter, doc, session=self._get_session())
    if not result.matched_count:
      raise UpdateConditionException()


  def get(self, data_type, data_id):
    collection = self._get_collection(data_type)
    result = collection.find_one({ '_id': data_id }, session=self._get_session())
    if not result:
      raise NotFoundException("ID '%s' not found" % data_id)
    return self._build_data(result)


  def get_batch(self, data_type, data_ids):
    collection = self._get_collection(data_type)
    results = collection.find({ '_id': { '$in': list(data_ids) } }, session=self._get_session())
    return [ (x['_id'], self._build_data(x)) for x in results ]


  def delete(self, data_type, data_id):
    collection = self._get_collection(data_type)
    collection.delete_one({ '_id': data_id })


  def iter(self, data_type, *conditions, descending=False, limit=None):
    collection = self._get_collection(data_type)
    sort_direction = pymongo.DESCENDING if descending else pymongo.ASCENDING
    db_filter, sort = self._build_verified_db_filter_sort(data_type, conditions, sort_direction)
    results = collection.find(db_filter, sort=sort, limit=limit or 0, session=self._get_session())
    yield from [ (x['_id'], self._build_data(x)) for x in results ]


  def list(self, data_type, *conditions, descending=False, limit=None):
    return list(self.iter(data_type, *conditions, descending=descending, limit=limit))

  # For cloning data from other databases

  def _import_items(self, data_type, items):
    # 'items' is a list as returned by a DB's implementation of the 'list' function
    if not items:
      return
    collection = self._get_collection(data_type)
    db_items = [ self._build_serializable_document(data_type, i, j) for i, j in items ]
    collection.insert_many(db_items)

  # For tests

  def _drop(self, collection_name):
    collection = pymongo.collection.Collection(self._db, collection_name)
    collection.delete_many({})
    self._db.drop_collection(collection_name)

  def _remove_computed(self):
    self._drop('_computed_field')

  def _remove_data_type(self, data_type):
    self._drop(_encode_data_type(data_type))
