r"""
SQL Connection object (:mod:`qiita_db.sql_connection`)
======================================================

.. currentmodule:: qiita_db.sql_connection

This modules provides wrappers for the psycopg2 module to allow easy use of
transaction blocks and SQL execution/data retrieval.

This module provides the variable TRN, which is the transaction available
to use in the system. The singleton pattern is applied and this works as long
as the system remains single-threaded.

Classes
-------

.. autosummary::
   :toctree: generated/

   Transaction
"""
# -----------------------------------------------------------------------------
# Copyright (c) 2014--, The Qiita Development Team.
#
# Distributed under the terms of the BSD 3-clause License.
#
# The full license is in the file LICENSE, distributed with this software.
# -----------------------------------------------------------------------------
from contextlib import contextmanager
from itertools import chain
from functools import wraps

from psycopg2 import (connect, ProgrammingError, Error as PostgresError,
                      OperationalError, errorcodes)
from psycopg2.extras import DictCursor
from psycopg2.extensions import TRANSACTION_STATUS_IDLE

from qiita_core.qiita_settings import qiita_config


def _checker(func):
    """Decorator to check that methods are executed inside the context"""
    @wraps(func)
    def wrapper(self, *args, **kwargs):
        if self._contexts_entered == 0:
            raise RuntimeError(
                "Operation not permitted. Transaction methods can only be "
                "invoked within the context manager.")
        return func(self, *args, **kwargs)
    return wrapper


class Transaction(object):
    """A context manager that encapsulates a DB transaction

    A transaction is defined by a series of consecutive queries that need to
    be applied to the database as a single block.

    Raises
    ------
    RuntimeError
        If the transaction methods are invoked outside a context.

    Notes
    -----
    When the execution leaves the context manager, any remaining queries in
    the transaction will be executed and committed.
    """
    def __init__(self, admin=False):
        self._queries = []
        self._results = []
        self._contexts_entered = 0
        self._connection = None
        self._post_commit_funcs = []
        self._post_rollback_funcs = []
        self.admin = admin

    def _open_connection(self):
        # If the connection already exists and is not closed, don't do anything
        if self._connection is not None and self._connection.closed == 0:
            return

        try:
            if self.admin:
                self._connection = connect(
                    user=qiita_config.admin_user,
                    password=qiita_config.admin_password,
                    host=qiita_config.host,
                    port=qiita_config.port)
                self._connection.autocommit = True
            else:
                self._connection = connect(user=qiita_config.user,
                                           password=qiita_config.password,
                                           database=qiita_config.database,
                                           host=qiita_config.host,
                                           port=qiita_config.port)
        except OperationalError as e:
            # catch three known common exceptions and raise runtime errors
            try:
                etype = str(e).split(':')[1].split()[0]
            except IndexError:
                # we recieved a really unanticipated error without a colon
                etype = ''
            if etype == 'database':
                etext = ('This is likely because the database `%s` has not '
                         'been created or has been dropped.' %
                         qiita_config.database)
            elif etype == 'role':
                etext = ('This is likely because the user string `%s` '
                         'supplied in your configuration file `%s` is '
                         'incorrect or not an authorized postgres user.' %
                         (qiita_config.user, qiita_config.conf_fp))
            elif etype == 'Connection':
                etext = ('This is likely because postgres isn\'t '
                         'running. Check that postgres is correctly '
                         'installed and is running.')
            else:
                # we recieved a really unanticipated error with a colon
                etext = ''
            ebase = ('An OperationalError with the following message occured'
                     '\n\n\t%s\n%s For more information, review `INSTALL.md`'
                     ' in the Qiita installation base directory.')
            raise RuntimeError(ebase % (str(e), etext))

    def close(self):
        if self._connection is not None:
            self._connection.close()

    @contextmanager
    def _get_cursor(self):
        """Returns a postgres cursor

        Returns
        -------
        psycopg2.cursor
            The psycopg2 cursor

        Raises
        ------
        RuntimeError
            if the cursor cannot be created
        """
        self._open_connection()

        try:
            with self._connection.cursor(cursor_factory=DictCursor) as cur:
                yield cur
        except PostgresError as e:
            raise RuntimeError("Cannot get postgres cursor: %s" % e)

    def __enter__(self):
        self._open_connection()
        self._contexts_entered += 1
        return self

    def _clean_up(self, exc_type):
        if exc_type is not None:
            # An exception occurred during the execution of the transaction
            # Make sure that we leave the DB w/o any modification
            self.rollback()
        elif self._queries:
            # There are still queries to be executed, execute them
            # It is safe to use the execute method here, as internally is
            # wrapped in a try/except and rollbacks in case of failure
            self.execute()
            self.commit()
        elif self._connection.get_transaction_status() != \
                TRANSACTION_STATUS_IDLE:
            # There are no queries to be executed, however, the transaction
            # is still not committed. Commit it so the changes are not lost
            self.commit()

    def __exit__(self, exc_type, exc_value, traceback):
        # We only need to perform some action if this is the last context
        # that we are entering
        if self._contexts_entered == 1:
            # We need to wrap the entire function in a try/finally because
            # at the end we need to decrement _contexts_entered
            try:
                self._clean_up(exc_type)
            finally:
                self._contexts_entered -= 1
        else:
            self._contexts_entered -= 1

    def _raise_execution_error(self, sql, sql_args, error):
        """Rollbacks the current transaction and raises a useful error
        The error message contains the name of the transaction, the failed
        query, the arguments of the failed query and the error generated.

        Raises
        ------
        ValueError
        """
        self.rollback()

        try:
            ec_lu = errorcodes.lookup(error.pgcode)
            raise ValueError(
                "Error running SQL: %s. MSG: %s\n" % (ec_lu, str(error)))
        except (KeyError, AttributeError):
            raise ValueError("Error running SQL query: %s" % str(error))

    @_checker
    def add(self, sql, sql_args=None, many=False):
        """Add a sql query to the transaction

        Parameters
        ----------
        sql : str
            The sql query
        sql_args : list, tuple or dict of objects, optional
            The arguments to the sql query
        many : bool, optional
            Whether or not we should add the query multiple times to the
            transaction

        Raises
        ------
        TypeError
            If `sql_args` is provided and is not a list, tuple or dict
        RuntimeError
            If invoked outside a context

        Notes
        -----
        If `many` is true, `sql_args` should be a list of lists, tuples or
        dicts, in which each element of the list contains the parameters for
        one SQL query of the many. Each element on the list is all the
        parameters for a single one of the many queries added. The amount of
        SQL queries added to the list is len(sql_args).
        """
        if not many:
            sql_args = [sql_args]

        for args in sql_args:
            if args:
                if not isinstance(args, (list, tuple, dict)):
                    raise TypeError("sql_args should be a list, tuple or dict."
                                    " Found %s" % type(args))
            self._queries.append((sql, args))

    def _execute(self):
        """Internal function that actually executes the transaction
        The `execute` function exposed in the API wraps this one to make sure
        that we catch any exception that happens in here and we rollback the
        transaction
        """
        with self._get_cursor() as cur:
            for sql, sql_args in self._queries: