# Copyright DataStax, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from cassandra import InvalidRequest
from cassandra.cluster import Cluster
from cassandra.cluster import NoHostAvailable
from cassandra.cqlengine import columns, CQLEngineException
from cassandra.cqlengine import connection as conn
from cassandra.cqlengine.management import drop_keyspace, sync_table, drop_table, create_keyspace_simple
from cassandra.cqlengine.models import Model, QuerySetDescriptor
from cassandra.cqlengine.query import ContextQuery, BatchQuery, ModelQuerySet
from tests.integration.cqlengine import setup_connection, DEFAULT_KEYSPACE
from tests.integration.cqlengine.base import BaseCassEngTestCase
from tests.integration.cqlengine.query import test_queryset
from tests.integration import local, CASSANDRA_IP


class TestModel(Model):

    __keyspace__ = 'ks1'

    partition = columns.Integer(primary_key=True)
    cluster = columns.Integer(primary_key=True)
    count = columns.Integer()
    text = columns.Text()


class AnotherTestModel(Model):

    __keyspace__ = 'ks1'

    partition = columns.Integer(primary_key=True)
    cluster = columns.Integer(primary_key=True)
    count = columns.Integer()
    text = columns.Text()

class ContextQueryConnectionTests(BaseCassEngTestCase):

    @classmethod
    def setUpClass(cls):
        super(ContextQueryConnectionTests, cls).setUpClass()
        create_keyspace_simple('ks1', 1)

        conn.unregister_connection('default')
        conn.register_connection('fake_cluster', ['1.2.3.4'], lazy_connect=True, retry_connect=True, default=True)
        conn.register_connection('cluster', [CASSANDRA_IP])

        with ContextQuery(TestModel, connection='cluster') as tm:
            sync_table(tm)

    @classmethod
    def tearDownClass(cls):
        super(ContextQueryConnectionTests, cls).tearDownClass()

        with ContextQuery(TestModel, connection='cluster') as tm:
            drop_table(tm)
        drop_keyspace('ks1', connections=['cluster'])


        # reset the default connection
        conn.unregister_connection('fake_cluster')
        conn.unregister_connection('cluster')
        setup_connection(DEFAULT_KEYSPACE)

    def setUp(self):
        super(BaseCassEngTestCase, self).setUp()

    def test_context_connection_priority(self):
        """
        Tests to ensure the proper connection priority is honored.

        Explicit connection should have higest priority,
        Followed by context query connection
        Default connection should be honored last.

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result priorities should be honored

        @test_category object_mapper
        """
        # model keyspace write/read

        # Set the default connection on the Model
        TestModel.__connection__ = 'cluster'
        with ContextQuery(TestModel) as tm:
            tm.objects.create(partition=1, cluster=1)

        # ContextQuery connection should have priority over default one
        with ContextQuery(TestModel, connection='fake_cluster') as tm:
            with self.assertRaises(NoHostAvailable):
                tm.objects.create(partition=1, cluster=1)

        # Explicit connection should have priority over ContextQuery one
        with ContextQuery(TestModel, connection='fake_cluster') as tm:
            tm.objects.using(connection='cluster').create(partition=1, cluster=1)

        # Reset the default conn of the model
        TestModel.__connection__ = None

        # No model connection and an invalid default connection
        with ContextQuery(TestModel) as tm:
            with self.assertRaises(NoHostAvailable):
                tm.objects.create(partition=1, cluster=1)

    def test_context_connection_with_keyspace(self):
        """
        Tests to ensure keyspace param is honored

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result Invalid request is thrown

        @test_category object_mapper
        """

        # ks2 doesn't exist
        with ContextQuery(TestModel, connection='cluster', keyspace='ks2') as tm:
            with self.assertRaises(InvalidRequest):
                tm.objects.create(partition=1, cluster=1)


