"""
CQLEngine-Session

Your cqlengine model must inherit from cqlengine_session.SessionModel instead
of cqlengine.model.Model

SessionModel will replace youl SessionModel declarations with classes of type
IdMapModel.  Your model module will get classes of type Model with an
underscore prefixed to the name.

example:
class Foo(SessionModel):
    pass

results in Foo being a IdMapModel, and _Foo being a Model.

Note that making blind handles requires you pass a key.
blind = Foo(key)

you can make changes and save them (without first needing to load the object.)
blind.title = u'new title'
save()

To create new object use create
foo = Foo.create()

"""

import copy
from datetime import date, datetime
import importlib
import json
import threading
from uuid import UUID

from cassandra.query import SimpleStatement
from cqlengine import columns
import cqlengine.connection
from cqlengine.exceptions import ValidationError
from cqlengine.management import get_fields, sync_table
from cqlengine.models import ColumnQueryEvaluator, Model, ModelMetaClass
from cqlengine.operators import EqualsOperator
from cqlengine.query import BatchQuery, ModelQuerySet
from cqlengine.statements import WhereClause, SelectStatement, DeleteStatement, UpdateStatement, AssignmentClause, InsertStatement, BaseCQLStatement, MapUpdateClause, MapDeleteClause, ListUpdateClause, SetUpdateClause, CounterUpdateClause


class AttributeUnavailable(Exception):
    pass


class SessionManager(object):
    def get_session(self):
        """Return current session for this context."""
        raise NotImplementedError

    def set_session(self, session):
        """Make the given session the current session for this context."""
        raise NotImplementedError


class ThreadLocalSessionManager(SessionManager):
    def __init__(self):
        self.storage = threading.local()

    def get_session(self):
        return getattr(self.storage, 'session', None)

    def set_session(self, session):
        self.storage.session = session


SESSION_MANAGER = ThreadLocalSessionManager()


def set_session_manager(manager):
    global SESSION_MANAGER
    SESSION_MANAGER = manager


def clear():
    """Empty the current session"""
    # xxx what happens to the existing id-map objects?  this is dangerous.
    # (also, the dev is not expected to call this.)
    SESSION_MANAGER.set_session(None)


def save(*objects):
    "Write all pending changes from session to Cassandra."
    session = SESSION_MANAGER.get_session()
    if session is not None:
        session.save(*objects)


def get_session(create_if_missing=True):
    session = SESSION_MANAGER.get_session()
    if session is None:
        session = Session()
        SESSION_MANAGER.set_session(session)
    return session


def add_call_after_save(callable, *args, **kwargs):
    """Call callable with given args and kwargs after next save."""
    get_session().call_after_save.append((callable, args, kwargs,))


