from collections.abc import Mapping
from typing import Any, Optional, TypeAlias, TypeVar

import sqlalchemy as sa
from lark import Lark, LarkError, Transformer, Tree
from lark.lexer import Token

from ai.backend.manager.errors.resource import DataTransformationFailed

from . import (
    ArrayFieldItem,
    EnumFieldItem,
    FieldSpecItem,
    JSONFieldItem,
    ORMFieldItem,
    get_col_from_table,
)

__all__ = (
    "FieldSpecType",
    "FilterableSQLQuery",
    "QueryFilterParser",
)

_grammar = r"""
    ?start: expr
    value: string
           | number
           | array
           | ATOM -> atom
    ATOM       : "null" | "true" | "false"
    COMBINE_OP : "&" | "|"
    UNARY_OP   : "!"
    BINARY_OP  : "==" | "!="
               | ">" | ">="
               | "<" | "<="
               | "contains" | "in"
               | "isnot" | "is"
               | "like" | "ilike"
    expr: UNARY_OP expr         -> unary_expr
        | CNAME BINARY_OP value -> binary_expr
        | expr COMBINE_OP expr  -> combine_expr
        | "(" expr ")"          -> paren_expr
    array  : "[" [value ("," value)*] "]"
    string : ESCAPED_STRING
    number : SIGNED_NUMBER
    %import common.CNAME
    %import common.ESCAPED_STRING
    %import common.SIGNED_NUMBER
    %import common.WS
    %ignore WS
"""
_parser = Lark(
    _grammar,
    parser="lalr",
    maybe_placeholders=False,
)

FilterableSQLQuery: TypeAlias = sa.sql.Select | sa.sql.Update | sa.sql.Delete
_TQuery = TypeVar("_TQuery", sa.sql.Select, sa.sql.Update, sa.sql.Delete)
FieldSpecType: TypeAlias = Mapping[str, FieldSpecItem] | None
WhereClauseType: TypeAlias = sa.sql.expression.ColumnElement[bool]


