import base64
import pickle
import pymongo

from bson import Decimal128
from bson.binary import UuidRepresentation
from bson.codec_options import CodecOptions
from collections import defaultdict
from datetime import date, datetime, timezone
from decimal import Decimal
from pymongo.errors import BulkWriteError, DuplicateKeyError
from uuid import UUID, uuid4

from .abstract_base_db import AbstractBaseDB
from .exceptions import *
from .field import Field
from .utils import INDEXABLE_TYPES, dereference


# The actual structure of our documents looks like the following:
# {
#   '_id': uuid
#   '+': encoded data (see _encode / _decode functions)
#   '-': hints for deserializing data (see _encode / _decode functions)
# }
_FIELD_DATA_KEY = '+'
_FIELD_HINTS_KEY = '-'


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


# 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, and we want to preserve some semblance of
# legibility just in case we need to interact directly with the DB for some reason.
# Example:
# 'my$data?type?D' -> 'my?Ddata??type?D'

_DATA_TYPE_CHAR_ESC = [ ('?', '?'), ('$', 'D'), ('\0', '0') ]

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')
  encoded = ''
  for c in data_type:
    for esc, sub in _DATA_TYPE_CHAR_ESC:
      if c == esc:
        encoded += '?' + sub
        break
    else:
      encoded += c
  return ('?' + encoded) if encoded.startswith('system.') else encoded

def _decode_data_type(encoded):
  decoded = ''
  escape = False
  for c in encoded:
    if escape:
      escape = False
      for esc, sub in _DATA_TYPE_CHAR_ESC:
        if c == sub:
          decoded += esc
          break
      else:
        if c != 's':
          raise Exception(f'Error decoding data type name, unexpected escaped character: {c}')
    elif c == '?':
      escape = True
    else:
      decoded += c
  return decoded


# 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, like a tuple in the following example.
#
# Example:
# { 'int': 1, 'a.b': 2, 'a?b': 3, <uuid>: 4, ('a', 'b'): 5, object': object(), 'date': <python_date> } ->
# (
#   { 'int': 1, 'a?_b': 2, 'a??b': 3, '?$<pickled_uuid>': 4, 'a.b': 5, 'object': { '?': object() }, 'nested': { date': <python_date> } },
#   { 'nested': { 'date': 'D' } } # Hints
# )

# A tuple of (to_replace, replace_with, _start_only)
_FIELD_CHAR_ESC = [ ('?', '?'), ('.', '_'), ('\0', '0'), ('$', 'D') ]

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:
    encoded = ''
    for c in field:
      for esc, sub in _DATA_TYPE_CHAR_ESC:
        if c == esc:
          encoded += '?' + sub
          break
      else:
        encoded += c
    return encoded
  return '?$' + _encode_field(base64.b85encode(pickle.dumps(field)).decode())

def _decode_field(encoded):
  decoded = ''
  escape = False
  for c in encoded:
    if escape:
      escape = False
      for esc, sub in _DATA_TYPE_CHAR_ESC:
        if c == sub:
          decoded += esc
          break
      else:
        if c == '$':
          return pickle.loads(base64.b85decode(_decode_field(encoded[2:]).encode()))
        else:
          raise Exception(f'Error decoding field name, unexpected escaped character: {c}')
    elif c == '?':
      escape = True
    else:
      decoded += c
  return decoded


# A map of data type => tuple of (hint, translate_to_fn, translate_from_fn)
# Don't use '?' as the type identifier as that is used for unsupported data types
_TRANSLATED_TYPES = [
  {
    'type': date,
    'hint': 'd',
    'to': lambda x: datetime(x.year, x.month, x.day, 0, 0),
    'from': lambda x: x.date()
  },
  {
    'type': Decimal,
    'hint': 'D',
    'to': Decimal128,
    'from': lambda x: x.to_decimal()
  }
]
_TRANSLATED_TYPES_BY_TYPE = { x['type']: x for x in _TRANSLATED_TYPES }
_TRANSLATED_TYPES_FROM_BY_HINT = { x['hint']: x['from'] for x in _TRANSLATED_TYPES }

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() }
  translate = _TRANSLATED_TYPES_BY_TYPE.get(data_type)
  if translate:
    hints[path] = translate['hint']
    return translate['to'](data)
  if data_type in INDEXABLE_TYPES:
    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-string 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 _TRANSLATED_TYPES_FROM_BY_HINT[hint](data)
  return data

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

def _encode_document(data_id, data):
  encoded, hints = _encode(data)
  return { '_id': data_id, _FIELD_DATA_KEY: encoded, _FIELD_HINTS_KEY: hints }

def _decode_document(encoded):
  decoded = _decode(encoded[_FIELD_DATA_KEY], encoded[_FIELD_HINTS_KEY])
  return decoded

