import logging
from typing import Optional, Tuple, cast

import sqlglot.expressions as sge

from sqlglot.errors import OptimizeError
from sqlglot.schema import Schema as SQLGlotSchema, ensure_schema
from sqlglot.optimizer.annotate_types import annotate_types
from sqlglot.optimizer.qualify import qualify

import dlt
from dlt.common import logger
from dlt.common.libs.sqlglot import (
    to_sqlglot_type,
    from_sqlglot_type,
    set_metadata,
    get_metadata,
    TSqlGlotDialect,
)
from dlt.common.schema.typing import (
    TTableSchemaColumns,
    TColumnSchema,
)

from dlt.dataset.exceptions import LineageFailedException


def create_sqlglot_schema(
    schema: dlt.Schema,
    dataset_name: str,
    dialect: TSqlGlotDialect,
) -> SQLGlotSchema:
    """Create an SQLGlot schema using a dlt Schema and the destination capabilities.

    The SQLGlot schema automatically scopes the tables to the `dataset_name`.
    This can allow cross-dataset transformations on the same physical location.
    No name translation nor case folding is performing. All identifiers correspond
    to identifiers in dlt schema.
    """

    sqlglot_schema = {}  # MappingSchema(empty_schema, normalize=False)

    for table_name in schema.tables.keys():
        column_mapping = {}
        # skip not materialized columns
        for column_name, column in schema.get_table_columns(
            table_name, include_incomplete=False
        ).items():
            sqlglot_type = to_sqlglot_type(
                dlt_type=column["data_type"],
                nullable=column.get("nullable"),
                precision=column.get("precision"),
                scale=column.get("scale"),
                timezone=column.get("timezone"),
            )
            sqlglot_type = set_metadata(sqlglot_type, column)
            # column_name = sql_client.capabilities.casefold_identifier(column_name)
            column_mapping[column_name] = sqlglot_type
        # skip tables without columns
        if column_mapping:
            # table_name = sql_client.make_qualified_table_name_path(table_name, quote=False, casefold=False)[-1]
            sqlglot_schema[table_name] = column_mapping

    # ensure proper nesting with db and catalog
    nested_schema = {dataset_name: sqlglot_schema}

    return ensure_schema(nested_schema, dialect=dialect, normalize=False)


# NOTE even if `infer_sqlglot_schema=True`, some queries can have undetermined final columns
# TODO refactor such that it doesn't return a tuple
def compute_columns_schema(
    expression: sge.Expression,
    sqlglot_schema: SQLGlotSchema,
    dialect: TSqlGlotDialect,
    infer_sqlglot_schema: bool = True,
    allow_anonymous_columns: bool = True,
    allow_partial: bool = True,
) -> Tuple[TTableSchemaColumns, Optional[sge.Query]]:
    """Compute the expected dlt columns schema for the output of an SQL SELECT query. No case-folding or
    quoting is performed on the query.

    Args:
        infer_sqlglot_schema (bool): If False, all columns and tables referenced must be derived from the SQLGlot schema.
            If True, allow columns and tables not found in SQLGlot schema
        allow_anonymous_columns (bool): If False, all columns in final selection must have an explicit name or alias.
            If True, the name of columns from the final selection can be generated by the dialect
        allow_partial (bool): If False, raise exceptions if the schema returned is incomplete.
            If True, this function always returns a dictionary, even in cases of
            SQL parsing errors, missing table reference, unresolved `SELECT *`, etc.

    Returns: tuple of dlt columns schema and qualified `sql_query`

    """
    # make sure we don't modify the original expression
    expression = expression.copy()

    # the only thing we care is a list of selects
    if not isinstance(expression, sge.Query):
        if allow_partial:
            logger.debug(
                "Parsed SQL query is not a SQL query. Returning empty table schema because"
                " `allow_fail=True`"
            )
            return {}, None

        raise LineageFailedException(
            "Parsed SQL query is not a SELECT statement. Received SQL expression of type"
            f" {expression.type}.",
        )
    else:
        select_expression = expression

    # prevent normalization
    select_expression.meta["case_sensitive"] = True

    # false-y values `schema={}` or `schema=None` are identical to `infer_schema=True`
    try:
        select_expression = cast(
            sge.Query,
            qualify(
                select_expression,
                schema=sqlglot_schema,
                dialect=dialect,
                infer_schema=infer_sqlglot_schema,
                quote_identifiers=False,
                expand_stars=True,
            ),
        )
    except OptimizeError as e:
        raise LineageFailedException(
            f"Failed to resolve SQL query against the schema received: {e}"
        ) from e

    expression = annotate_types(expression, schema=sqlglot_schema)

    # NOTE: this has to be fixed
    if allow_anonymous_columns is False:
        for col in select_expression.selects:
            if col.alias_or_name == "" and False:
                raise LineageFailedException(
                    "Found anonymous column in SELECT statement. Use"
                    f" `allow_anonymous_columns=True` for permissive handling.\nColumn:\n\t{col}",
                )

    dlt_table_schema: dict[str, TColumnSchema] = {}
    for col in select_expression.selects:
        if col.output_name == "*":
            if allow_partial is True:
                logger.debug(
                    "SELECT statement includes a `*` selection that can't be resolved. "
                    "Returning empty table schema because `allow_fail=True`"
                )
                continue

            raise LineageFailedException(
                "SELECT statement includes a `*` selection that can't be resolved. Modify the"
                " query to select columns explicitly or limit `*` to known tables (e.g., `SELECT"
                " known_table.*`). Use `allow_fail=True` to return a partial dlt schema with"
                f" resolvable column hints.\nColumn:\n\t{col}",
            )

        data_type_hints = from_sqlglot_type(sqlglot_type=col.type)
        additional_hints = get_metadata(sqlglot_type=col.type)
        # get original name for queries that aliased the source dlt column
        propagated_name = additional_hints.pop("name", None)
        # NOTE dictionary unpacking order matters; unpacking `data_type_hints` last ensures precedence.
        dlt_table_schema[col.output_name] = {
            "name": col.output_name,
            **additional_hints,
            **data_type_hints,
        }
        if propagated_name and col.output_name != propagated_name:
            dlt_table_schema[col.output_name]["x-original-name"] = propagated_name  # type: ignore[typeddict-unknown-key]

    return dlt_table_schema, expression