class Session(object):
    """Identity map objects and support for implicit batch save."""
    def __init__(self):
        self.instances_by_class = {}
        self.call_after_save = []
        #self.deletes = set()

    def save(self, *objects):
        """Flush all pending changes to Cassandra.

        objects -- if not None, only operate on this or these object(s)

        """
        updates = set()
        counter_updates = set()
        creates = set()
        counter_creates = set()
        for model_class, by_key in self.instances_by_class.iteritems():
            for key, instance in by_key.iteritems():
                if hasattr(instance, '_created') and instance._created:
                    if model_class.id_mapped_class._has_counter:
                        counter_creates.add(instance)
                    else:
                        creates.add(instance)
                elif hasattr(instance, '_dirties'):
                    if model_class.id_mapped_class._has_counter:
                        counter_updates.add(instance)
                    else:
                        updates.add(instance)
        if objects:
            updates = updates and objects
            counter_updates = counter_updates and objects
            creates = creates and objects
            counter_creates = counter_creates and objects

        with BatchQuery() as batch:
            for create in creates:
                # Note we skip a lot of cqlengine code and create the
                # insert statement directly.
                # (this is the non-optimized code that is replaced below)
                #key_names = create.id_mapped_class._columns.keys()
                #arg = {name: getattr(create, name) for name in key_names}
                #create.id_mapped_class.batch(batch).create(**arg)
                # (end non-optimized code)
                # (begin optimized)
                # note: it might save time to memoize column family name
                # note: cqlengine-session doesn't yet support 'ttl'
                insert = InsertStatement(create.id_mapped_class.column_family_name())#, ttl=self._ttl)
                for name, col in create.id_mapped_class._columns.items():
                    val = col.validate(getattr(create, name))
                    if col._val_is_null(val):
                        continue
                    insert.add_assignment_clause(AssignmentClause(
                        col.db_field_name,
                        col.to_database(val)))
                # skip query execution if it's empty
                # caused by pointless update queries
                if not insert.is_empty:
                    batch.add_query(insert)
                # (end optimized)
                del create._created
                try:
                    del create._dirties
                except AttributeError:
                    pass
            for update in updates:
                key_names = update._primary_keys.keys()
                arg = {name: getattr(update, name) for name in key_names}
                dirties = update._dirties
                update.id_mapped_class.objects(**arg).batch(batch).update(**dirties)
                del update._dirties
        # It would seem that batch does not work with counter?
        #with BatchQuery() as batch:
        for create in counter_creates:
            primary_key_names = create.id_mapped_class._primary_keys.keys()
            arg = {name: getattr(create, name) for name in primary_key_names}
            instance = create.id_mapped_class.create(**arg)
            for name, col in create.id_mapped_class._columns.items():
                if isinstance(col, columns.Counter):
                    val = getattr(create, name)
                    setattr(instance, name, val)
            del create._created
            try:
                del create._dirties
            except AttributeError:
                pass
            instance.update()
        for update in counter_updates:
            statement = UpdateStatement(update.id_mapped_class.column_family_name())#, ttl=self._ttl)
            for name, value in update._dirties.items():
                col = update.id_mapped_class._columns[name]
                clause = CounterUpdateClause(col.db_field_name, value, 0, column=col)
                statement.add_assignment_clause(clause)
            for name, col in update.id_mapped_class._primary_keys.items():
                statement.add_where_clause(WhereClause(
                    col.db_field_name,
                    EqualsOperator(),
                    col.to_database(getattr(update, name))
                ))
            params = statement.get_context()
            statement = SimpleStatement(str(statement))
            cqlengine.connection.get_session().execute(statement, params)
            del update._dirties
#            for delete in self.deletes:
#                raise NotImplementedError
        for callable, args, kwargs in self.call_after_save:
            callable(*args, **kwargs)
        self.call_after_save = []


class SessionModelMetaClass(ModelMetaClass):

    def __new__(cls, name, bases, attrs):
        if attrs.get('__abstract__'):
            return super(SessionModelMetaClass, cls).__new__(cls,
                                                             name,
                                                             bases,
                                                             attrs)
        if len(bases) > 1:
            raise TypeError('SessionModel does not allow multiple inheritance')
        # Take the result of the base class's __new__ and assign it to the
        # module using a prefixed underscore in the name.
        new_name = '_' + name
        # Note: at this point attrs has only those actually declared in
        # the class declaration (and not in any parent class declaration)
        base = super(SessionModelMetaClass, cls).__new__(cls,
                                                         new_name,
                                                         bases,
                                                         attrs)
        # Note: at this point, attrs has had a bunch of things added by
        # cqlengine.models.ModelMetaClass
        module = importlib.import_module(cls.__module__)
        setattr(module, new_name, base)

        # Copy attrs from the base class because this class won't actually
        # inherit from these base classes.
        base_attrs = {}
        copyable_bases = []
        for klass in bases[0].mro():
            if klass == SessionModel:
                break
            copyable_bases.append(klass)
        for klass in reversed(copyable_bases):
            base_attrs.update(klass.__dict__)
        base_attrs.update(attrs)
        base_attrs['id_mapped_class'] = base
        base_attrs['_promotable_column_names'] = set([cname for cname, c in base_attrs['_columns'].iteritems() if not c.primary_key])
        # Make descriptors for the columns so the instances will get/set
        # using a ColumnDescriptor instance.
        for col_name, col in base._columns.iteritems():
            if isinstance(col, columns.Counter):
                base_attrs[col_name] = CounterColumnDescriptor(col)
            else:
                base_attrs[col_name] = ColumnDescriptor(col)
        return IdMapMetaClass(name, (IdMapModel,), base_attrs)


