import pickle
import uuid
from datetime import datetime
from datetime import timedelta

from werkzeug.datastructures import CallbackDict
from flask.sessions import SessionMixin
from flask.sessions import SessionInterface


class MongoDBSession(CallbackDict, SessionMixin):
    def __init__(self, initial=None, sid=None, new=True):
        def on_update(this):
            this.modified = True
        if initial:
            initial = pickle.loads(str(initial))
        CallbackDict.__init__(self, initial, on_update)
        self.sid = sid
        self.new = new
        self.modified = False

    def pickle(self):
        return pickle.dumps(dict(self))


class MongoDBSessionInterface(SessionInterface):
    session_class = MongoDBSession

    def __init__(self, app, db, collection_name):
        self._db = db
        self._collection_name = collection_name

        if app is not None:
            self.app = app
            self.init_app(app)
        else:
            self.app = None

    def init_app(self, app):
        if not hasattr(app, 'extensions'):
            app.extensions = {}
        app.extensions['mongodb-sessions'] = self

    def open_session(self, app, request):
        sid = request.cookies.get(app.session_cookie_name)
        if not sid:
            sid = self.__generate_sid()
            return self.session_class(sid=sid)

        doc = self.__get_collection().find_one({'_id': sid})
        # Check if exists and not expired.
        # It's ok to remove tzinfo here, because utcnow() returns UTC time,
        # despite it's in naive form (without tzinfo).
        if doc and doc['exp'].replace(tzinfo=None) > datetime.utcnow():
            session = self.session_class(initial=doc['d'], sid=sid)
        else:
            # If the SID doesn't exist - create a new one to avoid possibility
            # of user-generated SID with invalid format (e.g. "abc123").
            sid = self.__generate_sid()
            session = self.session_class(sid=sid)
        return session

    def save_session(self, app, session, response):
        # cookie_domain = self.get_cookie_domain(app)
        # cookie_path = self.get_cookie_path(app)
        cookie_exp = self.get_expiration_time(app, session)

        if not session:
            self.__get_collection().remove({'_id': session.sid})
            if session.modified:
                response.delete_cookie(key=app.session_cookie_name,
                                       #path=cookie_path,
                                       #domain=cookie_domain
                                       )
            return

        # If session isn't permanent if will be considered valid for 1 day
        # (but not cookie which will be deleted by browser after exit).
        session_exp = cookie_exp or datetime.utcnow()+timedelta(days=1)
        self.__get_collection().update(
            {'_id': session.sid},
            {'$set': {
                'd': session.pickle(),
                'exp': session_exp,
            }},
            upsert=True)

        response.set_cookie(key=app.session_cookie_name,
                            value=session.sid,
                            expires=cookie_exp,
                            #path=cookie_path,
                            #domain=cookie_domain,
                            secure=self.get_cookie_secure(app),
                            httponly=self.get_cookie_httponly(app))

    def __get_collection(self):
        return self._db[self._collection_name]

    def __generate_sid(self):
        return uuid.uuid4().hex
