from sqlalchemy.ext.compiler import compiles

from .. import util
from .impl import DefaultImpl
from .base import alter_table, AddColumn, ColumnName, \
    format_table_name, format_column_name, ColumnNullable, alter_column,\
    format_server_default,ColumnDefault, format_type, ColumnType

class MSSQLImpl(DefaultImpl):
    __dialect__ = 'mssql'
    transactional_ddl = True
    batch_separator = "GO"

    def __init__(self, *arg, **kw):
        super(MSSQLImpl, self).__init__(*arg, **kw)
        self.batch_separator = self.context_opts.get(
                                "mssql_batch_separator",
                                self.batch_separator)

    def _exec(self, construct, *args, **kw):
        super(MSSQLImpl, self)._exec(construct, *args, **kw)
        if self.as_sql and self.batch_separator:
            self.static_output(self.batch_separator)

    def emit_begin(self):
        self.static_output("BEGIN TRANSACTION" + self.command_terminator)

    def emit_commit(self):
        super(MSSQLImpl, self).emit_commit()
        if self.as_sql and self.batch_separator:
            self.static_output(self.batch_separator)

    def alter_column(self, table_name, column_name,
                        nullable=None,
                        server_default=False,
                        name=None,
                        type_=None,
                        schema=None,
                        autoincrement=None,
                        existing_type=None,
                        existing_server_default=None,
                        existing_nullable=None,
                        existing_autoincrement=None
                    ):

        if nullable is not None and existing_type is None:
            if type_ is not None:
                existing_type = type_
                # the NULL/NOT NULL alter will handle
                # the type alteration
                type_ = None
            else:
                raise util.CommandError(
                        "MS-SQL ALTER COLUMN operations "
                        "with NULL or NOT NULL require the "
                        "existing_type or a new type_ be passed.")

        super(MSSQLImpl, self).alter_column(
                        table_name, column_name,
                        nullable=nullable,
                        type_=type_,
                        schema=schema,
                        autoincrement=autoincrement,
                        existing_type=existing_type,
                        existing_nullable=existing_nullable,
                        existing_autoincrement=existing_autoincrement
        )

        if server_default is not False:
            if existing_server_default is not False or \
                server_default is None:
                self._exec(
                    _exec_drop_col_constraint(self,
                            table_name, column_name,
                            'sys.default_constraints')
                )
            if server_default is not None:
                super(MSSQLImpl, self).alter_column(
                                table_name, column_name,
                                schema=schema,
                                server_default=server_default)

        if name is not None:
            super(MSSQLImpl, self).alter_column(
                                table_name, column_name,
                                schema=schema,
                                name=name)

    def bulk_insert(self, table, rows):
        if self.as_sql:
            self._exec(
                "SET IDENTITY_INSERT %s ON" %
                    self.dialect.identifier_preparer.format_table(table)
            )
            super(MSSQLImpl, self).bulk_insert(table, rows)
            self._exec(
                "SET IDENTITY_INSERT %s OFF" %
                    self.dialect.identifier_preparer.format_table(table)
            )
        else:
            super(MSSQLImpl, self).bulk_insert(table, rows)


    def drop_column(self, table_name, column, **kw):
        drop_default = kw.pop('mssql_drop_default', False)
        if drop_default:
            self._exec(
                _exec_drop_col_constraint(self,
                        table_name, column,
                        'sys.default_constraints')
            )
        drop_check = kw.pop('mssql_drop_check', False)
        if drop_check:
            self._exec(
                _exec_drop_col_constraint(self,
                        table_name, column,
                        'sys.check_constraints')
            )
        drop_fks = kw.pop('mssql_drop_foreign_key', False)
        if drop_fks:
            self._exec(
                _exec_drop_col_fk_constraint(self,
                        table_name, column)
            )
        super(MSSQLImpl, self).drop_column(table_name, column)

def _exec_drop_col_constraint(impl, tname, colname, type_):
    # from http://www.mssqltips.com/sqlservertip/1425/working-with-default-constraints-in-sql-server/
    # TODO: needs table formatting, etc.
    return """declare @const_name varchar(256)
select @const_name = [name] from %(type)s
where parent_object_id = object_id('%(tname)s')
and col_name(parent_object_id, parent_column_id) = '%(colname)s'
exec('alter table %(tname)s drop constraint ' + @const_name)""" % {
        'type': type_,
        'tname': tname,
        'colname': colname
    }

def _exec_drop_col_fk_constraint(impl, tname, colname):
    return """declare @const_name varchar(256)
select @const_name = [name] from
    sys.foreign_keys fk join sys.foreign_key_columns fkc
    on fk.object_id=fkc.constraint_object_id
where fkc.parent_object_id = object_id('%(tname)s')
and col_name(fkc.parent_object_id, fkc.parent_column_id) = '%(colname)s'
exec('alter table %(tname)s drop constraint ' + @const_name)""" % {
        'tname': tname,
        'colname': colname
    }



@compiles(AddColumn, 'mssql')
def visit_add_column(element, compiler, **kw):
    return "%s %s" % (
        alter_table(compiler, element.table_name, element.schema),
        mssql_add_column(compiler, element.column, **kw)
    )

def mssql_add_column(compiler, column, **kw):
    return "ADD %s" % compiler.get_column_specification(column, **kw)

@compiles(ColumnNullable, 'mssql')
def visit_column_nullable(element, compiler, **kw):
    return "%s %s %s %s" % (
        alter_table(compiler, element.table_name, element.schema),
        alter_column(compiler, element.column_name),
        format_type(compiler, element.existing_type),
        "NULL" if element.nullable else "NOT NULL"
    )

@compiles(ColumnDefault, 'mssql')
def visit_column_default(element, compiler, **kw):
    # TODO: there can also be a named constraint
    # with ADD CONSTRAINT here
    return "%s ADD DEFAULT %s FOR %s" % (
        alter_table(compiler, element.table_name, element.schema),
        format_server_default(compiler, element.default),
        format_column_name(compiler, element.column_name)
    )

@compiles(ColumnName, 'mssql')
def visit_rename_column(element, compiler, **kw):
    return "EXEC sp_rename '%s.%s', %s, 'COLUMN'" % (
        format_table_name(compiler, element.table_name, element.schema),
        format_column_name(compiler, element.column_name),
        format_column_name(compiler, element.newname)
    )

@compiles(ColumnType, 'mssql')
def visit_column_type(element, compiler, **kw):
    return "%s %s %s" % (
        alter_table(compiler, element.table_name, element.schema),
        alter_column(compiler, element.column_name),
        format_type(compiler, element.type_)
    )