def _build_db_filter_for_conditions(conditions):
  db_filter = defaultdict(lambda: defaultdict(dict))
  for condition in conditions:
    if not condition._filter:
      continue
    obj_path = _FIELD_DATA_KEY + '.' + _encode_field(condition._keys)
    for filter_key, operand in condition._filter.items():
      operator, needs_encode = _CONDITION_FILTERS[filter_key]
      if db_filter[obj_path][operator]:
        raise FieldFilterException(f"Duplicate '{filter_key}' operators")
      db_filter[obj_path][operator] = _encode(operand)[0] if needs_encode else operand
  return db_filter


_DB_CODEC_OPTIONS = CodecOptions(tz_aware=True, tzinfo=timezone.utc, uuid_representation=UuidRepresentation.STANDARD)


class DB(AbstractBaseDB):

  def __init__(self, connection_string, database):
    self._client = pymongo.MongoClient(connection_string)
    self._db = pymongo.database.Database(self._client, database, codec_options=_DB_CODEC_OPTIONS)


  def list_data_types(self):
    return self._db.list_collection_names()


  def _get_collection(self, data_type):
    return pymongo.collection.Collection(self._db, _encode_data_type(data_type))


  def create(self, data_type, data, *unique_fields):
    collection = self._get_collection(data_type)
    data_id = uuid4()
    encoded = _encode_document(data_id, data)
    if unique_fields:
      for field in unique_fields:
        if field._filter:
          raise Exception('unique_fields passed to create should not have conditions')
      conditions = [ Field(*x._keys) == dereference(data, x._keys) for x in unique_fields ]
      db_filter = _build_db_filter_for_conditions(conditions)
      response = collection.update_one(db_filter, { '$setOnInsert': encoded }, upsert=True)
      if not response.upserted_id:
        raise DuplicateKeyException()
    else:
      try:
        collection.insert_one(encoded)
      except DuplicateKeyError as e:
        raise DuplicateKeyException() from e
    return data_id


  def update(self, data_type, data_id, data, *conditions):
    collection = self._get_collection(data_type)
    db_filter = _build_db_filter_for_conditions(conditions)
    db_filter['_id'] = data_id
    try:
      result = collection.replace_one(db_filter, _encode_document(data_id, data))
    except DuplicateKeyError as e:
      raise DuplicateKeyException() from e
    if not result.matched_count:
      raise NotFoundException()


  def get(self, data_type, *data_ids):
    if not data_ids:
      raise Exception('No data IDs provided')
    collection = self._get_collection(data_type)
    results = list(collection.find({ '_id': { '$in': list(data_ids) } }))
    if len(results) != len(data_ids):
      raise NotFoundException()
    if len(data_ids) == 1:
      return _decode_document(results[0])
    by_id = { x['_id']: _decode_document(x) for x in results }
    return [ by_id[x] for x in data_ids ]


  def delete(self, data_type, *data_ids):
    if not len(data_ids):
      raise Exception('No data IDs provided')
    collection = self._get_collection(data_type)
    # TODO: add some checking to this?
    collection.delete_many({ '_id': { '$in': data_ids } })


  def iter(self, data_type, *conditions, reverse=False, limit=None):
    db_filter = _build_db_filter_for_conditions(conditions)
    sort_dir = -1 if reverse else 1
    sort = { (_FIELD_DATA_KEY + '.' + _encode_field(x._keys)): sort_dir for x in conditions }
    sort['_id'] = sort_dir
    results = self._get_collection(data_type).find(db_filter, sort=list(sort.items()), limit=limit or 0)

    # TODO: figure out how to log a warning here when the query uses a COLLSCAN
    # https://www.mongodb.com/docs/manual/reference/explain-results/
    # print()
    # print(results.explain()['queryPlanner']['winningPlan']['inputStage']['stage'])
    # print()

    for result in results:
      yield result['_id'], _decode_document(result)


  def count(self, data_type, *conditions, limit=None):
    collection = self._get_collection(data_type)
    db_filter = _build_db_filter_for_conditions(conditions)
    return collection.count_documents(db_filter, limit=limit) if limit else collection.count_documents(db_filter)


  def create_index(self, data_type, *fields):
    collection = self._get_collection(data_type)
    encoded_fields = [ _FIELD_DATA_KEY + '.' + _encode_field(x._keys) for x in fields ]
    collection.create_index(encoded_fields, name='op-' + ('-').join(encoded_fields))


  def list_indexes(self, data_type):
    indexes = self._get_collection('user').list_indexes()
    output = []
    for index in indexes:
      as_dict = index.to_dict()
      encoded_fields = list(as_dict['key'].keys())
      if as_dict['name'] != 'op-' + ('-').join(encoded_fields):
        # This index was not created using 'create_index'
        continue
      output.append({
        'fields': [ Field(*[ _decode_field(x) for x in y.split('.')[1:] ]) for y in encoded_fields ],
        'unique': as_dict.get('unique', False)
      })
    return output


  def drop_index(self, data_type, *fields):
    collection = self._get_collection(data_type)
    encoded_fields = [ _FIELD_DATA_KEY + '.' + _encode_field(x._keys) for x in fields ]
    collection.drop_index('op-' + ('-').join(encoded_fields))
