# Copyright (C) 2019 Apple Inc. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1.  Redistributions of source code must retain the above copyright
#     notice, this list of conditions and the following disclaimer.
# 2.  Redistributions in binary form must reproduce the above copyright
#     notice, this list of conditions and the following disclaimer in the
#     documentation and/or other materials provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY APPLE INC. AND ITS CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL APPLE INC. OR ITS CONTRIBUTORS BE LIABLE FOR
# ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

import mock
import re
import time

from cassandra.metadata import ColumnMetadata, KeyspaceMetadata, Metadata, TableMetadataV3
from cassandra.util import OrderedDict
from collections import defaultdict
from resultsdbpy.model.cassandra_context import CassandraContext


class MockCluster(object):

    metadata = Metadata()
    database = {}

    def __init__(self, contact_points=["127.0.0.1"], port=9042, auth_provider=None):
        assert isinstance(contact_points, list)
        self.contact_points = contact_points
        self.port = port
        self.sessions = []

    def connect(self, keyspace=None):
        result = MockSession(self, keyspace)
        self.sessions.append(result)
        return result

    def shutdown(self):
        self.sessions = []


class MockSession(object):

    def __init__(self, cluster, keyspace=None):
        self.cluster = cluster
        self.keyspace = keyspace

    def execute(self, *args):
        pass


class MockCQLEngineContext():

    connections = {}

    def __init__(self):
        self._depth = 0
        self.patches = [
            mock.patch('resultsdbpy.model.cassandra_context.Cluster', new=MockCluster),
            mock.patch('resultsdbpy.model.cassandra_context.register_connection', new=self.register_connection),
            mock.patch('resultsdbpy.model.cassandra_context.unregister_connection', new=self.unregister_connection),
            mock.patch('resultsdbpy.model.cassandra_context.get_cluster', new=self.get_cluster),
            mock.patch('resultsdbpy.model.cassandra_context.create_keyspace_simple', new=self.create_keyspace_simple),
            mock.patch('resultsdbpy.model.cassandra_context.create_keyspace_network_topology', new=self.create_keyspace_network_topology),
            mock.patch('resultsdbpy.model.cassandra_context.drop_keyspace', new=self.drop_keyspace),
            mock.patch('resultsdbpy.model.cassandra_context.sync_table', new=self.sync_table),
        ]

    def __enter__(self):
        if self._depth == 0:
            for patch in self.patches:
                patch.__enter__()
        self._depth += 1

    def __exit__(self, *args, **kwargs):
        self._depth -= 1
        if self._depth <= 0:
            for patch in self.patches:
                patch.__exit__(*args, **kwargs)

    @classmethod
    def register_connection(cls, name, session, **kwargs):
        if session is None:
            session = MockCluster().connect()
        cls.connections[name] = (session.cluster, session)

    @classmethod
    def unregister_connection(cls, name):
        if name in cls.connections:
            del cls.connections[name]

    @classmethod
    def get_cluster(cls, connection):
        return cls.connections.get(connection, [None])[0]

    @classmethod
    def create_keyspace_simple(cls, name, replication_factor=1, durable_writes=True, connections=[]):
        for connection in connections:
            assert connection in cls.connections
            if name not in cls.get_cluster(connection).metadata.keyspaces:
                cls.get_cluster(connection).metadata.keyspaces[name] = KeyspaceMetadata(name, durable_writes, None, None)
                MockCluster.database[name] = {}

    @classmethod
    def create_keyspace_network_topology(cls, name, dc_replication_map={}, durable_writes=True, connections=[]):
        return cls.create_keyspace_simple(name, durable_writes=durable_writes, connections=connections)

    @classmethod
    def drop_keyspace(cls, name, connections=[]):
        for connection in connections:
            assert connection in cls.connections
            if name in cls.get_cluster(connection).metadata.keyspaces:
                del cls.get_cluster(connection).metadata.keyspaces[name]
                del MockCluster.database[name]

    @classmethod
    def sync_table(cls, model, keyspaces=[], connections=[]):
        for connection in connections:
            assert connection in cls.connections

            for keyspace in keyspaces:
                keyspace_metadata = MockCQLEngineContext.get_cluster(connection).metadata.keyspaces.get(keyspace)
                assert keyspace_metadata is not None

                primary_keys = []
                column_keys = []
                for key, value in model._columns.items():
                    if value.partition_key or value.primary_key:
                        primary_keys.append(key)
                    else:
                        column_keys.append(key)

                columns = OrderedDict()
                partition_keys = []
                clustering_keys = []
                for key in primary_keys + column_keys:
                    value = model._columns[key]
                    meta_data = ColumnMetadata(None, key, value.db_type, is_reversed=(value.clustering_order == 'DESC'))
                    columns[key] = meta_data

                    if value.partition_key:
                        partition_keys.append(meta_data)
                    if value.primary_key:
                        clustering_keys.append(meta_data)

                MockCluster.database[keyspace][model._raw_column_family_name()] = defaultdict(lambda: {})
                keyspace_metadata.tables[model._raw_column_family_name()] = TableMetadataV3(
                    keyspace, model._raw_column_family_name(),
                    columns=columns, partition_key=partition_keys, clustering_key=clustering_keys,
                )


