'''
Use this template to test that a store implementation abides by the op-storage interface. Override the global
'STORE_TYPE' and 'CONFIG' variables to specify a specific implementation. The default is to test using an 'in-memory'
store, as that has no downstream dependencies.

For some implementations of the op-storage interface, our downstream dependencies are quite complex, like the AWS RDS
data API. Mocking those dependencies correctly and capturing all the nuances around things like eventual consistency is
difficult to get right. Therefore, I favor this pattern of running integration tests. The downside is that because you
need to have your dependencies properly initialized, running these tests does take some manual setup.
'''


STORE_TYPE = 'in_memory'
STORE_CONFIG = {}


from datetime import datetime, timezone
from freezegun import freeze_time
from functools import wraps
from unittest import TestCase
from uuid import UUID

from op_storage import factory, Index, Condition
from op_storage.exceptions import *


def dict_sort_key(d):
  return tuple((i, d[i]) for i in sorted(d.keys()))


class Test(TestCase):

  # Note: I don't like this pattern for 'setUp' and 'tearDown' of manipulating persistent / global variables. I'd much
  # rather structure them as decorators that delivered additional properties to the test functions - then the scope of
  # any created variables is limited.

  def setUp(self):
    # Create a DB instance specifically for setting up our tests - I want to make sure this doesn't interfere with the
    # actual instance under test
    set_up_db_instance = factory.build(STORE_TYPE, STORE_CONFIG)
    set_up_db_instance._remove_indexes()
    set_up_db_instance._remove_data_type('test_user')
    set_up_db_instance._remove_data_type('!@#$%^&*():.,/')

    global db
    db = factory.build(STORE_TYPE, STORE_CONFIG)

  def test_init(self):
    db.init('test_user')
    self.assertEqual(db.describe(), { 'test_user': [] })

  def test_init_duplicate(self):
    db.init('test_user')
    with self.assertRaises(DuplicateDataTypeException):
      db.init('test_user')

  def test_special_characters(self):
    db.init('!@#$%^&*():.,/', indexes=[ '!@#$%^&*():.,/' ])
    data_id = db.create('!@#$%^&*():.,/', { '!@#$%^&*():.,/': 'value' })
    data = db.get('!@#$%^&*():.,/', data_id)
    self.assertEqual(data['!@#$%^&*():.,/'], 'value')

  def test_get_missing_data_type(self):
    with self.assertRaises(MissingDataTypeException):
      db.get('test_user', UUID(int=0))

  def test_get_not_found(self):
    db.init('test_user')
    with self.assertRaises(NotFoundException):
      db.get('test_user', UUID(int=0))

  def test_create_get(self):
    db.init('test_user')
    user = { 'name': 'Alex', 'email': 'AlexExarhos@gmail.com', 'age': 31 }
    user_id = db.create('test_user', user)
    retrieved_user = db.get('test_user', user_id)
    self.assertEqual(user, retrieved_user)

  def test_get_batch_not_found(self):
    db.init('test_user')
    retrieved_users = db.get_batch('test_user', [ UUID(int=0) ])
    self.assertEqual(len(retrieved_users), 0)

  def test_create_get_batch(self):
    db.init('test_user')
    user_1 = { 'name': 'Alex', 'age': 31 }
    user_2 = { 'name': 'Kelly', 'age': 29 }
    user_1_id = db.create('test_user', user_1)
    user_2_id = db.create('test_user', user_2)
    retrieved_users = db.get_batch('test_user', [ user_1_id, user_2_id ])
    self.assertEqual(len(retrieved_users), 2)
    self.assertEqual(dict(retrieved_users)[user_1_id], user_1)
    self.assertEqual(dict(retrieved_users)[user_2_id], user_2)

  def test_create_update_get(self):
    db.init('test_user')
    user = { 'name': 'Alex', 'email': 'AlexExarhos@gmail.com', 'age': 31 }
    user_updated = { 'name': 'Alex', 'email': 'AlexExarhos@gmail.com', 'age': 32 }
    user_id = db.create('test_user', user)
    db.update('test_user', user_id, user_updated)
    retrieved_user = db.get('test_user', user_id)
    self.assertEqual(retrieved_user, user_updated)

  def test_create_delete_get(self):
    db.init('test_user')
    user = { 'name': 'Alex', 'email': 'AlexExarhos@gmail.com', 'age': 31 }
    user_id = db.create('test_user', user)
    db.delete('test_user', user_id)
    with self.assertRaises(NotFoundException):
      db.get('test_user', UUID(int=0))

  def test_init_indexes(self):
    db.init('test_user', indexes=[ Index('age') ])
    self.assertEqual(db.describe(), { 'test_user': [ 'age' ] })

  def test_init_indexes_as_str(self):
    db.init('test_user', indexes=[ 'age' ])
    self.assertEqual(db.describe(), { 'test_user': [ 'age' ] })

  def test_init_duplicate_indexes_with_str(self):
    with self.assertRaises(DuplicateIndexException):
      db.init('test_user', indexes=[ 'age', Index('age') ])

  def test_index_added_externally(self):
    db.init('test_user', indexes=[ 'name' ])
    db.create('test_user', { 'name': 'Kelly', 'age': 29 })
    db_2 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_2.init('test_user', indexes=[ 'age' ])
    db_2.create('test_user', { 'name': 'Alex', 'age': 31 })
    self.assertEqual(len(db.list('test_user', db['name'] != '')), 2)
    self.assertEqual(len(db_2.list('test_user', db_2['age'] > 10)), 2)

  def test_index_same_added_externally(self):
    db.init('test_user', indexes=[ 'name' ])
    db.create('test_user', { 'name': 'Kelly', 'age': 29 })
    db_2 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_2.init('test_user', indexes=[ 'name' ])
    db_2.create('test_user', { 'name': 'Alex', 'age': 31 })
    self.assertEqual(len(db.list('test_user', db['name'] != '')), 2)
    self.assertEqual(len(db_2.list('test_user', db_2['name'] != '')), 2)

  def test_index_same_name_different_fn_added_externally(self):
    db.init('test_user', indexes=[ 'name' ])
    db.create('test_user', { 'name': 'Kelly', 'age': 29 })
    db_2 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_2.init('test_user', indexes=[ Index('name', lambda data: data['name'].lower()) ])
    db_2.create('test_user', { 'name': 'Alex', 'age': 31 })
    self.assertEqual(len(db.list('test_user', db['name'] == 'Alex')), 1)
    self.assertEqual(len(db.list('test_user', db['name'] == 'alex')), 0)
    self.assertEqual(len(db_2.list('test_user', db_2['name'] == 'Alex')), 0)
    self.assertEqual(len(db_2.list('test_user', db_2['name'] == 'alex')), 1)

  def test_list_index_no_data(self):
    db.init('test_user', indexes=[ 'age' ])
    self.assertEqual(db.list('test_user', db['age'] > 3), [])

  # The following tests are concerned with querying users. To make them simpler, we provide a 'query' decorator that
  # calls the wrapped test function with both the 'db.list' and 'db.iter' functions - those functions should behave
  # identically except that one returns a list and one an iterable. It also creates the users from the static
  # 'TEST_USERS' structure and passes the corresponding IDs into the test function as an additional parameter.

  TEST_USERS = [
    { 'name': 'Alex', 'age': 31 },
    { 'name': 'Kelly', 'age': 29 },
    { 'name': 'Ben', 'age': 27 },
    { 'name': 'Jared', 'age': 29 }
  ]

  def query(indexes=[]):
    def decorator(test_fn):
      @wraps(test_fn)
      def wrapper(self):
        db.init('test_user', indexes=indexes)
        user_ids = [ db.create('test_user', i) for i in self.TEST_USERS ]
        test_fn(self, db.list, user_ids)
        test_fn(self, lambda *args, **kwargs: list(db.iter(*args, **kwargs)), user_ids)
      return wrapper
    return decorator

  @query()
  def test_query_no_conditions(self, query_fn, user_ids):
    retrieved_users = [ j for i, j in query_fn('test_user') ]
    self.assertEqual(sorted(retrieved_users, key=dict_sort_key), sorted(self.TEST_USERS, key=dict_sort_key))

  @query()
  def test_query_order_no_conditions(self, query_fn, user_ids):
    retrieved_user_ids = [ i for i, j in query_fn('test_user') ]
    self.assertEqual(retrieved_user_ids, sorted(user_ids))

  @query()
  def test_query_order_descending_no_conditions(self, query_fn, user_ids):
    retrieved_user_ids = [ i for i, j in query_fn('test_user', descending=True) ]
    self.assertEqual(retrieved_user_ids, sorted(user_ids, reverse=True))

  @query()
  def test_query_limit(self, query_fn, user_ids):
    retrieved_users = query_fn('test_user', limit=2)
    self.assertEqual(len(retrieved_users), 2)

  @query([ 'age' ])
  def test_query_index(self, query_fn, user_ids):
    results = query_fn('test_user', db['age'] > 30)
    retrieved_users = [ j for i, j in results ]
    self.assertEqual(retrieved_users, [ { 'name': 'Alex', 'age': 31 } ])

  @query([ 'age', 'name' ])
  def test_query_index_range(self, query_fn, user_ids):
    results = query_fn('test_user', 30 < db['age'] < 35)
    retrieved_users = [ j for i, j in results ]
    self.assertEqual(retrieved_users, [ { 'name': 'Alex', 'age': 31 } ])

  # @query([ 'age', 'name' ])
  # def test_query_multiple_indexes(self, query_fn, user_ids):
  #   results = query_fn('test_user', db['age'] > 20, db['age'] < 30, db['name'] == 'Kelly')
  #   retrieved_users = [ j for i, j in results ]
  #   self.assertEqual(retrieved_users, [ { 'name': 'Kelly', 'age': 29 } ])

  @query([ Index('name-lower', key=lambda data: data['name'].lower()) ])
  def test_query_index_fn(self, query_fn, user_ids):
    results = query_fn('test_user', db['name-lower'] == 'alex')
    retrieved_users = [ j for i, j in results ]
    assert retrieved_users == [ { 'name': 'Alex', 'age': 31 } ]

  # This test creates multiple database instances that all depend on differing indexes running simultaneously - think
  # of these as different versions of your code that may be deployed and running over the same backing store

  def test_index_lifecycle_1(self):
    # In version 1, we decide we want to be able to query by 'age' 
    db_1 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_1.init('test_user', indexes=[ 'age' ])

    # When we create a user, the 'age' index is populated automatically and available to query
    db_1.create('test_user', { 'name': 'Alex', 'age': 31 })
    users = [ user for user_id, user in db_1.list('test_user', db_1['age'] == 31) ]
    self.assertEqual(users, [ { 'name': 'Alex', 'age': 31 } ])

    # In version 2, we realized we don't actually need the index on 'age' so we remove it
    db_2 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_2.init('test_user')

    # Even though 'db_2' never explicitly initialized the 'age' index, 'db_1' can still use it to query and pull in the
    # new object created by 'db_2'
    db_2.create('test_user', { 'name': 'Kelly', 'age': 29 })
    users = [ user for user_id, user in db_1.list('test_user', db_1['age'] == 29) ]
    self.assertEqual(users, [ { 'name': 'Kelly', 'age': 29 } ])

    # In version 3, we want to query by 'name'
    db_3 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_3.init('test_user', indexes=[ 'name' ])

    # The 'name' index did not exist when 'Alex' and 'Kelly' were created, but we can still retrieve those objects by
    # name
    users = [ user for user_id, user in db_3.list('test_user', db_3['name'] == 'Alex') ]
    self.assertEqual(users, [ { 'name': 'Alex', 'age': 31 } ])

    # In version 4, we want our name queries to be case-insensitive, se we change the 'name' index. Note that I'd
    # generally consider this bad practice - a better name for the index would be something like 'name-lower'.
    db_4 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_4.init('test_user', indexes=[ Index('name', key=lambda user: user['name'].lower()) ])

    # Importantly, the behavior does not change for version 3
    users = [ user for user_id, user in db_3.list('test_user', db_3['name'] == 'Alex') ]
    self.assertEqual(users, [ { 'name': 'Alex', 'age': 31 } ])
    # But it does for version 4 (note that I'm querying with the lowercased string 'alex')
    users = [ user for user_id, user in db_4.list('test_user', db_4['name'] == 'alex') ]
    self.assertEqual(users, [ { 'name': 'Alex', 'age': 31 } ])

    # Emulate a production deployment scenario - eventually, older versions of the schema will stop running and the
    # indexes to which only they refer will no longer be relevant. On a new initialization, those indexes will be
    # automatically removed to avoid needlessly taking up storage space. Note that you can still go back and run an
    # older version of the schema at any point, initialization will just take some time to rebuild the index.
    # db_5 = factory.build(STORE_TYPE, STORE_CONFIG)
    
    # # TODO: uncomment the following two lines to test rds-data-postgres
    # from datetime import timedelta
    # db_5._drop_stale_indexes('test_user', stale=timedelta(days=-1))

    # db_5.init('test_user')
    # db_5.create('test_user', { 'name': 'Ben', 'age': 27 })
    # assert db_5.describe()['test_user'] == []

  def test_index_lifecycle_2(self):
    db_1 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_2 = factory.build(STORE_TYPE, STORE_CONFIG)
    db_3 = factory.build(STORE_TYPE, STORE_CONFIG)

    db_1.init('test_user', indexes=[ 'age' ])
    db_2.init('test_user', indexes=[ 'name' ])
    db_3.init('test_user', indexes=[ Index('name', key=lambda x: x['name'].lower()) ])

    user_id = db_1.create('test_user', { 'name': 'Alex', 'age': 31 })

    results = db_3.list('test_user', Condition('name') == 'alex')
    self.assertEqual(len(results), 1)
    results = db_3.list('test_user', Condition('name') == 'Alex')
    self.assertEqual(len(results), 0)
    results = db_2.list('test_user', Condition('name') == 'alex')
    self.assertEqual(len(results), 0)
    results = db_2.list('test_user', Condition('name') == 'Alex')
    self.assertEqual(len(results), 1)
