"""
Collection of query wrappers / abstractions to both facilitate data
retrieval and to reduce dependency on DB-specific API.
"""
from datetime import datetime

import numpy as np
import traceback

from pandas.core.datetools import format as date_format
from pandas.core.api import DataFrame, isnull

#------------------------------------------------------------------------------
# Helper execution function


def execute(sql, con, retry=True, cur=None, params=None):
    """
    Execute the given SQL query using the provided connection object.

    Parameters
    ----------
    sql: string
        Query to be executed

    Returns
    -------
    Cursor object
    """
    try:
        if cur is None:
            cur = con.cursor()

        if params is None:
            cur.execute(sql)
        else:
            cur.execute(sql, params)
        return cur
    except Exception:
        try:
            con.rollback()
        except Exception:  # pragma: no cover
            pass

        print 'Error on sql %s' % sql
        raise


def _safe_fetch(cur):
    try:
        result = cur.fetchall()
        if not isinstance(result, list):
            result = list(result)
        return result
    except Exception, e:  # pragma: no cover
        excName = e.__class__.__name__
        if excName == 'OperationalError':
            return []


def tquery(sql, con=None, cur=None, retry=True):
    """
    Returns list of tuples corresponding to each row in given sql
    query.

    If only one column selected, then plain list is returned.

    Parameters
    ----------
    sql: string
        SQL query to be executed
    con: SQLConnection or DB API 2.0-compliant connection
    cur: DB API 2.0 cursor

    Provide a specific connection or a specific cursor if you are executing a
    lot of sequential statements and want to commit outside.
    """
    cur = execute(sql, con, cur=cur)
    result = _safe_fetch(cur)

    if con is not None:
        try:
            cur.close()
            con.commit()
        except Exception, e:
            excName = e.__class__.__name__
            if excName == 'OperationalError':  # pragma: no cover
                print 'Failed to commit, may need to restart interpreter'
            else:
                raise

            traceback.print_exc()
            if retry:
                return tquery(sql, con=con, retry=False)

    if result and len(result[0]) == 1:
        # python 3 compat
        result = list(list(zip(*result))[0])
    elif result is None:  # pragma: no cover
        result = []

    return result


def uquery(sql, con=None, cur=None, retry=True, params=()):
    """
    Does the same thing as tquery, but instead of returning results, it
    returns the number of rows affected.  Good for update queries.
    """
    cur = execute(sql, con, cur=cur, retry=retry, params=params)

    result = cur.rowcount
    try:
        con.commit()
    except Exception, e:
        excName = e.__class__.__name__
        if excName != 'OperationalError':
            raise

        traceback.print_exc()
        if retry:
            print 'Looks like your connection failed, reconnecting...'
            return uquery(sql, con, retry=False)
    return result


def read_frame(sql, con, index_col=None, coerce_float=True):
    """
    Returns a DataFrame corresponding to the result set of the query
    string.

    Optionally provide an index_col parameter to use one of the
    columns as the index. Otherwise will be 0 to len(results) - 1.

    Parameters
    ----------
    sql: string
        SQL query to be executed
    con: DB connection object, optional
    index_col: string, optional
        column name to use for the returned DataFrame object.
    """
    cur = execute(sql, con)
    rows = _safe_fetch(cur)
    columns = [col_desc[0] for col_desc in cur.description]

    cur.close()
    con.commit()

    result = DataFrame.from_records(rows, columns=columns,
                                    coerce_float=coerce_float)

    if index_col is not None:
        result = result.set_index(index_col)

    return result

frame_query = read_frame


def write_frame(frame, name=None, con=None, flavor='sqlite', append=False):
    """
    Write records stored in a DataFrame to SQLite. The index will currently be
    dropped
    """
    if flavor == 'sqlite':
        schema = get_sqlite_schema(frame, name)
    else:
        raise NotImplementedError

    if not append and not has_table(name, con):
        con.execute(schema)

    wildcards = ','.join(['?'] * len(frame.columns))
    insert_sql = 'INSERT INTO %s VALUES (%s)' % (name, wildcards)
    data = [tuple(x) for x in frame.values]
    con.executemany(insert_sql, data)


def has_table(name, con):
    sqlstr = "SELECT name FROM sqlite_master WHERE type='table' AND name='%s'" % name
    rs = tquery(sqlstr, con)
    return len(rs) > 0


def get_sqlite_schema(frame, name, dtypes=None, keys=None):
    template = """
CREATE TABLE %(name)s (
  %(columns)s%(keystr)s
);"""

    column_types = []

    frame_types = frame.dtypes
    for k in frame_types.index:
        dt = frame_types[k]

        if issubclass(dt.type, (np.integer, np.bool_)):
            sqltype = 'INTEGER'
        elif issubclass(dt.type, np.floating):
            sqltype = 'REAL'
        else:
            sqltype = 'TEXT'

        if dtypes is not None:
            sqltype = dtypes.get(k, sqltype)

        column_types.append((k, sqltype))
    columns = ',\n  '.join('[%s] %s' % x for x in column_types)

    keystr = ''
    if keys is not None:
        if isinstance(keys, basestring):
            keys = (keys,)
        keystr = ', PRIMARY KEY (%s)' % ','.join(keys)
    return template % {'name': name, 'columns': columns, 'keystr': keystr}


#------------------------------------------------------------------------------
# Query formatting

_formatters = {
    datetime: lambda dt: "'%s'" % date_format(dt),
    str: lambda x: "'%s'" % x,
    np.str_: lambda x: "'%s'" % x,
    unicode: lambda x: "'%s'" % x,
    float: lambda x: "%.8f" % x,
    int: lambda x: "%s" % x,
    type(None): lambda x: "NULL",
    np.float64: lambda x: "%.10f" % x,
    bool: lambda x: "'%s'" % x,
}


def format_query(sql, *args):
    """

    """
    processed_args = []
    for arg in args:
        if isinstance(arg, float) and isnull(arg):
            arg = None

        formatter = _formatters[type(arg)]
        processed_args.append(formatter(arg))

    return sql % tuple(processed_args)
