# -*- coding: utf-8 -*-
"""This module holds a ConnectionWrapper that is used with a
   JDBC Connection. The module should only be used when running Jython.
"""

import java.sql as jdbc

from copy import copy as pcopy
from datetime import datetime
from sys import modules
from threading import Thread
from Queue import Queue

import pygrametl
from pygrametl.FIFODict import FIFODict

# NOTE: This module is made for Jython.

__author__ = "Christian Thomsen"
__maintainer__ = "Sílex Sistemas Ltda."
__version__ = '0.0.1'
__all__ = ['JDBCConnectionWrapper', 'BackgroundJDBCConnectionWrapper']


class JDBCConnectionWrapper(object):
    """Wrap a JDBC Connection.

       All Dimension and FactTable communicate with the data warehouse using 
       a ConnectionWrapper. In this way, the code for loading the DW does not
       have to care about which parameter format is used.  
       This ConnectionWrapper is a special one for JDBC in Jython.
    """

    def __init__(self, jdbcconn, stmtcachesize=20):
        """Create a ConnectionWrapper around the given JDBC connection.

           If no default ConnectionWrapper already exists, the new
           ConnectionWrapper is set to be the default ConnectionWrapper.

           Arguments:
           - jdbcconn: An open JDBC Connection (not a PEP249 Connection)
           - stmtcachesize: The maximum number of PreparedStatements kept
             open. Default: 20.
        """
        if not isinstance(jdbcconn, jdbc.Connection):
            raise TypeError, '1st argument must implement java.sql.Connection'
        if jdbcconn.isClosed():
            raise ValueError, '1st argument must be an open Connection'
        self.__jdbcconn = jdbcconn
        # Add a finalizer to __prepstmts to close PreparedStatements when
        # they are pushed out
        self.__prepstmts = FIFODict(stmtcachesize, lambda k, v: v[0].close())
        self.__resultmeta = FIFODict(stmtcachesize)
        self.__resultset = None
        self.__resultnames = None
        self.__resulttypes = None
        self.nametranslator = lambda s: s
        self.__jdbcconn.setAutoCommit(False)
        if pygrametl._defaulttargetconnection is None:
            pygrametl._defaulttargetconnection = self

    def __preparejdbcstmt(self, sql):
        # Find pyformat arguments and change them to question marks while
        # appending the attribute names to a list
        names = []
        newsql = sql
        while True:
            start = newsql.find('%(')
            if start == -1:
                break
            end = newsql.find(')s', start)
            if end == -1: 
                break
            name = newsql[start+2 : end]
            names.append(name)
            newsql = newsql.replace(newsql[start:end+2], '?', 1)

        ps = self.__jdbcconn.prepareStatement(newsql)

        # Find parameter types
        types = []
        parmeta = ps.getParameterMetaData()
        for i in range(len(names)):
            types.append(parmeta.getParameterType(i+1))

        self.__prepstmts[sql] = (ps, names, types)

    def __executejdbcstmt(self, sql, args):
        if self.__resultset:
            self.__resultset.close()

        if sql not in self.__prepstmts:
            self.__preparejdbcstmt(sql)
        (ps, names, types) = self.__prepstmts[sql]

        for pos in range(len(names)): # Not very Pythonic, but we're doing Java
            if args[names[pos]] is None:
                ps.setNull(pos + 1, types[pos])
            else:
                ps.setObject(pos + 1, args[names[pos]], types[pos])

        if ps.execute():
            self.__resultset = ps.getResultSet()
            if sql not in self.__resultmeta:
                self.__resultmeta[sql] = \
                    self.__extractresultmetadata(self.__resultset)
            (self.__resultnames, self.__resulttypes) = self.__resultmeta[sql]
        else:
            self.__resultset = None
            (self.__resultnames, self.__resulttypes) = (None, None)

    def __extractresultmetadata(self, resultset):
        # Get jdbc resultset metadata. extract names and types
        # and add it to self.__resultmeta
        meta = resultset.getMetaData()
        names = []
        types = []
        for col in range(meta.getColumnCount()):
            names.append(meta.getColumnName(col+1))
            types.append(meta.getColumnType(col+1))
        return (names, types)

    def __readresultrow(self):
        if self.__resultset is None:
            return None
        result = []
        for i in range(len(self.__resulttypes)):
            e = self.__resulttypes[i] # Not Pythonic, but we need i for JDBC
            if e in (jdbc.Types.CHAR, jdbc.Types.VARCHAR, 
                     jdbc.Types.LONGVARCHAR):
                result.append(self.__resultset.getString(i+1))
            elif e in (jdbc.Types.BIT, jdbc.Types.BOOLEAN):
                result.append(self.__resultset.getBool(i+1))
            elif e in (jdbc.Types.TINYINT, jdbc.Types.SMALLINT, 
                       jdbc.Types.INTEGER):
                result.append(self.__resultset.getInt(i+1))
            elif e in (jdbc.Types.BIGINT, ):
                result.append(self.__resultset.getLong(i+1))
            elif e in (jdbc.Types.DATE, ):
                result.append(self.__resultset.getDate(i+1))
            elif e in (jdbc.Types.TIMESTAMP, ):
                result.append(self.__resultset.getTimestamp(i+1))
            elif e in (jdbc.Types.TIME, ):
                result.append(self.__resultset.getTime(i+1))
            else:
                # Try this and hope for the best...
                result.append(self.__resultset.getString(i+1))
        return tuple(result)

    def execute(self, stmt, arguments=None, namemapping=None, ignored=None):
        """Execute a statement.

           Arguments:
           - stmt: the statement to execute
           - arguments: a mapping with the arguments (default: None)
           - namemapping: a mapping of names such that if stmt uses %(arg)s
             and namemapping[arg]=arg2, the value arguments[arg2] is used 
             instead of arguments[arg]
           - ignored: An ignored argument only present to accept the same 
             number of arguments as ConnectionWrapper.execute
        """
        if namemapping and arguments:
            arguments = pygrametl.copy(arguments, **namemapping)
        self.__executejdbcstmt(stmt, arguments)

    def executemany(self, stmt, params, ignored=None):
        """Execute a sequence of statements.

           Arguments:
           - stmt: the statement to execute
           - params: a sequence of arguments
           - ignored: An ignored argument only present to accept the same 
             number of arguments as ConnectionWrapper.executemany
        """
        for paramset in params:
            self.__executejdbcstmt(stmt, paramset)

    def rowfactory(self, names=None):
        """Return a generator object returning result rows (i.e. dicts)."""
        if names is None:
            if self.__resultnames is None:
                return
            else:
                names = [self.nametranslator(t[0]) for t in self.__resultnames]
        empty = (None, ) * len(self.__resultnames)
        while True:
            tuple = self.fetchonetuple()
            if tuple == empty:
                return
            yield dict(zip(names, tuple))

    def fetchone(self, names=None):
        """Return one result row (i.e. dict)."""
        if self.__resultset is None:
            return {}
        if names is None:
            names = [self.nametranslator(t[0]) for t in self.__resultnames]
        values = self.fetchonetuple()
        return dict(zip(names, values))

    def fetchonetuple(self):
        """Return one result tuple."""
        if self.__resultset is None:
            return ()
        if not self.__resultset.next():
            return (None, ) * len(self.__resultnames)
        else:
            return self.__readresultrow()

    def fetchmanytuples(self, cnt):
        """Return cnt result tuples."""
        if self.__resultset is None:
            return []
        empty = (None, ) * len(self.__resultnames)
        result = []
        for i in range(cnt):
            tmp = self.fetchonetuple()
            if tmp == empty:
                break
            result.append(tmp)
        return result

    def fetchalltuples(self):
        """Return all result tuples"""
        if self.__resultset is None:
            return []
        result = []
        empty = (None, ) * len(self.__resultnames)
        while True:
            tmp = self.fetchonetuple()
            if tmp == empty:
                return result
            result.append(tmp)

    def rowcount(self):
        """Not implemented. Return 0. Should return the size of the result."""
        return 0 

    def getunderlyingmodule(self):
        """Return a reference to the underlying connection's module."""
        return modules[self.__class__.__module__]

    def commit(self):
        """Commit the transaction."""
        pygrametl.endload()
        self.__jdbcconn.commit()

    def close(self):
        """Close the connection to the database,"""
        self.__jdbcconn.close()

    def rollback(self):
        """Rollback the transaction."""
        self.__jdbcconn.rollback()

    def setasdefault(self):
        """Set this ConnectionWrapper as the default connection."""
        pygrametl._defaulttargetconnection = self

    def cursor(self):
        """Not implemented for this JDBC connection wrapper!"""
        raise NotImplementedError, ".cursor() not supported"

    def resultnames(self):
        if self.__resultnames is None:
            return None
        else:
            return tuple(self.__resultnames)

