import itertools

from sqlglot import alias, exp
from sqlglot.errors import OptimizeError
from sqlglot.optimizer.schema import ensure_schema
from sqlglot.optimizer.scope import traverse_scope


def qualify_columns(expression, schema):
    """
    Rewrite sqlglot AST to have fully qualified columns.

    Example:
        >>> import sqlglot
        >>> schema = {"tbl": {"col": "INT"}}
        >>> expression = sqlglot.parse_one("SELECT col FROM tbl")
        >>> qualify_columns(expression, schema).sql()
        'SELECT tbl.col AS col FROM tbl'

    Args:
        expression (sqlglot.Expression): expression to qualify
        schema (dict|sqlglot.optimizer.Schema): Database schema
    Returns:
        sqlglot.Expression: qualified expression
    """
    schema = ensure_schema(schema)

    for scope in traverse_scope(expression):
        _check_union_outputs(scope)
        _pop_table_column_aliases(scope.ctes)
        _pop_table_column_aliases(scope.derived_tables)
        _qualify_columns(scope, schema)
        _expand_stars(scope, schema)
        _qualify_outputs(scope)
        _check_unknown_tables(scope)

    return expression


def _check_union_outputs(scope):
    """Assert that the outputs of both sides of a UNION are the same"""
    if not isinstance(scope.expression, exp.Union):
        return
    left, right = scope.union
    if left.expression.named_selects != right.expression.named_selects:
        raise OptimizeError(
            f"UNION outputs not equal: {left.expression.named_selects} vs. {left.expression.named_selects}"
        )


def _pop_table_column_aliases(derived_tables):
    """
    Remove table column aliases.

    (e.g. SELECT ... FROM (SELECT ...) AS foo(col1, col2)
    """
    for derived_table in derived_tables:
        table_alias = derived_table.args.get("alias")
        if table_alias:
            table_alias.args.pop("columns", None)


def _qualify_columns(scope, schema):
    """Disambiguate columns, ensuring each column specifies a source"""
    unambiguous_columns = None  # lazily loaded

    for column in scope.columns:
        column_table = column.table
        column_name = column.name

        if (
            column_table
            and column_table in scope.sources
            and column_name
            not in _get_source_columns(column_table, scope.sources, schema)
        ):
            raise OptimizeError(f"Unknown column: {column_name}")

        if not column_table:
            if unambiguous_columns is None:
                source_columns = {
                    k: _get_source_columns(k, scope.sources, schema)
                    for k in scope.selected_sources
                }

                unambiguous_columns = _get_unambiguous_columns(source_columns)

            column_table = unambiguous_columns.get(column_name)
            if not column_table and not scope.is_subquery:
                raise OptimizeError(f"Ambiguous column: {column_name}")
            column.set("table", exp.to_identifier(column_table))


def _expand_stars(scope, schema):
    """Expand stars to lists of column selections"""

    new_selections = []

    for expression in scope.selects:
        if isinstance(expression, exp.Star):
            tables = list(scope.selected_sources)
        elif isinstance(expression, exp.Column) and isinstance(
            expression.this, exp.Star
        ):
            tables = [expression.table]
        else:
            new_selections.append(expression)
            continue

        for table in tables:
            if table not in scope.sources:
                raise OptimizeError(f"Unknown table: {table}")
            columns = _get_source_columns(table, scope.sources, schema)
            for column in columns:
                new_selections.append(exp.column(column, table))

    scope.expression.set("expressions", new_selections)


def _qualify_outputs(scope):
    """Ensure all output columns are aliased"""

    new_selections = []

    for i, (selection, aliased_column) in enumerate(
        itertools.zip_longest(scope.selects, scope.outer_column_list)
    ):
        if isinstance(selection, exp.Column):
            # convoluted setter because a simple selection.replace(alias) would require a copy
            alias_ = alias(exp.column(""), alias=selection.name)
            alias_.set("this", selection)
            selection = alias_
        elif not isinstance(selection, exp.Alias):
            alias_ = alias(exp.column(""), f"_col_{i}")
            alias_.set("this", selection)
            selection = alias_

        if aliased_column:
            selection.set("alias", exp.to_identifier(aliased_column))

        new_selections.append(selection)

    scope.expression.set("expressions", new_selections)


def _check_unknown_tables(scope):
    if scope.external_columns and not scope.is_correlated_subquery:
        raise OptimizeError(f"Unknown table: {scope.external_columns[0].text('table')}")


def _get_unambiguous_columns(source_columns):
    """
    Find all the unambiguous columns in sources.

    Args:
        source_columns (dict): Mapping of names to source columns
    Returns:
        dict: Mapping of column name to source name
    """
    if not source_columns:
        return {}

    source_columns = list(source_columns.items())

    first_table, first_columns = source_columns[0]
    unambiguous_columns = {
        col: first_table for col in _find_unique_columns(first_columns)
    }

    for table, columns in source_columns[1:]:
        unique = _find_unique_columns(columns)
        ambiguous = set(unambiguous_columns).intersection(unique)
        for column in ambiguous:
            unambiguous_columns.pop(column)
        for column in unique.difference(ambiguous):
            unambiguous_columns[column] = table

    return unambiguous_columns


def _find_unique_columns(columns):
    """
    Find the unique columns in a list of columns.

    Example:
        >>> sorted(_find_unique_columns(["a", "b", "b", "c"]))
        ['a', 'c']

    This is necessary because duplicate column names are ambiguous.
    """
    counts = {}
    for column in columns:
        counts[column] = counts.get(column, 0) + 1
    return {column for column, count in counts.items() if count == 1}


def _get_source_columns(name, sources, schema):
    """Resolve the source columns for a given source `name`"""
    if name not in sources:
        raise OptimizeError(f"Unknown table: {name}")

    source = sources[name]

    # If referencing a table, return the columns from the schema
    if isinstance(source, exp.Table):
        try:
            return schema.column_names(source)
        except Exception as e:
            raise OptimizeError(str(e)) from e

    # Otherwise, if referencing another scope, return that scope's named selects
    return source.expression.named_selects
