# coding: utf-8


from time import time

from sqlalchemy import create_engine
from sqlalchemy.orm import sessionmaker, scoped_session

from .models import Session, VmmasterLogStep, SessionLogStep

from ..utils.utils import to_thread


def threaded_transaction(func):
    @to_thread
    def wrapper(self, *args, **kwargs):
        session = sessionmaker(bind=self.engine)()
        try:
            return func(self, session=session, *args, **kwargs)
        except:
            session.rollback()
            raise
        finally:
            session.close()
    return wrapper


def transaction(func):
    def wrapper(self, *args, **kwargs):
        session = self.Session()
        try:
            return func(self, session=session, *args, **kwargs)
        except:
            session.rollback()
            raise
        finally:
            session.close()
    return wrapper


class Database(object):
    _instance = None

    def __new__(cls, *args, **kwargs):
        if not cls._instance:
            cls._instance = super(Database, cls).__new__(cls, *args, **kwargs)
        return cls._instance

    def __init__(self, connection_string):
        self.engine = create_engine(connection_string)
        self.Session = scoped_session(sessionmaker(bind=self.engine))
        from vmmaster import migrations
        migrations.run(connection_string)

    @transaction
    def create_session(self, status="running", name=None, time=time(), session=None):
        _session = Session(status=status, time=time)
        session.add(_session)
        session.commit()
        if name is None:
            _session.name = _session.id
        else:
            _session.name = name
        session.add(_session)
        session.commit()
        created_session = session.query(Session).filter_by(id=_session.id).first()
        session.flush()
        return created_session

    @transaction
    def get_session(self, session_id, session=None):
        return session.query(Session).filter_by(id=session_id).first()

    @transaction
    def create_vmmaster_log_step(self, session_id, control_line, body, screenshot="", time=time(), session=None):
        _log_step = VmmasterLogStep(session_id=session_id,
                                    control_line=control_line,
                                    body=body,
                                    screenshot=screenshot,
                                    time=time)
        session.add(_log_step)
        session.commit()
        created_log_step = session.query(VmmasterLogStep).filter_by(id=_log_step.id).first()
        session.flush()
        return created_log_step

    @transaction
    def get_vmmaster_log_step(self, log_step_id, session=None):
        return session.query(VmmasterLogStep).filter_by(id=log_step_id).first()

    @transaction
    def create_session_log_step(self, vmmaster_log_step_id, control_line, body, time=time(), session=None):
        _log_step = SessionLogStep(vmmaster_log_step_id=vmmaster_log_step_id,
                                   control_line=control_line,
                                   body=body,
                                   time=time)
        session.add(_log_step)
        session.commit()
        created_log_step = session.query(SessionLogStep).filter_by(id=_log_step.id).first()
        session.flush()
        return created_log_step

    @transaction
    def get_session_log_step(self, log_step_id, session=None):
        return session.query(SessionLogStep).filter_by(id=log_step_id).first()

    @transaction
    def update(self, obj, session=None):
        session.merge(obj)
        session.commit()
        updated_obj = session.query(type(obj)).filter_by(id=obj.id).first()
        session.flush()
        return updated_obj

database = None
@@ -629,32 +629,213 @@
 rgs, **kwargs):%0A
+        # Try to use passed database session%0A        if 'session' in kwargs.keys():%0A            if kwargs%5B'session'%5D is not None:%0A                return func(self, *args, **kwargs)%0A
         session 
@@ -851,16 +851,16 @@
 ssion()%0A
-
         
@@ -1944,61 +1944,51 @@
 = se
-ssion.query(Session).filter_by(id=_session.id).first(
+lf.get_session(_session.id, session=session
 )%0A  
@@ -2149,29 +2149,20 @@
 ession).
-filter_by(id=
+get(
 session_
@@ -2160,32 +2160,24 @@
 (session_id)
-.first()
 %0A%0A    @trans
@@ -2661,70 +2661,62 @@
 = se
-ssion.query(VmmasterLogStep).filter_by(id=_log_step.id).first(
+lf.get_vmmaster_log_step(_log_step.id, session=session
 )%0A  
@@ -2893,37 +2893,28 @@
 terLogStep).
-filter_by(id=
+get(
 log_step_id)
@@ -2909,32 +2909,24 @@
 log_step_id)
-.first()
 %0A%0A    @trans
@@ -3361,69 +3361,61 @@
 = se
-ssion.query(SessionLogStep).filter_by(id=_log_step.id).first(
+lf.get_session_log_step(_log_step.id, session=session
 )%0A  
@@ -3594,29 +3594,20 @@
 ogStep).
-filter_by(id=
+get(
 log_step
@@ -3610,24 +3610,16 @@
 step_id)
-.first()
 %0A%0A    @t
@@ -3773,35 +3773,18 @@
 j)).
-filter_by(id=obj.id).first(
+get(obj.id
 )%0A  

Enhance database layer