import psycopg2 as _psycopg2
from psycopg2.extras import DictCursor as _cursor_factory

from speaker import Speaker
from commons import OKI_COL, INF_COL, ERR_COL
from eccez import QueryError

"""
The Mother DB engine.
"""

def _connectStr(loc):

    dbhost= loc['DB_HOST']
    dbname= loc['DB_NAME']
    dbpasswd= loc['DB_PASSWD']
    dbuser= loc['DB_USER']

    # Unix Sockets?
    if not dbhost:
        return "dbname=%s user=%s password=%s" % (dbname, dbuser, dbpasswd)

    dbport= loc.get('DB_PORT', 5432)

    return ( "dbname=%s user=%s host=%s password=%s port=%d" 
            % (dbname, dbuser, dbhost, dbpasswd, dbport) )

class Psygres(Speaker):

    _db_initialized= False
    _connect_str= None
    trans_level= 0
    pg_conn= None
    pg_cursor= None

    @staticmethod
    def _connect(pg_conn= None):

        if Psygres._db_initialized:
            Speaker.log_info("Psygres already connected.")
            return

        if pg_conn:
            Speaker.log_info(
                    "Initializing Psygres with own db connection.")
            Psygres.pg_conn= pg_conn
            try:
                # Zope WorkAround
                Psygres.pg_cursor= pg_conn.getcursor()
            except:
                Psygres.pg_cursor= pg_conn.cursor()
        else:
            Psygres.pg_conn= _psycopg2.connect(Psygres._connect_str)
            Psygres.pg_cursor= Psygres.pg_conn.cursor(
                    cursor_factory= _cursor_factory)

        Psygres._db_initialized= True

        return

    @staticmethod
    def beginTrans():
        Psygres.trans_level+= 1
        Speaker.log_debug(
                'Incremented transaction level: %s', 
                INF_COL(Psygres.trans_level))

    @staticmethod
    def rollback():
        if not Psygres.trans_level:
            Speaker.log_debug(
                    "Nothing to rollback: "
                    "nested rollback?")
            return

        Psygres.pg_conn.rollback()
        Psygres.trans_level= 0

        return

    @staticmethod
    def commit():
        lvl= Psygres.trans_level
        if not lvl:
            Speaker.log_error(
                    "Nothing to commit: "
                    "nested commit?")
            return

        if lvl== 1:
            Psygres.pg_conn.commit()
            Psygres.trans_level-= 1
            Speaker.log_debug(
                    "Queries committed.")
            return

        Psygres.trans_level-= 1
        Speaker.log_debug(
                "Decremented transaction: now %s",
                INF_COL(Psygres.trans_level))
        return

    @staticmethod
    def mogrify(s, d):
        return Psygres.pg_cursor.mogrify(s, d)

    @staticmethod
    def lastrowid():
        return Psygres.pg_cursor.lastrowid


    @staticmethod
    def _get_query(s):

        Speaker.log_info("QSQL- %s", s)
        cursor= Psygres.pg_cursor

        try:
            cursor.execute(s)
            res= cursor.fetchall()

        except Exception, es:
            Psygres.pg_conn.rollback()
            Speaker.log_raise(
                    "Unable to execute query: %s" % 
                    ERR_COL(es), QueryError)

        dict_list= []
        for r in res:
            d={}
            d.update(r)
            dict_list.append(d)
        return dict_list

    @staticmethod
    def _quiet_query(s):

        Speaker.log_info("QSQL- %s", s)
        cursor= Psygres.pg_cursor

        try:
            cursor.execute(s)
            if not Psygres.trans_level:
                Psygres.pg_conn.commit()
            else:
                Speaker.log_debug(
                        "Inside Trans: query not committed.")

        except Exception, es:
            Psygres.pg_conn.rollback()
            Speaker.log_raise(
                    "Unable to execute query: %s" % 
                    ERR_COL(es), QueryError)

    @staticmethod
    def oc_query(s):
        """ One Commit Query: no returns."""
        Psygres._quiet_query(s)

    @staticmethod
    def mr_query(s):
        """ Multiple Records Query: return a dict list."""
        return Psygres._get_query(s)

    @staticmethod
    def or_query(s):
        """ One Record Query: returns a dict."""

        res= Psygres.mr_query(s)
        
        if len(res)<> 1:
            Speaker.log_raise(
                    "or_query() returned %s records." % 
                    ERR_COL(len(res)), QueryError)

        return res[0]

    @staticmethod
    def ov_query(s):
        """ One Value Query: returns one only value."""
        try:
            res= Psygres.or_query(s)
        except:
            Speaker.log_raise(
                    "ov_query() returned %s records." % 
                    ERR_COL(len(res)), QueryError)

        res= res.values()
        if len(res)<> 1:
            Speaker.log_raise(
                    "ov_query() returned %s values." % 
                    ERR_COL(len(res)), QueryError)
        return res[0]