# declare your models with this so that SessionModelMetaClass is the metaclass.
class SessionModel(Model):
    __abstract__ = True
    __metaclass__ = SessionModelMetaClass

class IdMapMetaClass(type):

#    def __new__(cls, name, bases, attrs):
#        return None
#        return type(name, bases, attrs)

    def __call__(cls, *key):
        """If instance is in the id-map, return it, else make and return it."""
        session = get_session()
        try:
            instance_by_key = session.instances_by_class[cls]
            try:
                return instance_by_key[key]
            except KeyError:
                pass
        except KeyError:
            instance_by_key = {}
            session.instances_by_class[cls] = instance_by_key
        instance = super(IdMapMetaClass, cls).__call__(*key)
        instance_by_key[key] = instance
        return instance


# this is copied from cqlengine, may need more modification..
class QuerySetDescriptor(object):
    def __get__(self, instance, session_class):
        return WrappedQuerySet(instance, session_class)


class IdMapModel(object):

    __metaclass__ = IdMapMetaClass

    objects = QuerySetDescriptor()

    def __init__(self, *key):
        self.key = key
        key_names = self.id_mapped_class._primary_keys.keys()
        for name, value in zip(key_names, key):
            self._promote(name, value)

    @classmethod
    def all(cls):
        return cls.objects.all()

    @classmethod
    def filter(cls, *args, **kwargs):
        return cls.objects.filter(*args, **kwargs)

    @classmethod
    def get(cls, *args, **kwargs):
        return cls.objects.get(*args, **kwargs)

    @classmethod
    def create(cls, **kwargs):
        column_names = cls.id_mapped_class._columns.keys()
        extra_columns = set(kwargs.keys()) - set(column_names)
        if extra_columns:
            raise ValidationError(
                    "Incorrect columns passed: {}".format(extra_columns))

        primary_keys = cls.id_mapped_class._primary_keys
        uncleaned_values = {}
        for name, col in cls.id_mapped_class._columns.items():
            try:
                value = kwargs[name]
            except KeyError:
                if col.default:
                    if callable(col.default):
                        value = col.default()
                    else:
                        value = col.default
                elif isinstance(col, columns.Counter):
                    value = 0
                elif name in primary_keys:
                    raise ValueError(u"Can't create {} without providing primary key {}".format(cls.__name__, name))
                else:
                    # Container columns have non-None empty cases.
                    value = None
            uncleaned_values[name] = value

        key = []
        for name, col in primary_keys.items():
            key.append(col.to_python(uncleaned_values[name]))
        instance = cls(*key)
        instance._created = True
        for name, col in cls.id_mapped_class._columns.items():
            if name in primary_keys:
                continue
            value = uncleaned_values[name]
            if isinstance(col, columns.BaseContainerColumn):
                if isinstance(col, columns.Set):
                    value = OwnedSet(instance, name, col.to_python(value))
                elif isinstance(col, columns.List):
                    value = OwnedList(instance, name, col.to_python(value))
                elif isinstance(col, columns.Map):
                    value = OwnedMap(instance, name, col.to_python(value))
            elif value is not None:
                value = col.to_python(value)
            instance._promote(name, value)
        return instance

    def promote(self, **kwargs):
        """Set kwargs on entity without marking as dirty

        Invalid column names in kwargs raises an exception

        Promoting the value of a key raises an exception

        """
        extra_columns = set(kwargs.keys()) - self._promotable_column_names
        if extra_columns:
            raise ValidationError("Incorrect columns passed: {}".format(extra_columns))

        for col_name, col_value in kwargs.items():
            self._promote(col_name, col_value)

    def _promote(self, name, value):
        """set without marking attribute as dirty."""
        try:
            self._values[name] = value
        except AttributeError:
            self._values = {name: value}

    def _mark_dirty(self, name, value):
        """mark an attribute as dirty."""
        try:
            self._dirties[name] = value
        except AttributeError:
            self._dirties = {name: value}

    @classmethod
    def sync_table(cls):
        sync_table(cls.id_mapped_class)

    @classmethod
    def _construct_instance(cls, values):