class ManagementConnectionTests(BaseCassEngTestCase):

    keyspaces = ['ks1', 'ks2']
    conns = ['cluster']

    @classmethod
    def setUpClass(cls):
        super(ManagementConnectionTests, cls).setUpClass()
        conn.unregister_connection('default')
        conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True)
        conn.register_connection('cluster', [CASSANDRA_IP])

    @classmethod
    def tearDownClass(cls):
        super(ManagementConnectionTests, cls).tearDownClass()

        # reset the default connection
        conn.unregister_connection('fake_cluster')
        conn.unregister_connection('cluster')
        setup_connection(DEFAULT_KEYSPACE)

    def setUp(self):
        super(BaseCassEngTestCase, self).setUp()

    def test_create_drop_keyspace(self):
        """
        Tests drop and create keyspace with connections explicitly set

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result keyspaces should be created and dropped

        @test_category object_mapper
        """

        # No connection (default is fake)
        with self.assertRaises(NoHostAvailable):
            create_keyspace_simple(self.keyspaces[0], 1)

        # Explicit connections
        for ks in self.keyspaces:
            create_keyspace_simple(ks, 1, connections=self.conns)

        for ks in self.keyspaces:
            drop_keyspace(ks, connections=self.conns)

    def test_create_drop_table(self):
        """
        Tests drop and create Table with connections explicitly set

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result Tables should be created and dropped

        @test_category object_mapper
        """
        for ks in self.keyspaces:
            create_keyspace_simple(ks, 1, connections=self.conns)

        # No connection (default is fake)
        with self.assertRaises(NoHostAvailable):
            sync_table(TestModel)

        # Explicit connections
        sync_table(TestModel, connections=self.conns)

        # Explicit drop
        drop_table(TestModel, connections=self.conns)

        # Model connection
        TestModel.__connection__ = 'cluster'
        sync_table(TestModel)
        TestModel.__connection__ = None

        # No connection (default is fake)
        with self.assertRaises(NoHostAvailable):
            drop_table(TestModel)

        # Model connection
        TestModel.__connection__ = 'cluster'
        drop_table(TestModel)
        TestModel.__connection__ = None

        # Model connection
        for ks in self.keyspaces:
            drop_keyspace(ks, connections=self.conns)

    def test_connection_creation_from_session(self):
        """
        Test to ensure that you can register a connection from a session
        @since 3.8
        @jira_ticket PYTHON-649
        @expected_result queries should execute appropriately

        @test_category object_mapper
        """
        cluster = Cluster([CASSANDRA_IP])
        session = cluster.connect()
        connection_name = 'from_session'
        conn.register_connection(connection_name, session=session)
        self.assertIsNotNone(conn.get_connection(connection_name).cluster.metadata.get_host(CASSANDRA_IP))
        self.addCleanup(conn.unregister_connection, connection_name)
        cluster.shutdown()

    def test_connection_from_hosts(self):
        """
        Test to ensure that you can register a connection from a list of hosts
        @since 3.8
        @jira_ticket PYTHON-692
        @expected_result queries should execute appropriately

        @test_category object_mapper
        """
        connection_name = 'from_hosts'
        conn.register_connection(connection_name, hosts=[CASSANDRA_IP])
        self.assertIsNotNone(conn.get_connection(connection_name).cluster.metadata.get_host(CASSANDRA_IP))
        self.addCleanup(conn.unregister_connection, connection_name)

    def test_connection_param_validation(self):
        """
        Test to validate that invalid parameter combinations for registering connections via session are not tolerated
        @since 3.8
        @jira_ticket PYTHON-649
        @expected_result queries should execute appropriately

        @test_category object_mapper
        """
        cluster = Cluster([CASSANDRA_IP])
        session = cluster.connect()
        with self.assertRaises(CQLEngineException):
            conn.register_connection("bad_coonection1", session=session, consistency="not_null")
        with self.assertRaises(CQLEngineException):
            conn.register_connection("bad_coonection2", session=session, lazy_connect="not_null")
        with self.assertRaises(CQLEngineException):
            conn.register_connection("bad_coonection3", session=session, retry_connect="not_null")
        with self.assertRaises(CQLEngineException):
            conn.register_connection("bad_coonection4", session=session, cluster_options="not_null")
        with self.assertRaises(CQLEngineException):
            conn.register_connection("bad_coonection5", hosts="not_null", session=session)
        cluster.shutdown()

        cluster.shutdown()


        cluster.shutdown()

