from sqlalchemy import schema as sa_schema, types as sqltypes
from sqlalchemy.engine.reflection import Inspector
from sqlalchemy import event
from ..operations import ops
import logging
from .. import util
from ..util import compat
from ..util import sqla_compat
from sqlalchemy.util import OrderedSet
import re
from .render import _user_defined_render
import contextlib
from alembic.ddl.base import _fk_spec

log = logging.getLogger(__name__)


def _populate_migration_script(autogen_context, migration_script):
    upgrade_ops = migration_script.upgrade_ops_list[-1]
    downgrade_ops = migration_script.downgrade_ops_list[-1]

    _produce_net_changes(autogen_context, upgrade_ops)
    upgrade_ops.reverse_into(downgrade_ops)


comparators = util.Dispatcher(uselist=True)


def _produce_net_changes(autogen_context, upgrade_ops):

    connection = autogen_context.connection
    include_schemas = autogen_context.opts.get('include_schemas', False)

    inspector = Inspector.from_engine(connection)

    default_schema = connection.dialect.default_schema_name
    if include_schemas:
        schemas = set(inspector.get_schema_names())
        # replace default schema name with None
        schemas.discard("information_schema")
        # replace the "default" schema with None
        schemas.discard(default_schema)
        schemas.add(None)
    else:
        schemas = [None]

    comparators.dispatch("schema", autogen_context.dialect.name)(
        autogen_context, upgrade_ops, schemas
    )


@comparators.dispatch_for("schema")
def _autogen_for_tables(autogen_context, upgrade_ops, schemas):
    inspector = autogen_context.inspector

    metadata = autogen_context.metadata

    conn_table_names = set()

    version_table_schema = \
        autogen_context.migration_context.version_table_schema
    version_table = autogen_context.migration_context.version_table

    for s in schemas:
        tables = set(inspector.get_table_names(schema=s))
        if s == version_table_schema:
            tables = tables.difference(
                [autogen_context.migration_context.version_table]
            )
        conn_table_names.update(zip([s] * len(tables), tables))

    metadata_table_names = OrderedSet(
        [(table.schema, table.name) for table in metadata.sorted_tables]
    ).difference([(version_table_schema, version_table)])

    _compare_tables(conn_table_names, metadata_table_names,
                    inspector, metadata, upgrade_ops, autogen_context)


def _compare_tables(conn_table_names, metadata_table_names,
                    inspector, metadata, upgrade_ops, autogen_context):

    default_schema = inspector.bind.dialect.default_schema_name

    # tables coming from the connection will not have "schema"
    # set if it matches default_schema_name; so we need a list
    # of table names from local metadata that also have "None" if schema
    # == default_schema_name.  Most setups will be like this anyway but
    # some are not (see #170)
    metadata_table_names_no_dflt_schema = OrderedSet([
        (schema if schema != default_schema else None, tname)
        for schema, tname in metadata_table_names
    ])

    # to adjust for the MetaData collection storing the tables either
    # as "schemaname.tablename" or just "tablename", create a new lookup
    # which will match the "non-default-schema" keys to the Table object.
    tname_to_table = dict(
        (
            no_dflt_schema,
            metadata.tables[sa_schema._get_table_key(tname, schema)]
        )
        for no_dflt_schema, (schema, tname) in zip(
            metadata_table_names_no_dflt_schema,
            metadata_table_names)
    )
    metadata_table_names = metadata_table_names_no_dflt_schema

    for s, tname in metadata_table_names.difference(conn_table_names):
        name = '%s.%s' % (s, tname) if s else tname
        metadata_table = tname_to_table[(s, tname)]
        if autogen_context.run_filters(
                metadata_table, tname, "table", False, None):
            upgrade_ops.ops.append(
                ops.CreateTableOp.from_table(metadata_table))
            log.info("Detected added table %r", name)
            modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)

            comparators.dispatch("table")(
                autogen_context, modify_table_ops,
                s, tname, None, metadata_table
            )
            if not modify_table_ops.is_empty():
                upgrade_ops.ops.append(modify_table_ops)

    removal_metadata = sa_schema.MetaData()
    for s, tname in conn_table_names.difference(metadata_table_names):
        name = sa_schema._get_table_key(tname, s)
        exists = name in removal_metadata.tables
        t = sa_schema.Table(tname, removal_metadata, schema=s)

        if not exists:
            event.listen(
                t,
                "column_reflect",
                autogen_context.migration_context.impl.
                _compat_autogen_column_reflect(inspector))
            inspector.reflecttable(t, None)
        if autogen_context.run_filters(t, tname, "table", True, None):
            upgrade_ops.ops.append(
                ops.DropTableOp.from_table(t)
            )
            log.info("Detected removed table %r", name)

    existing_tables = conn_table_names.intersection(metadata_table_names)

    existing_metadata = sa_schema.MetaData()
    conn_column_info = {}
    for s, tname in existing_tables:
        name = sa_schema._get_table_key(tname, s)
        exists = name in existing_metadata.tables
        t = sa_schema.Table(tname, existing_metadata, schema=s)
        if not exists:
            event.listen(
                t,
                "column_reflect",
                autogen_context.migration_context.impl.
                _compat_autogen_column_reflect(inspector))
            inspector.reflecttable(t, None)
        conn_column_info[(s, tname)] = t

    for s, tname in sorted(existing_tables, key=lambda x: (x[0] or '', x[1])):
        s = s or None
        name = '%s.%s' % (s, tname) if s else tname
        metadata_table = tname_to_table[(s, tname)]
        conn_table = existing_metadata.tables[name]

        if autogen_context.run_filters(
                metadata_table, tname, "table", False,
                conn_table):

            modify_table_ops = ops.ModifyTableOps(tname, [], schema=s)
            with _compare_columns(
                s, tname,
                conn_table,
                metadata_table,
                    modify_table_ops, autogen_context, inspector):

                comparators.dispatch("table")(
                    autogen_context, modify_table_ops,
                    s, tname, conn_table, metadata_table
                )

            if not modify_table_ops.is_empty():
                upgrade_ops.ops.append(modify_table_ops)


