from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
from __future__ import unicode_literals

from sqlalchemy.engine import default
from sqlalchemy.sql import compiler
from sqlalchemy import types

import pydruid.db
from pydruid.db import exceptions


RESERVED_SCHEMAS = ['INFORMATION_SCHEMA']


type_map = {
    'char': types.String,
    'varchar': types.String,
    'float': types.Float,
    'decimal': types.Float,
    'real': types.Float,
    'double': types.Float,
    'boolean': types.Boolean,
    'tinyint': types.BigInteger,
    'smallint': types.BigInteger,
    'integer': types.BigInteger,
    'bigint': types.BigInteger,
    'timestamp': types.TIMESTAMP,
    'date': types.DATE,
    'other': types.BLOB,
}


class UniversalSet(object):
    def __contains__(self, item):
        return True


class DruidIdentifierPreparer(compiler.IdentifierPreparer):
    reserved_words = UniversalSet()


class DruidCompiler(compiler.SQLCompiler):
    pass


class DruidTypeCompiler(compiler.GenericTypeCompiler):
    def visit_REAL(self, type_, **kwargs):
        return "DOUBLE"

    def visit_NUMERIC(self, type_, **kwargs):
        return "LONG"

    visit_DECIMAL = visit_NUMERIC
    visit_INTEGER = visit_NUMERIC
    visit_SMALLINT = visit_NUMERIC
    visit_BIGINT = visit_NUMERIC
    visit_BOOLEAN = visit_NUMERIC
    visit_TIMESTAMP = visit_NUMERIC
    visit_DATE = visit_NUMERIC

    def visit_CHAR(self, type_, **kwargs):
        return "STRING"

    visit_NCHAR = visit_CHAR
    visit_VARCHAR = visit_CHAR
    visit_NVARCHAR = visit_CHAR
    visit_TEXT = visit_CHAR

    def visit_DATETIME(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type DATETIME is not supported')

    def visit_TIME(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type TIME is not supported')

    def visit_BINARY(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type BINARY is not supported')

    def visit_VARBINARY(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type VARBINARY is not supported')

    def visit_BLOB(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type BLOB is not supported')

    def visit_CLOB(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type CBLOB is not supported')

    def visit_NCLOB(self, type_, **kwargs):
        raise exceptions.NotSupportedError('Type NCBLOB is not supported')


class DruidDialect(default.DefaultDialect):

    name = 'druid'
    scheme = 'http'
    driver = 'rest'
    preparer = DruidIdentifierPreparer
    statement_compiler = DruidCompiler
    type_compiler = DruidTypeCompiler
    supports_alter = False
    supports_pk_autoincrement = False
    supports_default_values = False
    supports_empty_insert = False
    supports_unicode_statements = True
    supports_unicode_binds = True
    returns_unicode_strings = True
    description_encoding = None
    supports_native_boolean = True

    @classmethod
    def dbapi(cls):
        return pydruid.db

    def create_connect_args(self, url):
        kwargs = {
            'host': url.host,
            'port': url.port or 8082,
            'path': url.database,
            'scheme': self.scheme,
        }
        return ([], kwargs)

    def get_schema_names(self, connection, **kwargs):
        # Each Druid datasource appears as a table in the "druid" schema. This
        # is also the default schema, so Druid datasources can be referenced as
        # either druid.dataSourceName or simply dataSourceName.
        result = connection.execute(
            'SELECT SCHEMA_NAME FROM INFORMATION_SCHEMA.SCHEMATA')

        return [
            row.SCHEMA_NAME for row in result
            if row.SCHEMA_NAME not in RESERVED_SCHEMAS
        ]

    def has_table(self, connection, table_name, schema=None):
        query = """
            SELECT COUNT(*) > 0 AS exists_
              FROM INFORMATION_SCHEMA.TABLES
             WHERE TABLE_NAME = '{table_name}'
        """.format(table_name=table_name)

        result = connection.execute(query)
        return result.fetchone().exists_

    def get_table_names(self, connection, schema=None, **kwargs):
        query = "SELECT TABLE_NAME FROM INFORMATION_SCHEMA.TABLES"
        if schema:
            query = "{query} WHERE TABLE_SCHEMA = '{schema}'".format(
                query=query, schema=schema)

        result = connection.execute(query)
        return [row.TABLE_NAME for row in result]

    def get_view_names(self, connection, schema=None, **kwargs):
        return []

    def get_table_options(self, connection, table_name, schema=None, **kwargs):
        return {}

    def get_columns(self, connection, table_name, schema=None, **kwargs):
        query = """
            SELECT COLUMN_NAME,
                   DATA_TYPE,
                   IS_NULLABLE,
                   COLUMN_DEFAULT
              FROM INFORMATION_SCHEMA.COLUMNS
             WHERE TABLE_NAME = '{table_name}'
        """.format(table_name=table_name)
        if schema:
            query = "{query} AND TABLE_SCHEMA = '{schema}'".format(
                query=query, schema=schema)

        result = connection.execute(query)

        return [
            {
                'name': row.COLUMN_NAME,
                'type': type_map[row.DATA_TYPE.lower()],
                'nullable': get_is_nullable(row.IS_NULLABLE),
                'default': get_default(row.COLUMN_DEFAULT),
            }
            for row in result
        ]

    def get_pk_constraint(self, connection, table_name, schema=None, **kwargs):
        return {'constrained_columns': [], 'name': None}

    def get_foreign_keys(self, connection, table_name, schema=None, **kwargs):
        return []

    def get_check_constraints(
        self,
        connection,
        table_name,
        schema=None,
        **kwargs
    ):
        return []

    def get_table_comment(self, connection, table_name, schema=None, **kwargs):
        return {'text': ''}

    def get_indexes(self, connection, table_name, schema=None, **kwargs):
        return []

    def get_unique_constraints(
        self,
        connection,
        table_name,
        schema=None,
        **kwargs
    ):
        return []

    def get_view_definition(
        self,
        connection,
        view_name,
        schema=None,
        **kwargs
    ):
        pass

    def do_rollback(self, dbapi_connection):
        pass

    def _check_unicode_returns(self, connection, additional_tests=None):
        return True

    def _check_unicode_description(self, connection):
        return True


DruidHTTPDialect = DruidDialect


class DruidHTTPSDialect(DruidDialect):

    scheme = 'https'


def get_is_nullable(druid_is_nullable):
    # this should be 'YES' or 'NO'; we default to no
    return druid_is_nullable.lower() == 'yes'


def get_default(druid_column_default):
    # currently unused, returns ''
    return str(druid_column_default) if druid_column_default != '' else None