class PsygresFly(Speaker):

    _default_session= 'PsyFly'

    exported_methods= [
            'mogrify',
            'lastrowid',
            'oc_query',
            'ov_query',
            'mr_query',
            'or_query',
            'beginTrans',
            'commit',
            'rollback',
            'endSession'
            ]
    
    def __init__(self, pg_conn= None, name= None):

        self.session_name= name or self._default_session
        self._connect(pg_conn)

    def _connect(self, pg_conn= None):

        if pg_conn:
            Speaker.log_info(
                    "Initializing PsygresFly with own db connection.")
            self.pg_conn= pg_conn
            try:
                # Zope WorkAround
                self.pg_cursor= pg_conn.getcursor()
            except:
                self.pg_cursor= pg_conn.cursor()
        else:
            self.pg_conn= _psycopg2.connect(Psygres._connect_str)
            self.pg_cursor= self.pg_conn.cursor(
                    cursor_factory= _cursor_factory)

        return

    def _import(self, cls):
        for m in PsygresFly.exported_methods:
            setattr(cls, m, getattr(self, m))

    def mogrify(self, s, d):
        return self.pg_cursor.mogrify(s, d)

    def lastrowid(self):
        return self.pg_cursor.lastrowid


    def _get_query(self, s):

        Speaker.log_info("QSQL (%s) - %s", 
                INF_COL(self.session_name), s)
        cursor= self.pg_cursor

        try:
            cursor.execute(s)
            res= cursor.fetchall()

        except Exception, es:
            Speaker.log_raise(
                    "Unable to execute query: %s" % 
                    ERR_COL(es), QueryError)

        dict_list= []
        for r in res:
            d={}
            d.update(r)
            dict_list.append(d)
        return dict_list

    def _quiet_query(self, s):

        Speaker.log_info("QSQL (%s) - %s", 
                INF_COL(self.session_name), s)
        cursor= self.pg_cursor

        try:
            cursor.execute(s)

        except Exception, es:
            Speaker.log_raise(
                    "Unable to execute query: %s" % 
                    ERR_COL(es), QueryError)

    def oc_query(self, s):
        """ One Commit Query: no returns."""
        self._quiet_query(s)

    def mr_query(self, s):
        """ Multiple Records Query: return a dict list."""
        return self._get_query(s)

    def or_query(self, s):
        """ One Record Query: returns a dict."""

        res= self.mr_query(s)
        
        if len(res)<> 1:
            Speaker.log_raise(
                    "or_query() returned %s records." % 
                    ERR_COL(len(res)), QueryError)

        return res[0]

    def ov_query(self, s):
        """ One Value Query: returns one only value."""
        
        res= self.or_query(s)

        if len(res)<> 1:
            Speaker.log_raise(
                    "ov_query() returned %s values.",
                    ERR_COL(len(res)), QueryError)

        return res.values()[0]
    
    def beginTrans(self):
        self.log_noise(
                "Transactions is the default "
                "inside Sessions (session= %s).", 
                INF_COL(self.session_name))

    def endSession(self):
        self.log_info(
                "Terminating Session %s", 
                OKI_COL(self.session_name))

        if not hasattr(self, '_pool_queue'):
            self.log_insane("Session %s dies here.", 
                    OKI_COL(self.session_name))
            self.pg_conn.commit()
            return

        res= self.commit()

        if res<>0:
            PsygresPool.discard(self)
            self.log_int_raise(
                    "Session %s, unable to commit queries(): %s. "\
                    "Broken session removd from the pool.",
                    ERR_COL(self.session_name), ERR_COL(res))

        PsygresPool.backHome(self)
        self.log_debug(
                "Session %s back to the Pool.", 
                OKI_COL(self.session_name))

    def rollback(self):
        self.log_debug("Rollbacking queries for Session %s", 
                self.session_name)
        self.pg_conn.rollback()

    def commit(self):
        try:
            self.pg_conn.commit()
            self.log_debug("Queries Committed for Session %s.", 
                    OKI_COL(self.session_name))
            return 0
        except Exception, ss:
            return ss


