import typing as t

from snsql._ast.ast import Query
from snsql.metadata import Metadata
from snsql.sql.parse import QueryParser

from sarus_sql.dialects import SQLDialect
import sarus_sql.translator as tr


def parse_query(
    query_str: str,
    metadata: t.Optional[t.Union[t.Dict[str, t.Any], Metadata]] = None,
    dialect: t.Optional[SQLDialect] = None,
) -> Query:
    """Parses the string query. if dialect is provided, we suppose that the
    the query is belonging to that dialect and the outcome will be
    the translation from the dilalect to postgres.
    """
    try:
        if dialect:
            parsed_query = tr.translate_to_postgres(
                QueryParser().query(query_str), dialect
            )
        else:
            parsed_query = QueryParser().query(query_str)
    except BaseException as err:
        raise type(err)(
            f"Query parsing failed for `{query_str}`\n with error:\n{err}"
        )
    if metadata is not None:
        meta = (
            Metadata.from_dict(metadata)
            if isinstance(metadata, t.Dict)
            else metadata
        )
        load_symbols_in_query(
            parsed_query,
            meta,
        )
    return parsed_query


def load_symbols_in_query(
    query: Query,
    meta: Metadata,
    row_privacy: t.Optional[bool] = None,
    ctes: t.Optional[t.Dict[str, "Query"]] = None,
) -> Query:
    """Load the metadata in the query and set the row_privacy
    (Override the load_symbols)"""
    try:
        if ctes is not None:
            query.load_symbols(meta, ctes=ctes)
        else:
            query.load_symbols(meta)
    except BaseException as err:
        raise type(err)(
            "Please check the syntax of your query.\n"
            "Loading metadata into the parsed query failed with error:\n"
            f"{err}\nquery:{query}"
            f"\nGot metadata: {meta}"
        )

    if row_privacy:
        query.row_privacy = row_privacy
    return query