from __future__ import annotations

import operator
from datetime import datetime, timedelta
from typing import TYPE_CHECKING, Any, Callable, Iterable, Sequence, Union, overload

if TYPE_CHECKING:
    from pandas import DataFrame, Series

# Type for operands in expressions. These are dtypes compatible with pandas DataFrame
# columns. (see https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.dtypes.html)
Operand = Union["_Key", "Operator", str, float, int, datetime, timedelta]
Numeric = float | int


class _SupportsOperators:
    def __lt__(self, other: Operand) -> "Operator": ...
    def __le__(self, other: Operand) -> "Operator": ...
    def __eq__(self, other: Operand) -> "Operator": ...
    def __ne__(self, other: Operand) -> "Operator": ...
    def __gt__(self, other: Operand) -> "Operator": ...
    def __ge__(self, other: Operand) -> "Operator": ...
    def __add__(self, other: Operand) -> "Operator": ...
    def __sub__(self, other: Operand) -> "Operator": ...
    def __mul__(self, other: Operand) -> "Operator": ...
    def __truediv__(self, other: Operand) -> "Operator": ...
    def __floordiv__(self, other: Operand) -> "Operator": ...
    def __mod__(self, other: Operand) -> "Operator": ...
    def __pow__(self, other: Operand) -> "Operator": ...
    def __or__(self, other: Operand) -> "Operator": ...
    def __and__(self, other: Operand) -> "Operator": ...


def add_operators(cls):
    ops = {
        "__lt__": operator.lt,
        "__le__": operator.le,
        "__eq__": operator.eq,
        "__ne__": operator.ne,
        "__gt__": operator.gt,
        "__ge__": operator.ge,
        "__add__": operator.add,
        "__sub__": operator.sub,
        "__mul__": operator.mul,
        "__truediv__": operator.truediv,
        "__floordiv__": operator.floordiv,
        "__mod__": operator.mod,
        "__pow__": operator.pow,
        "__or__": operator.or_,
        "__and__": operator.and_,
    }

    def make_op(op_func):
        def method(self, other):
            return Operator(op_func, self, other)

        return method

    for name, func in ops.items():
        setattr(cls, name, make_op(func))

    return cls


@add_operators
class Operator(_SupportsOperators):
    @overload
    def __init__(
        self,
        op: Callable[[Any, Any], bool],
        a: Numeric | str | _Key | Operator,
        b: Numeric | str | _Key | Operator,
    ) -> None: ...
    @overload
    def __init__(
        self, op: Callable[[Any], bool], a: Numeric | str | _Key | Operator
    ) -> None: ...
    def __init__(self, op, a, b=None):
        self._op = op
        self._a = a
        self._b = b

    def evaluate(self, item: Any) -> Any:
        def resolve(val):
            if isinstance(val, _Key):
                return item[val._value]
            elif isinstance(val, Operator):
                return val.evaluate(item)
            return val

        if self._b is None:
            return self._op(resolve(self._a))
        return self._op(resolve(self._a), resolve(self._b))

    def __repr__(self) -> str:
        return f"Operation({self._a} {self._op.__name__} {self._b})"


@add_operators
class _Key(_SupportsOperators):
    def __init__(self, key: str) -> None:
        self._value = key

    def __repr__(self) -> str:
        return f"Key({self._value})"


class Filter:
    """Filter applied to a collection."""

    def __init__(self, *operators: Operator, tags: Iterable[str] | None = None) -> None:
        self._ops: Sequence[Operator] = operators
        self._tags: set[str] = set(tags) if tags else set()

    def __and__(self, other: Filter | None) -> Filter:
        return (
            Filter(*self._ops, *other._ops, tags=self._tags.union(other._tags))
            if other
            else self
        )

    def apply(self, df: DataFrame) -> DataFrame | Series:
        import pandas as pd

        if df.empty:
            return df

        # filter by operators
        mask = pd.Series(True, index=df.index)
        for op in self._ops:
            mask &= op.evaluate(df)

        # filter by tags if specified
        if self._tags:
            tag_mask = df["tags"].map(lambda x: self._tags.issubset(x))
            mask &= tag_mask

        return df[mask]

    def __repr__(self) -> str:
        return "Filter({})".format(" & ".join(str(op) for op in self._ops))


class SortInstruction:
    def __init__(self, key: str, ascending: bool = True) -> None:
        self.key = key
        self.ascending = ascending

    def __repr__(self) -> str:
        order = "ASC" if self.ascending else "DESC"
        return f"SortInstruction({self.key} {order})"


class Sorter:
    """Sorter applied to a collection."""

    def __init__(self, *instructions: SortInstruction) -> None:
        self._instructions: Sequence[SortInstruction] = instructions

    def __and__(self, other: Sorter | None) -> Sorter:
        return Sorter(*self._instructions, *other._instructions) if other else self

    def apply(self, df: DataFrame) -> DataFrame:
        if not self._instructions:
            return df
        by = [instr.key for instr in self._instructions]
        ascending = [instr.ascending for instr in self._instructions]
        return df.sort_values(by=by, ascending=ascending)

    def __repr__(self) -> str:
        return "Sorter({})".format(
            ", ".join(str(instr) for instr in self._instructions)
        )