class MockCassandraContext(CassandraContext):

    @classmethod
    def drop_keyspace(cls, *args, **kwargs):
        with MockCQLEngineContext():
            CassandraContext.drop_keyspace(*args, **kwargs)

    def __init__(self, *args, **kwargs):
        self._mock_cql_engine = MockCQLEngineContext()
        with self._mock_cql_engine:
            super(MockCassandraContext, self).__init__(*args, **kwargs)

    def __enter__(self):
        self._mock_cql_engine.__enter__()
        super(MockCassandraContext, self).__enter__()

    def __exit__(self, *args, **kwargs):
        super(MockCassandraContext, self).__exit__(*args, **kwargs)
        self._mock_cql_engine.__exit__(*args, **kwargs)

    @CassandraContext.AssertConnectedDecorator()
    def insert_row(self, table_name, ttl=None, **kwargs):
        if table_name not in self._models:
            raise self.SchemaException(f'{table_name} does not exist in the database')

        item = self._models[table_name](**kwargs)
        schema = self.schema_for_table(table_name)

        partition = []
        for column in schema.partition_key:
            partition.append(getattr(item, column.name))
        cluster = []
        for column in schema.primary_key:
            if getattr(item, column.name) not in partition:
                cluster.append(getattr(item, column.name))

        self.cluster.database[self.keyspace][table_name][tuple(partition)][tuple(cluster)] = (None if ttl is None else time.time() + ttl, item)

    @CassandraContext.AssertConnectedDecorator()
    def select_from_table(self, table_name, limit=10000, **kwargs):
        result = []
        schema = self.schema_for_table(table_name)

        partition_names = []
        for column in schema.partition_key:
            partition_names.append(column.name)

        # Sort kwargs so they are in the same order as the columns.
        ordered_kwargs = OrderedDict()
        for name in self._models[table_name]._columns.keys():
            for key, value in kwargs.items():
                if key.split('__')[0] == name and value is not None:
                    ordered_kwargs[key] = value

        # Convert arguments to filters
        partitions = [[]]
        filters = []
        for key, value in ordered_kwargs.items():
            key_value = key.split('__')[0]
            operator = None if len(key.split('__')) == 1 else key.split('__')[1]

            if key_value in partition_names:
                if operator == 'in':
                    original_partitions = list(partitions)
                    partitions = []
                    for _ in range(len(value)):
                        for element in original_partitions:
                            partitions.append(list(element))
                    for multiplier in range(len(value)):
                        for index in range(len(original_partitions)):
                            partitions[multiplier * len(original_partitions) + index].append(value[multiplier])
                else:
                    for partition in partitions:
                        partition.append(value)

            filters.append(self.filter_for_argument(key, value))

        if len(partitions[0]) == 0:
            partitions = self.cluster.database[self.keyspace][table_name].keys()

        candidate_elements = []
        for partition in partitions:
            for key, row in self.cluster.database[self.keyspace][table_name][tuple(partition)].items():
                does_match = True
                for filter in filters:
                    if not filter(row[1]):
                        does_match = False
                        break
                if not does_match:
                    continue
                candidate_elements.append((partition, key))

        # To reverse, they all need to agree. This might not work with compound clustering keys, but we really
        # shouldn't use those anyways
        is_reversed = True
        for column in schema.primary_key:
            if column in schema.partition_key:
                continue
            if not column.is_reversed:
                is_reversed = False
        candidate_elements.sort(reverse=is_reversed)

        for candidate in candidate_elements:
            element = self.cluster.database[self.keyspace][table_name][tuple(candidate[0])][candidate[1]]
            deadline = element[0]
            if deadline is None or deadline > time.time():
                result.append(element[1])
                limit -= 1
                if limit <= 0:
                    return result
            else:
                del self.cluster.database[self.keyspace][table_name][tuple(candidate[0])][candidate[1]]
        return result