# BackgroundJDBCConnectionWrapper is added for experiments. It is quite similar
# to JDBCConnectionWrapper and one of them may be removed.


class BackgroundJDBCConnectionWrapper(object):
    """Wrap a JDBC Connection and do all DB communication in the background.

       All Dimension and FactTable communicate with the data warehouse using 
       a ConnectionWrapper. In this way, the code for loading the DW does not
       have to care about which parameter format is used.
       This ConnectionWrapper is a special one for JDBC in Jython and does DB 
       communication from a Thread.
    """

    def __init__(self, jdbcconn, stmtcachesize=20):
        """Create a ConnectionWrapper around the given JDBC connection """
        self.__jdbcconn = jdbcconn
        # Add a finalizer to __prepstmts to close PreparedStatements when
        # they are pushed out
        self.__prepstmts = FIFODict(stmtcachesize, lambda k, v: v[0].close())
        self.__resultmeta = FIFODict(stmtcachesize)
        self.__resultset = None
        self.__resultnames = None
        self.__resulttypes = None
        self.nametranslator = lambda s: s
        self.__jdbcconn.setAutoCommit(False)
        self.__queue = Queue(5000)
        t = Thread(target=self.__worker)
        t.setDaemon(True)  # NB: "t.daemon = True" does NOT work...
        t.setName('BackgroundJDBCConnectionWrapper')
        t.start()

    def __worker(self):
        while True:
            (sql, args) = self.__queue.get()
            self.__executejdbcstmt(sql, args)
            self.__queue.task_done()

    def __preparejdbcstmt(self, sql):
        # Find pyformat arguments and change them to question marks while
        # appending the attribute names to a list
        names = []
        newsql = sql
        while True:
            start = newsql.find('%(')
            if start == -1:
                break
            end = newsql.find(')s', start)
            if end == -1: 
                break
            name = newsql[start+2 : end]
            names.append(name)
            newsql = newsql.replace(newsql[start:end+2], '?', 1)

        ps = self.__jdbcconn.prepareStatement(newsql)

        # Find parameter types
        types = []
        parmeta = ps.getParameterMetaData()
        for i in range(len(names)):
            types.append(parmeta.getParameterType(i+1))

        self.__prepstmts[sql] = (ps, names, types)

    def __executejdbcstmt(self, sql, args):
        if self.__resultset:
            self.__resultset.close()

        if sql not in self.__prepstmts:
            self.__preparejdbcstmt(sql)
        (ps, names, types) = self.__prepstmts[sql]

        for pos in range(len(names)): # Not very Pythonic, but we're doing Java
            if args[names[pos]] is None:
                ps.setNull(pos + 1, types[pos])
            else:
                ps.setObject(pos + 1, args[names[pos]], types[pos])

        if ps.execute():
            self.__resultset = ps.getResultSet()
            if sql not in self.__resultmeta:
                self.__resultmeta[sql] = \
                    self.__extractresultmetadata(self.__resultset)
            (self.__resultnames, self.__resulttypes) = self.__resultmeta[sql]
        else:
            self.__resultset = None
            (self.__resultnames, self.__resulttypes) = (None, None)

    def __extractresultmetadata(self, resultset):
        # Get jdbc resultset metadata. extract names and types
        # and add it to self.__resultmeta
        meta = resultset.getMetaData()
        names = []
        types = []
        for col in range(meta.getColumnCount()):
            names.append(meta.getColumnName(col+1))
            types.append(meta.getColumnType(col+1))
        return (names, types)

    def __readresultrow(self):
        if self.__resultset is None:
            return None
        result = []
        for i in range(len(self.__resulttypes)):
            e = self.__resulttypes[i] # Not Pythonic, but we need i for JDBC
            if e in (jdbc.Types.CHAR, jdbc.Types.VARCHAR, 
                     jdbc.Types.LONGVARCHAR):
                result.append(self.__resultset.getString(i+1))
            elif e in (jdbc.Types.BIT, jdbc.Types.BOOLEAN):
                result.append(self.__resultset.getBool(i+1))
            elif e in (jdbc.Types.TINYINT, jdbc.Types.SMALLINT, 
                       jdbc.Types.INTEGER):
                result.append(self.__resultset.getInt(i+1))
            elif e in (jdbc.Types.BIGINT, ):
                result.append(self.__resultset.getLong(i+1))
            elif e in (jdbc.Types.DATE, ):
                result.append(self.__resultset.getDate(i+1))
            elif e in (jdbc.Types.TIMESTAMP, ):
                result.append(self.__resultset.getTimestamp(i+1))
            elif e in (jdbc.Types.TIME, ):
                result.append(self.__resultset.getTime(i+1))
            else:
                # Try this and hope for the best...
                result.append(self.__resultset.getString(i+1))
        return tuple(result)

    def execute(self, stmt, arguments=None, namemapping=None, ignored=None):
        """Execute a statement.

           Arguments:
           - stmt: the statement to execute
           - arguments: a mapping with the arguments (default: None)
           - namemapping: a mapping of names such that if stmt uses %(arg)s
             and namemapping[arg]=arg2, the value arguments[arg2] is used 
             instead of arguments[arg]
           - ignored: An ignored argument only present to accept the same 
             number of arguments as ConnectionWrapper.execute
        """
        if namemapping and arguments:
            arguments = pygrametl.copy(arguments, **namemapping)
        else:
            arguments = pcopy(arguments)
        self.__queue.put((stmt, arguments))

    def executemany(self, stmt, params, ignored=None):
        """Execute a sequence of statements.

           Arguments:
           - stmt: the statement to execute
           - params: a sequence of arguments
           - ignored: An ignored argument only present to accept the same 
             number of arguments as ConnectionWrapper.executemany
        """
        for paramset in params:
            self.__queue.put((stmt, paramset))

    def rowfactory(self, names=None):
        """Return a generator object returning result rows (i.e. dicts)."""
        self.__queue.join()
        if names is None:
            if self.__resultnames is None:
                return
            else:
                names = [self.nametranslator(t[0]) for t in self.__resultnames]
        empty = (None, ) * len(self.__resultnames)
        while True:
            tuple = self.fetchonetuple()
            if tuple == empty:
                return
            yield dict(zip(names, tuple))

    def fetchone(self, names=None):
        """Return one result row (i.e. dict)."""
        self.__queue.join()
        if self.__resultset is None:
            return {}
        if names is None:
            names = [self.nametranslator(t[0]) for t in self.__resultnames]
        values = self.fetchonetuple()
        return dict(zip(names, values))

    def fetchonetuple(self):
        """Return one result tuple."""
        self.__queue.join()
        if self.__resultset is None:
            return ()
        if not self.__resultset.next():
            return (None, ) * len(self.__resultnames)
        else:
            return self.__readresultrow()

    def fetchmanytuples(self, cnt):
        """Return cnt result tuples."""
        self.__queue.join()
        if self.__resultset is None:
            return []
        empty = (None, ) * len(self.__resultnames)
        result = []
        for i in range(cnt):
            tmp = self.fetchonetuple()
            if tmp == empty:
                break
            result.append(tmp)
        return result

    def fetchalltuples(self):
        """Return all result tuples"""
        self.__queue.join()
        if self.__resultset is None:
            return []
        result = []
        empty = (None, ) * len(self.__resultnames)
        while True:
            tmp = self.fetchonetuple()
            if tmp == empty:
                return result
            result.append(tmp)

    def rowcount(self):
        """Not implemented. Return 0. Should return the size of the result."""
        return 0 

    def getunderlyingmodule(self):
        """Return a reference to the underlying connection's module."""
        return modules[self.__class__.__module__]

    def commit(self):
        """Commit the transaction."""
        pygrametl.endload()
        self.__queue.join()
        self.__jdbcconn.commit()

    def close(self):
        """Close the connection to the database,"""
        self.__queue.join()
        self.__jdbcconn.close()

    def rollback(self):
        """Rollback the transaction."""
        self.__queue.join()
        self.__jdbcconn.rollback()

    def setasdefault(self):
        """Set this ConnectionWrapper as the default connection."""
        pygrametl._defaulttargetconnection = self

    def cursor(self):
        """Not implemented for this JDBC connection wrapper!"""
        raise NotImplementedError, ".cursor() not supported"

    def resultnames(self):
        self.__queue.join()
        if self.__resultnames is None:
            return None
        else:
            return tuple(self.__resultnames)


def Date(year, month, day):
    date = '%s-%s-%s' % \
        (str(year).zfill(4), str(month).zfill(2), str(day).zfill(2)) 
    return jdbc.Date.valueOf(date)


def Timestamp(year, month, day, hour, minute, second):
    date = '%s-%s-%s %s:%s:%s' % \
        (str(year).zfill(4), str(month).zfill(2), str(day).zfill(2),
         str(hour).zfill(2), str(minute).zfill(2), str(second).zfill(2))
    return jdbc.Timestamp.valueOf(date)