import operator as op
from typing import Any, Callable, Optional, Sequence

import polars as pl
import pyparsing as pp

__all__ = ['query']

from numpy.random.mtrand import Sequence

ops = {'**': op.pow,
       '~': op.not_,
       '*': op.mul, '@': op.matmul, '/': op.truediv, '//': op.floordiv, '%': op.mod,
       '+': op.add, '-': op.sub,
       '<<': op.lshift, '>>': op.rshift,
       '&': op.and_, '^': op.xor, '|': op.or_,
       'in': lambda a, b: a in b, 'not in': lambda a, b: a not in b,
       'is': lambda a, b: a is b, 'is not': lambda a, b: a is not b,
       '<': op.lt, '<=': op.le, '>': op.gt, '>=': op.ge, '!=': op.ne, '==': op.eq,
       'and': lambda a, b: a and b, 'or': lambda a, b: a or b}


def parse_op(_a, _b, s: Sequence[Any]) -> list[Callable[[Any, Any], Any]]:
    return [ops[s[0]]]


def parse_chain(_a, _b, s: Sequence[Any]) -> list[Any]:
    return [s[1](s[0], s[2]) & s[3](s[2], s[4])]


def parse_col(_a, _b, s: Sequence[Any]) -> list[pl.col]:
    col = pl.col(s[0][len(s[0]) // 2])
    for t in s[1:]:
        if len(t) == 2:  # property
            col = getattr(col, t[1])
        elif len(t) == 5:  # function
            col = getattr(col, t[1])(t[3])
    return [col]


def parse_expr(_a, _b, s: Sequence[Any]) -> list[Any]:
    s = list(s)
    for operator in ops.values():
        while operator in s[1::2]:
            idx = 1 + 2 * s[1::2].index(operator)
            s = s[:idx - 1] + [operator(s[idx - 1], s[idx + 1])] + s[idx + 2:]
    return s


def parse_brackets(_a, _b, s: Sequence[Any]) -> list[Any]:
    return [s[1]]


def parse_string(_a, _b, s: Sequence[Any]) -> list[str]:
    return [s[0][1:-1]]


def parse_list(_a, _b, s: Sequence[Any]) -> list[Sequence[Any]]:
    return [s[1:-1]]


def parse_not(_a, _b, s: Sequence[Any]) -> list[Any]:
    return [not s[1]] if len(s) == 2 else [s[0]]


def parse_const(_a, _b, s: Sequence[Any]) -> list[Optional[bool]]:
    return [{'None': None, 'True': True, 'False': False}[s[0]]]


Num = pp.common.number
String = pp.dbl_quoted_string ^ pp.sgl_quoted_string
String.set_parse_action(parse_string)
Var = pp.Combine('@' + pp.Word(pp.alphas + "_", pp.alphanums + "_"))
Const = pp.Literal('None') ^ 'True' ^ 'False'
Const.set_parse_action(parse_const)
Not = pp.Opt(pp.Literal('not'))
Id = Num ^ String ^ Var ^ Const
Id0 = Not + Id
Id0.set_parse_action(parse_not)
List = ('(' + pp.DelimitedList(Id0, allow_trailing_delim=True) + ')' ^
        '[' + pp.DelimitedList(Id0, allow_trailing_delim=True) + ']')
List.set_parse_action(parse_list)
Col = (pp.Group(pp.Word(pp.alphanums + "_") ^
        pp.Literal('`') + pp.Word(pp.printables + ' ', exclude_chars='`') + pp.Literal('`'))
       + pp.ZeroOrMore(pp.Group(pp.Char('.') + pp.Opt(pp.Word(pp.alphas + '_', pp.alphanums + "_")) +
                       pp.Opt(pp.Literal('(') + pp.Opt(Id0) + pp.Literal(')')))))
Col.set_parse_action(parse_col)
Id = Id ^ List | Col
Id1 = Not + Id
Id1.set_parse_action(parse_not)
CompOp = pp.Combine(pp.Char('<>') + pp.Opt('=')) ^ '==' ^ '!='
CompOp.set_parse_action(parse_op)
CalcOp = pp.Char('+-/*%@&|^~') ^ '//' ^ '**' ^ '<<' ^ '>>' ^ 'in' ^ 'not in' ^ 'is' ^ 'is not'
CalcOp.set_parse_action(parse_op)
Op = CompOp ^ CalcOp
Expr = pp.Forward()
Brackets = pp.Literal('(') + Expr + pp.Literal(')')
Brackets.set_parse_action(parse_brackets)
Chain = Id1 + CompOp + Id1 + CompOp + Id1
Chain.set_parse_action(parse_chain)
Id = Brackets | Chain | Id
Id2 = Not + Id
Id2.set_parse_action(parse_not)
Expr << Id2 + pp.ZeroOrMore(Op + Id2)
Expr.set_parse_action(parse_expr)


def query(q: str, local_dict: dict[str, Any] = None) -> Any:
    """ use a query string in polars.filter as can be done in pandas.query """
    local_dict = local_dict or {}

    def parse_var(_a, _b, s: Sequence[Any]) -> list[Any]:
        return [local_dict[s[0][1:]]]

    Var.set_parse_action(parse_var)
    return Expr.parse_string(q, parse_all=True)[0]
