import traceback
import logging
import time

from sqlalchemy import event
from sqlalchemy.engine import Engine
from sqlalchemy.sql.compiler import Compiled
from sqlalchemy.exc import UnsupportedCompilationError
from sqlalchemy.dialects import postgresql

from ylog.context import log_context

from watcher.config import settings
from .state import context_state

logger = logging.getLogger(__name__)


class SqlRequestExecutedProvider:

    def __init__(self, connection, prepared_statement):
        self.connection = connection
        self.prepared_statement = prepared_statement

    def query_transaction_isolation_level(self):
        return self.connection.get_isolation_level()

    def query_transaction_status(self):
        return 1 if self.connection.in_transaction() else 0

    def query_prepared_sql(self):
        try:
            text = getattr(self.prepared_statement, 'text', None)
            if not text:
                if hasattr(self.prepared_statement, 'compile'):
                    text = str(
                        self.prepared_statement.compile(
                            dialect=postgresql.dialect()
                        )
                    )
                else:
                    text = str(self.prepared_statement)

            return text.strip()
        except UnsupportedCompilationError:
            return 'Unable to compile sql'

    def duration(self):
        query_start_time = self.connection.info.get('query_start_time')
        if query_start_time:
            return max(
                time.time() - query_start_time.pop(-1),
                0
            ) * 1000
        return 0

    def make_context(self):
        duration = self.duration()
        return {
            'is_select': self.query_prepared_sql().lower().startswith('select'),
            'is_slow': duration > settings.LOG_CONTEXT_SLOW_SQL,
            'vendor': 'database',
            'iso_level': self.query_transaction_isolation_level(),
            'trans_status': self.query_transaction_status(),
            'query_to_analyse': self.query_prepared_sql(),
            'duration': duration,
        }


class SqlAlchemyLogger:

    def structured_logs(self, conn, clauseelement, multiparams, params, result):
        """
        Записать в логи SQL Alchemy запрос.
        """
        if isinstance(clauseelement, Compiled):
            statement = clauseelement.statement
        else:
            statement = clauseelement

        state = context_state.get()
        context_provider = SqlRequestExecutedProvider(
            connection=conn,
            prepared_statement=statement,
        )
        context = context_provider.make_context()

        state.add_sql_time(context['duration'])

        if not settings.LOG_CONTEXT_ENABLE_DB_TRACKING:
            return

        if settings.LOG_CONTEXT_ENABLE_STACKTRACES:
            context['stacktrace'] = ''.join(
                i for i in traceback.format_stack()[:-1]
            )

        params = {
            'profiling': context,
        }

        with log_context(execution_time=int(context['duration']), **params):
            logger.info(
                '(%.1f msec) %s',
                context['duration'],
                context_provider.query_prepared_sql()
            )


sql_logger = SqlAlchemyLogger()


def set_start_time(conn, *args, **kwargs):
    conn.info.setdefault('query_start_time', []).append(time.time())


def enable_sql_instrumentation():
    event.listen(Engine, "after_execute", sql_logger.structured_logs)
    event.listen(Engine, "before_execute", set_start_time)