class QueryFilterTransformer(Transformer):
    def __init__(self, sa_table: sa.Table, fieldspec: Optional[FieldSpecType] = None) -> None:
        super().__init__()
        self._sa_table = sa_table
        self._fieldspec = fieldspec

    def string(self, token: list[Token]) -> str:
        s = token[0]
        # SQL-side escaping is handled by SQLAlchemy
        return s[1:-1].replace('\\"', '"')

    def number(self, token: list[Token]) -> int | float:
        n = token[0]
        if "." in n:
            return float(n)
        return int(n)

    array = list

    def atom(self, token: list[Token]) -> sa.sql.elements.ColumnElement[Any]:
        a = token[0]
        if a.value == "null":
            return sa.null()
        if a.value == "true":
            return sa.true()
        if a.value == "false":
            return sa.false()
        raise ValueError("Unknown/unsupported atomic token", a.value)

    def _transform_val_leaf(self, col_name: str, op: str, value: Any) -> Any:
        if self._fieldspec:
            try:
                func = self._fieldspec[col_name][1]
            except KeyError:
                raise ValueError("Unknown/unsupported field name", col_name)
            return func(value) if func is not None else value
        return value

    def _transform_val(self, col_name: str, op: str, value: Any) -> Any:
        if isinstance(value, Tree):
            val = self._transform_val(col_name, op, value.children[0])
        elif isinstance(value, list):
            val = [self._transform_val(col_name, op, v) for v in value]
        else:
            val = self._transform_val_leaf(col_name, op, value)
        return val

    def binary_expr(self, *args: Any) -> sa.sql.elements.ColumnElement[Any]:
        children: list[Token] = args[0]
        col_name = children[0].value
        op = children[1].value
        val = self._transform_val(col_name, op, children[2])

        def build_expr(op: str, col: Any, val: Any) -> sa.sql.elements.ColumnElement[Any]:
            match op:
                case "==":
                    expr = col == val
                case "!=":
                    expr = col != val
                case ">":
                    expr = col > val
                case ">=":
                    expr = col >= val
                case "<":
                    expr = col < val
                case "<=":
                    expr = col <= val
                case "contains":
                    expr = col.contains(val)
                case "in":
                    expr = col.in_(val)
                case "isnot":
                    expr = col.isnot(val)
                case "is":
                    expr = col.is_(val)
                case "like":
                    expr = col.like(val)
                case "ilike":
                    expr = col.ilike(val)
                case _:
                    raise ValueError(f"Unknown operator: {op}")
            return expr

        expr: sa.sql.elements.ColumnElement[Any]
        try:
            if self._fieldspec is not None:
                match self._fieldspec[children[0].value][0]:
                    case ArrayFieldItem(col_name):
                        # For array columns, let's apply the expression on every item,
                        # and select the row if anyone makes the result true.
                        col = get_col_from_table(self._sa_table, col_name)
                        unnested_col = sa.func.unnest(col).alias("item")
                        subq: sa.sql.Select[Any] = (
                            sa.select(sa.column("item"))
                            .select_from(unnested_col)
                            .where(build_expr(op, sa.column("item"), val))
                        )
                        expr = sa.exists(subq)
                    case JSONFieldItem(col_name, obj_key):
                        # For json columns, we additionally indicate the object key
                        # to retrieve the value used in the expression.
                        col = get_col_from_table(self._sa_table, col_name).op("->>")(obj_key)
                        expr = build_expr(op, col, val)
                    case EnumFieldItem(col_name, enum_cls):
                        col = get_col_from_table(self._sa_table, col_name)
                        # allow both key and value of enum to be specified on variable `val`
                        # fetch original enum pointer from given `val`
                        try:
                            enum_val = enum_cls(val)
                        except ValueError:
                            try:
                                enum_val = enum_cls[val]
                            except KeyError:
                                raise ValueError(f"Invalid enum value: {val}")
                        expr = build_expr(op, col, enum_val)
                    case ORMFieldItem(column):
                        expr = build_expr(op, column, val)
                    case str(col_name):
                        col = get_col_from_table(self._sa_table, col_name)
                        expr = build_expr(op, col, val)
            else:
                col = get_col_from_table(self._sa_table, col_name)
                expr = build_expr(op, col, val)
        except KeyError:
            raise ValueError("Unknown/unsupported field name", col_name)
        return expr

    def unary_expr(self, *args):
        children = args[0]
        op = children[0].value
        expr = children[1]
        if op in ("not", "!"):
            return sa.not_(expr)
        return args

    def combine_expr(self, *args):
        children = args[0]
        op = children[1].value
        expr1 = children[0]
        expr2 = children[2]
        if op == "&":
            return sa.and_(expr1, expr2)
        if op == "|":
            return sa.or_(expr1, expr2)
        return args

    def paren_expr(self, *args):
        children = args[0]
        return children[0]


class QueryFilterParser:
    def __init__(self, fieldspec: Optional[FieldSpecType] = None) -> None:
        self._fieldspec = fieldspec
        self._parser = _parser

    def parse_filter(
        self,
        table,
        filter_expr: str,
    ) -> WhereClauseType:
        try:
            ast = self._parser.parse(filter_expr)
            where_clause = QueryFilterTransformer(table, self._fieldspec).transform(ast)
        except LarkError as e:
            raise ValueError(f"Query filter parsing error: {e}")
        return where_clause

    def append_filter(
        self,
        sa_query: _TQuery,
        filter_expr: str,
    ) -> _TQuery:
        """
        Parse the given filter expression and build the where clause based on the first target table from
        the given SQLAlchemy query object.
        """
        if isinstance(sa_query, sa.sql.Select):
            table = sa_query.froms[0]
        elif isinstance(sa_query, (sa.sql.Delete, sa.sql.Update)):
            table = sa_query.table
        else:
            raise ValueError("Unsupported SQLAlchemy query object type")
        where_clause = self.parse_filter(table, filter_expr)
        final_query = sa_query.where(where_clause)
        if final_query is None:
            raise DataTransformationFailed("Failed to apply filter to query")
        return final_query