class PsygresPool:

    import threading
    from Queue import Empty as _empty_queue

    _pool_initialized= False

    _pool_max= 10
    _pool_min= 4
    _pool_timeout= 15
    _pool_current= 0
    _pool_orphans= {}
    _pool_queue= None

    _pool_current_mutex= threading.Lock()
    _pool_get_mutex= threading.Lock()
    _pool_orphans_mutex= threading.Lock()

    del threading

    @staticmethod
    def _add_orphan(sname, db):
        m= PsygresPool._pool_orphans_mutex
        m.acquire()
        PsygresPool._pool_orphans[id(db)]= sname
        m.release()

    @staticmethod
    def _del_orphan(db):
        m= PsygresPool._pool_orphans_mutex
        m.acquire()
        PsygresPool._pool_orphans.pop(id(db))
        m.release()

    @staticmethod
    def _full():
        m= PsygresPool._pool_current_mutex
        m.acquire()
        res= PsygresPool._pool_current== PsygresPool._pool_max
        m.release()
        return res

    @staticmethod
    def status():
        """ status() -> (available, total, min, max, orphaned)
        """
        mc= PsygresPool._pool_current_mutex
        mg= PsygresPool._pool_get_mutex
        mo= PsygresPool._pool_orphans_mutex
        mc.acquire()
        mg.acquire()
        mo.acquire()
        total= PsygresPool._pool_current
        available= PsygresPool._pool_queue.qsize()
        orphaned= PsygresPool._pool_orphans.values()
        mo.release()
        mg.release()
        mc.release()
        max= PsygresPool._pool_max
        min= PsygresPool._pool_min
        return (available, total, min, max, orphaned)

    @staticmethod
    def _add_conn(n=1, pg_conn=None):
        """ Never call me directly: use _addConnection!"""
        max= PsygresPool._pool_max
        cur= PsygresPool._pool_current
        dif= max- cur

        if not dif:
            Speaker.log_warning(
                    "No space left in PsygresPool "
                    "(max= %s)", ERR_COL(max))
            return dif

        elif dif< n:
            Speaker.log_warning(
                    "Adding only %s new connection(s) "
                    "(max= %s, cur= %s, req= %s)",
                    OKI_COL(dif), ERR_COL(max), 
                    INF_COL(cur), INF_COL(n))

        else:
            dif= n
            Speaker.log_info(
                    "Adding %s new connection(s) (max= %s, cur= %s)",
                    OKI_COL(dif), OKI_COL(max), OKI_COL(cur))

        for i in xrange(dif):
            p= PsygresFly(pg_conn)
            q= PsygresPool._pool_queue
            p._pool_queue= q
            q.put(p)

        PsygresPool._pool_current+= dif

        return dif

    ##
    # Note that this function could be called by:
    #  init_pool() -> no problem with mutexes
    #  newSession() -> no problem with mutexes, coz
    #                  lock is acquired.
    ##
    @staticmethod
    def _addConnection(n=1, pg_conn= None):
        """ 
        Don't call me directly, use newSession() instead!
        """
        m= PsygresPool._pool_current_mutex
        m.acquire()
        PsygresPool._add_conn(n, pg_conn)
        m.release()


    @staticmethod
    def _init_pool(n= None, pg_conn= None):

        n= n or PsygresPool._pool_min

        Speaker.log_info("Initializing connection Pool ...")
        from Queue import Queue
        PsygresPool._pool_queue= Queue()
        PsygresPool._addConnection(n, pg_conn)

        PsygresPool._pool_initialized= True

    @staticmethod
    def _get_session():
        """ Never call me directly: use newSession!"""

        db_pool= PsygresPool._pool_queue

        try:
            db= db_pool.get_nowait()
            return db
        except PsygresPool._empty_queue:
            pass

        if PsygresPool._full():
            return db_pool.get(True, PsygresPool._pool_timeout)

        PsygresPool._addConnection()

        return db_pool.get_nowait()


    @staticmethod
    def newSession(name= None, pg_conn= None):

        Speaker.log_info("Initializing session %s", INF_COL(name))

        if not PsygresPool._pool_initialized:
            return PsygresFly(pg_conn, name)

        m= PsygresPool._pool_get_mutex
        m.acquire()

        try:
            db= PsygresPool._get_session()

        except Exception, ss:
            m.release()
            Speaker.log_int_raise(
                    "Cannot retrieve Session from Pool (%s). FATAL.", 
                    ERR_COL(ss))

        sname= name or PsygresFly._default_session

        PsygresPool._add_orphan(sname, db)

        m.release()
        db.session_name= sname 

        return db

    @staticmethod
    def backHome(pg_conn):

        q= PsygresPool._pool_queue
        m= PsygresPool._pool_current_mutex
        sname= pg_conn.session_name

        m.acquire()

        try:
            q.put_nowait(pg_conn)
        except Exception, ss:
            PsygresPool._pool_current-= 1
            Speaker.log_warning(
                    "Removed connection %s from Pool: %s.", 
                    ERR_COL(pg_conn.session_name), ERR_COL(ss))

        PsygresPool._del_orphan(pg_conn)

        m.release()


    @staticmethod
    def discard(pg_conn):

        m= PsygresPool._pool_current_mutex

        m.acquire()
        PsygresPool._pool_current-= 1
        PsygresPool._del_orphan(pg_conn)
        m.release()
        Speaker.log_warning("Connection %s dropped from Pool.", 
                ERR_COL(pg_conn.session_name))