class BatchQueryConnectionTests(BaseCassEngTestCase):

    conns = ['cluster']

    @classmethod
    def setUpClass(cls):
        super(BatchQueryConnectionTests, cls).setUpClass()

        create_keyspace_simple('ks1', 1)
        sync_table(TestModel)
        sync_table(AnotherTestModel)

        conn.unregister_connection('default')
        conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True)
        conn.register_connection('cluster', [CASSANDRA_IP])

    @classmethod
    def tearDownClass(cls):
        super(BatchQueryConnectionTests, cls).tearDownClass()

        # reset the default connection
        conn.unregister_connection('fake_cluster')
        conn.unregister_connection('cluster')
        setup_connection(DEFAULT_KEYSPACE)

        drop_keyspace('ks1')

    def setUp(self):
        super(BaseCassEngTestCase, self).setUp()

    def test_basic_batch_query(self):
        """
        Test Batch queries with connections explicitly set

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result queries should execute appropriately

        @test_category object_mapper
        """

        # No connection with a QuerySet (default is a fake one)
        with self.assertRaises(NoHostAvailable):
            with BatchQuery() as b:
                TestModel.objects.batch(b).create(partition=1, cluster=1)

        # Explicit connection with a QuerySet
        with BatchQuery(connection='cluster') as b:
            TestModel.objects.batch(b).create(partition=1, cluster=1)

        # Get an object from the BD
        with ContextQuery(TestModel, connection='cluster') as tm:
            obj = tm.objects.get(partition=1, cluster=1)
            obj.__connection__ = None

        # No connection with a model (default is a fake one)
        with self.assertRaises(NoHostAvailable):
            with BatchQuery() as b:
                obj.count = 2
                obj.batch(b).save()

        # Explicit connection with a model
        with BatchQuery(connection='cluster') as b:
            obj.count = 2
            obj.batch(b).save()

    def test_batch_query_different_connection(self):
        """
        Test BatchQuery with Models that have a different connection

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result queries should execute appropriately

        @test_category object_mapper
        """

        # Testing on a model class
        TestModel.__connection__ = 'cluster'
        AnotherTestModel.__connection__ = 'cluster2'

        with self.assertRaises(CQLEngineException):
            with BatchQuery() as b:
                TestModel.objects.batch(b).create(partition=1, cluster=1)
                AnotherTestModel.objects.batch(b).create(partition=1, cluster=1)

        TestModel.__connection__ = None
        AnotherTestModel.__connection__ = None

        with BatchQuery(connection='cluster') as b:
            TestModel.objects.batch(b).create(partition=1, cluster=1)
            AnotherTestModel.objects.batch(b).create(partition=1, cluster=1)

        # Testing on a model instance
        with ContextQuery(TestModel, AnotherTestModel, connection='cluster') as (tm, atm):
            obj1 = tm.objects.get(partition=1, cluster=1)
            obj2 = atm.objects.get(partition=1, cluster=1)

            obj1.__connection__ = 'cluster'
            obj2.__connection__ = 'cluster2'

            obj1.count = 4
            obj2.count = 4

        with self.assertRaises(CQLEngineException):
            with BatchQuery() as b:
                obj1.batch(b).save()
                obj2.batch(b).save()

    def test_batch_query_connection_override(self):
        """
        Test that we cannot override a BatchQuery connection per model

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result Proper exceptions should be raised

        @test_category object_mapper
        """

        with self.assertRaises(CQLEngineException):
            with BatchQuery(connection='cluster') as b:
                TestModel.batch(b).using(connection='test').save()

        with self.assertRaises(CQLEngineException):
            with BatchQuery(connection='cluster') as b:
                TestModel.using(connection='test').batch(b).save()

        with ContextQuery(TestModel, AnotherTestModel, connection='cluster') as (tm, atm):
            obj1 = tm.objects.get(partition=1, cluster=1)
            obj1.__connection__ = None

        with self.assertRaises(CQLEngineException):
            with BatchQuery(connection='cluster') as b:
                obj1.using(connection='test').batch(b).save()

        with self.assertRaises(CQLEngineException):
            with BatchQuery(connection='cluster') as b:
                obj1.batch(b).using(connection='test').save()

class UsingDescriptorTests(BaseCassEngTestCase):

    conns = ['cluster']
    keyspaces = ['ks1', 'ks2']

    @classmethod
    def setUpClass(cls):
        super(UsingDescriptorTests, cls).setUpClass()

        conn.unregister_connection('default')
        conn.register_connection('fake_cluster', ['127.0.0.100'], lazy_connect=True, retry_connect=True, default=True)
        conn.register_connection('cluster', [CASSANDRA_IP])

    @classmethod
    def tearDownClass(cls):
        super(UsingDescriptorTests, cls).tearDownClass()

        # reset the default connection
        conn.unregister_connection('fake_cluster')
        conn.unregister_connection('cluster')
        setup_connection(DEFAULT_KEYSPACE)

        for ks in cls.keyspaces:
            drop_keyspace(ks)

    def setUp(self):
        super(BaseCassEngTestCase, self).setUp()

    def _reset_data(self):

        for ks in self.keyspaces:
            drop_keyspace(ks, connections=self.conns)
        for ks in self.keyspaces:
            create_keyspace_simple(ks, 1, connections=self.conns)
        sync_table(TestModel, keyspaces=self.keyspaces, connections=self.conns)

    def test_keyspace(self):
        """
        Test keyspace segregation when same connection is used

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result Keyspace segration is honored

        @test_category object_mapper
        """
        self._reset_data()

        with ContextQuery(TestModel, connection='cluster') as tm:

            # keyspace Model class
            tm.objects.using(keyspace='ks2').create(partition=1, cluster=1)
            tm.objects.using(keyspace='ks2').create(partition=2, cluster=2)

            with self.assertRaises(TestModel.DoesNotExist):
                tm.objects.get(partition=1, cluster=1)  # default keyspace ks1
            obj1 = tm.objects.using(keyspace='ks2').get(partition=1, cluster=1)

            obj1.count = 2
            obj1.save()

        with self.assertRaises(NoHostAvailable):
            TestModel.objects.using(keyspace='ks2').get(partition=1, cluster=1)

        obj2 = TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=1, cluster=1)
        self.assertEqual(obj2.count, 2)

        # Update test
        TestModel.objects(partition=2, cluster=2).using(connection='cluster', keyspace='ks2').update(count=5)
        obj3 = TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=2, cluster=2)
        self.assertEqual(obj3.count, 5)

        TestModel.objects(partition=2, cluster=2).using(connection='cluster', keyspace='ks2').delete()
        with self.assertRaises(TestModel.DoesNotExist):
            TestModel.objects.using(connection='cluster', keyspace='ks2').get(partition=2, cluster=2)

    def test_connection(self):
        """
        Test basic connection functionality

        @since 3.7
        @jira_ticket PYTHON-613
        @expected_result proper connection should be used

        @test_category object_mapper
        """
        self._reset_data()

        # Model class
        with self.assertRaises(NoHostAvailable):
            TestModel.objects.create(partition=1, cluster=1)

        TestModel.objects.using(connection='cluster').create(partition=1, cluster=1)
        TestModel.objects(partition=1, cluster=1).using(connection='cluster').update(count=2)
        obj1 = TestModel.objects.using(connection='cluster').get(partition=1, cluster=1)
        self.assertEqual(obj1.count, 2)

        obj1.using(connection='cluster').update(count=5)
        obj1 = TestModel.objects.using(connection='cluster').get(partition=1, cluster=1)
        self.assertEqual(obj1.count, 5)

        obj1.using(connection='cluster').delete()
        with self.assertRaises(TestModel.DoesNotExist):
            TestModel.objects.using(connection='cluster').get(partition=1, cluster=1)


class ModelQuerySetNew(ModelQuerySet):
    def __init__(self, *args, **kwargs):
        super(ModelQuerySetNew, self).__init__(*args, **kwargs)
        self._connection = "cluster"

class BaseConnectionTestNoDefault(object):
    conns = ['cluster']

    @classmethod
    def setUpClass(cls):
        conn.register_connection('cluster', [CASSANDRA_IP])
        test_queryset.TestModel.__queryset__ = ModelQuerySetNew
        test_queryset.IndexedTestModel.__queryset__ = ModelQuerySetNew
        test_queryset.CustomIndexedTestModel.__queryset__ = ModelQuerySetNew
        test_queryset.IndexedCollectionsTestModel.__queryset__ = ModelQuerySetNew
        test_queryset.TestMultiClusteringModel.__queryset__ = ModelQuerySetNew

        super(BaseConnectionTestNoDefault, cls).setUpClass()
        conn.unregister_connection('default')

    @classmethod
    def tearDownClass(cls):
        conn.unregister_connection('cluster')
        setup_connection(DEFAULT_KEYSPACE)
        super(BaseConnectionTestNoDefault, cls).tearDownClass()
        # reset the default connection

    def setUp(self):
        super(BaseCassEngTestCase, self).setUp()


class TestQuerySetOperationConnection(BaseConnectionTestNoDefault, test_queryset.TestQuerySetOperation):
    """
    Execute test_queryset.TestQuerySetOperation using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetDistinctNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetDistinct):
    """
    Execute test_queryset.TestQuerySetDistinct using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetOrderingNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetOrdering):
    """
    Execute test_queryset.TestQuerySetOrdering using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetCountSelectionAndIterationNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetCountSelectionAndIteration):
    """
    Execute test_queryset.TestQuerySetOrdering using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetSlicingNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetSlicing):
    """
    Execute test_queryset.TestQuerySetOrdering using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetValidationNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetValidation):
    """
    Execute test_queryset.TestQuerySetOrdering using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestQuerySetDeleteNoDefault(BaseConnectionTestNoDefault, test_queryset.TestQuerySetDelete):
    """
    Execute test_queryset.TestQuerySetDelete using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestValuesListNoDefault(BaseConnectionTestNoDefault, test_queryset.TestValuesList):
    """
    Execute test_queryset.TestValuesList using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass


class TestObjectsPropertyNoDefault(BaseConnectionTestNoDefault, test_queryset.TestObjectsProperty):
    """
    Execute test_queryset.TestObjectsProperty using non default connection

    @since 3.7
    @jira_ticket PYTHON-613
    @expected_result proper connection should be used

    @test_category object_mapper
    """
    pass