def _make_index(params, conn_table):
    # TODO: add .info such as 'duplicates_constraint'
    return sa_schema.Index(
        params['name'],
        *[conn_table.c[cname] for cname in params['column_names']],
        unique=params['unique']
    )


def _make_unique_constraint(params, conn_table):
    uq = sa_schema.UniqueConstraint(
        *[conn_table.c[cname] for cname in params['column_names']],
        name=params['name']
    )
    if 'duplicates_index' in params:
        uq.info['duplicates_index'] = params['duplicates_index']

    return uq


def _make_foreign_key(params, conn_table):
    tname = params['referred_table']
    if params['referred_schema']:
        tname = "%s.%s" % (params['referred_schema'], tname)

    options = params.get('options', {})

    const = sa_schema.ForeignKeyConstraint(
        [conn_table.c[cname] for cname in params['constrained_columns']],
        ["%s.%s" % (tname, n) for n in params['referred_columns']],
        onupdate=options.get('onupdate'),
        ondelete=options.get('ondelete'),
        deferrable=options.get('deferrable'),
        initially=options.get('initially'),
        name=params['name']
    )
    # needed by 0.7
    conn_table.append_constraint(const)
    return const


@contextlib.contextmanager
def _compare_columns(schema, tname, conn_table, metadata_table,
                     modify_table_ops, autogen_context, inspector):
    name = '%s.%s' % (schema, tname) if schema else tname
    metadata_cols_by_name = dict((c.name, c) for c in metadata_table.c)
    conn_col_names = dict((c.name, c) for c in conn_table.c)
    metadata_col_names = OrderedSet(sorted(metadata_cols_by_name))

    for cname in metadata_col_names.difference(conn_col_names):
        if autogen_context.run_filters(
                metadata_cols_by_name[cname], cname,
                "column", False, None):
            modify_table_ops.ops.append(
                ops.AddColumnOp.from_column_and_tablename(
                    schema, tname, metadata_cols_by_name[cname])
            )
            log.info("Detected added column '%s.%s'", name, cname)

    for colname in metadata_col_names.intersection(conn_col_names):
        metadata_col = metadata_cols_by_name[colname]
        conn_col = conn_table.c[colname]
        if not autogen_context.run_filters(
                metadata_col, colname, "column", False,
                conn_col):
            continue
        alter_column_op = ops.AlterColumnOp(
            tname, colname, schema=schema)

        comparators.dispatch("column")(
            autogen_context, alter_column_op,
            schema, tname, colname, conn_col, metadata_col
        )

        if alter_column_op.has_changes():
            modify_table_ops.ops.append(alter_column_op)

    yield

    for cname in set(conn_col_names).difference(metadata_col_names):
        if autogen_context.run_filters(
                conn_table.c[cname], cname,
                "column", True, None):
            modify_table_ops.ops.append(
                ops.DropColumnOp.from_column_and_tablename(
                    schema, tname, conn_table.c[cname]
                )
            )
            log.info("Detected removed column '%s.%s'", name, cname)


class _constraint_sig(object):

    def __eq__(self, other):
        return self.const == other.const

    def __ne__(self, other):
        return self.const != other.const

    def __hash__(self):
        return hash(self.const)


class _uq_constraint_sig(_constraint_sig):
    is_index = False
    is_unique = True

    def __init__(self, const):
        self.const = const
        self.name = const.name
        self.sig = tuple(sorted([col.name for col in const.columns]))

    @property
    def column_names(self):
        return [col.name for col in self.const.columns]


class _ix_constraint_sig(_constraint_sig):
    is_index = True

    def __init__(self, const):
        self.const = const
        self.name = const.name
        self.sig = tuple(sorted([col.name for col in const.columns]))
        self.is_unique = bool(const.unique)

    @property
    def column_names(self):
        return sqla_compat._get_index_column_names(self.const)


class _fk_constraint_sig(_constraint_sig):
    def __init__(self, const, include_options=False):
        self.const = const