def init_dbda(conf, pg_conn= None):

    if not hasattr(Psygres, '_db_initialized'):
        Speaker.log_info("Dbda already initialized.")
        return

    err= None
    if not isinstance(conf, dict):
        import speaker
        loc= {}
        try:
            execfile(conf, speaker.__dict__, loc)
        except Exception, ss:
            err= "Unable to read Mother configuration file %s: %s" % (ERR_COL(conf), ss)
    else:
        loc= conf

    if hasattr(Speaker, '_spkr_initialized'):
        from speaker import init_speaker
        init_speaker(loc)

    if err:
        Speaker.log_int_raise(err)

    Psygres._connect_str= _connectStr(loc)
    Psygres._connect(pg_conn)

    pool= loc.get('DB_POOL', False)

    if pool:
        pool_min= loc.get('DB_POOL_MIN', PsygresPool._pool_min)
        pool_max= loc.get('DB_POOL_MAX', PsygresPool._pool_max)
        pool_to= loc.get('DB_POOL_MIN', PsygresPool._pool_timeout)

        if pool_min> pool_max:
            Speaker.log_int_raise("Invalid configuration file: %s",
                    ERR_COL("DB_POOL_MIN> DB_POOL_MAX !!!"))

        PsygresPool._pool_min= pool_min
        PsygresPool._pool_max= pool_max
        PsygresPool._pool_timeout= pool_to

        PsygresPool._init_pool()

    del Psygres._db_initialized

__all__ = [
        'Psygres',
        'PsygresFly',
        'PsygresPool',
        'init_dbda'
        ]


