"""
Functions for building queries, from nodes or SQL.
"""

import ast
import datetime
import operator
import re
from typing import Any, Callable, Dict, List, Literal, Optional, Set, Tuple, cast

from dateutil.parser import parse
from sqlalchemy.engine.url import make_url
from sqlalchemy.schema import Column as SqlaColumn
from sqlalchemy.sql.elements import BinaryExpression
from sqlalchemy.sql.expression import ClauseElement
from sqlmodel import Session, select
from sqloxide import parse_sql

from datajunction.constants import DEFAULT_DIMENSION_COLUMN
from datajunction.errors import DJError, DJInvalidInputException, ErrorCode
from datajunction.models.node import Node, NodeType
from datajunction.models.query import QueryCreate
from datajunction.sql.dag import (
    get_database_for_nodes,
    get_dimensions,
    get_referenced_columns_from_sql,
    get_referenced_columns_from_tree,
)
from datajunction.sql.parse import (
    find_nodes_by_key,
    find_nodes_by_key_with_parent,
    get_expression_from_projection,
)
from datajunction.sql.transpile import get_query, get_select_for_node
from datajunction.typing import (
    Expression,
    Identifier,
    Join,
    Projection,
    Relation,
    Select,
)
from datajunction.utils import get_session

FILTER_RE = re.compile(r"([\w\./_]+)(<=|<|>=|>|!=|=)(.+)")
FilterOperator = Literal[">", ">=", "<", "<=", "=", "!="]
COMPARISONS: Dict[FilterOperator, Callable[[Any, Any], bool]] = {
    ">": operator.gt,
    ">=": operator.ge,
    "<": operator.lt,
    "<=": operator.le,
    "=": operator.eq,
    "!=": operator.ne,
}


def parse_filter(filter_: str) -> Tuple[str, FilterOperator, str]:
    """
    Parse a filter into name, op, value.
    """
    match = FILTER_RE.match(filter_)
    if not match:
        raise DJInvalidInputException(
            message=f'The filter "{filter_}" is invalid',
            errors=[
                DJError(
                    code=ErrorCode.INVALID_FILTER_PATTERN,
                    message=(
                        f'The filter "{filter_}" is not a valid filter. Filters should '
                        "consist of a dimension name, follow by a valid operator "
                        "(<=|<|>=|>|!=|=), followed by a value. If the value is a "
                        "string or date/time it should be enclosed in single quotes."
                    ),
                    debug={"context": {"filter": filter_}},
                ),
            ],
        )

    name, operator_, value = match.groups()
    operator_ = cast(FilterOperator, operator_)
    return name, operator_, value


def get_filter(columns: Dict[str, SqlaColumn], filter_: str) -> BinaryExpression:
    """
    Build a SQLAlchemy filter.
    """
    name, operator_, value = parse_filter(filter_)

    if name not in columns:
        raise Exception(f"Invalid column name: {name}")
    column = columns[name]
    comparison = COMPARISONS[operator_]

    if column.type.python_type in [datetime.date, datetime.datetime]:
        try:
            value = str(parse(value))
        except Exception as ex:
            raise Exception(f"Invalid date or datetime value: {value}") from ex
    else:
        try:
            value = ast.literal_eval(value)
        except Exception as ex:
            raise Exception(f"Invalid value: {value}") from ex

    return comparison(column, value)


def get_dimensions_from_filters(filters: List[str]) -> Set[str]:
    """
    Extract dimensions from filters passed to the metric API.
    """
    return {parse_filter(filter_)[0] for filter_ in filters}


def get_query_for_node(  # pylint: disable=too-many-locals
    session: Session,
    node: Node,
    groupbys: List[str],
    filters: List[str],
    database_id: Optional[int] = None,
) -> QueryCreate:
    """
    Return a DJ QueryCreate object from a given node.
    """
    # check that groupbys and filters are valid dimensions
    requested_dimensions = set(groupbys) | get_dimensions_from_filters(filters)
    valid_dimensions = set(get_dimensions(node))
    if not requested_dimensions <= valid_dimensions:
        invalid = sorted(requested_dimensions - valid_dimensions)
        plural = "s" if len(invalid) > 1 else ""
        raise Exception(f"Invalid dimension{plural}: {', '.join(invalid)}")

    # which columns are needed from the parents; this is used to determine the database
    # where the query will run
    referenced_columns = get_referenced_columns_from_sql(node.expression, node.parents)

    # extract all referenced dimensions so we can join the node with them
    dimensions: Dict[str, Node] = {}
    for dimension in requested_dimensions:
        name, column = dimension.rsplit(".", 1)
        if (
            name not in {parent.name for parent in node.parents}
            and name not in dimensions
        ):
            dimensions[name] = session.exec(select(Node).where(Node.name == name)).one()
            referenced_columns[name].add(column)

    # find database
    nodes = [node]
    nodes.extend(dimensions.values())
    database = get_database_for_nodes(session, nodes, referenced_columns, database_id)

    # base query
    node_select = get_select_for_node(node, database)
    source = node_select.froms[0]

    # join with dimensions
    for dimension in dimensions.values():
        subquery = get_select_for_node(
            dimension,
            database,
            referenced_columns[dimension.name],
        ).alias(dimension.name)
        condition = find_on_clause(node, source, dimension, subquery)
        node_select = node_select.select_from(source.join(subquery, condition))

    columns = {
        f"{column.table.name}.{column.name}": column
        for from_ in node_select.froms
        for column in from_.columns
    }

    # filter
    node_select = node_select.filter(
        *[get_filter(columns, filter_) for filter_ in filters]
    )

    # groupby
    node_select = node_select.group_by(*[columns[groupby] for groupby in groupbys])

    # add groupbys to projection as well
    for groupby in groupbys:
        node_select.append_column(columns[groupby])

    dialect = make_url(database.URI).get_dialect()
    sql = str(
        node_select.compile(dialect=dialect(), compile_kwargs={"literal_binds": True}),
    )

    return QueryCreate(database_id=database.id, submitted_query=sql)


def find_on_clause(
    node: Node,
    node_select: Select,
    dimension: Node,
    subquery: Select,
) -> ClauseElement:
    """
    Return the on clause for a node/dimension selects.
    """
    for parent in node.parents:
        for column in parent.columns:
            if column.dimension == dimension:
                dimension_column = column.dimension_column or DEFAULT_DIMENSION_COLUMN
                return (
                    node_select.columns[column.name]  # type: ignore
                    == subquery.columns[dimension_column]  # type: ignore
                )

    raise Exception(f"Node {node.name} has no columns with dimension {dimension.name}")


# pylint: disable=too-many-branches, too-many-locals, too-many-statements
def get_query_for_sql(sql: str) -> QueryCreate:
    """
    Return a query given a SQL expression querying the repo.

    Eg:

        SELECT
            "core.users.gender", "core.num_comments"
        FROM metrics
        WHERE "core.comments.user_id" > 1
        GROUP BY
            "core.users.gender"

    This works by converting metrics (``core.num_comments``) into their selection
    definition (``COUNT(*)``), updating the sources to include the metrics parents
    (including joining with dimensions), and updating column references in the
    ``WHERE``, ``GROUP BY``, etc.
    """
    session = next(get_session())

    tree = parse_sql(sql, dialect="ansi")
    query_select = tree[0]["Query"]["body"]["Select"]

    # fetch all metric and dimension nodes
    nodes = {node.name: node for node in session.exec(