import builtins
import typing
from enum import Enum

import topk_sdk.data

FlexibleExpr = typing.Union[str, int, float, bool, None, LogicalExpr]
Numeric = typing.Union[int, float, LogicalExpr]
Boolish = typing.Union[bool, LogicalExpr]
Stringy = typing.Union[str, LogicalExpr]

class LogicalExpr(Enum):
    def __repr__(self) -> builtins.str: ...
    def _expr_eq(self, other: LogicalExpr) -> builtins.bool: ...
    # Unary operators
    def is_null(self) -> LogicalExpr: ...
    def is_not_null(self) -> LogicalExpr: ...
    # Comparison operators
    def eq(self, other: FlexibleExpr) -> LogicalExpr: ...
    def __eq__(self, other: FlexibleExpr) -> LogicalExpr: ...
    def ne(self, other: FlexibleExpr) -> LogicalExpr: ...
    def __ne__(self, other: FlexibleExpr) -> LogicalExpr: ...
    def lt(self, other: Numeric) -> LogicalExpr: ...
    def __lt__(self, other: Numeric) -> LogicalExpr: ...
    def __rlt__(self, other: Numeric) -> LogicalExpr: ...
    def lte(self, other: Numeric) -> LogicalExpr: ...
    def __le__(self, other: Numeric) -> LogicalExpr: ...
    def __rle__(self, other: Numeric) -> LogicalExpr: ...
    def gt(self, other: Numeric) -> LogicalExpr: ...
    def __gt__(self, other: Numeric) -> LogicalExpr: ...
    def __rgt__(self, other: Numeric) -> LogicalExpr: ...
    def gte(self, other: Numeric) -> LogicalExpr: ...
    def __ge__(self, other: Numeric) -> LogicalExpr: ...
    def __rge__(self, other: Numeric) -> LogicalExpr: ...
    # Arithmetic operators
    def add(self, other: Numeric) -> LogicalExpr: ...
    def __add__(self, other: Numeric) -> LogicalExpr: ...
    def __radd__(self, other: Numeric) -> LogicalExpr: ...
    def sub(self, other: Numeric) -> LogicalExpr: ...
    def __sub__(self, other: Numeric) -> LogicalExpr: ...
    def __rsub__(self, other: Numeric) -> LogicalExpr: ...
    def mul(self, other: Numeric) -> LogicalExpr: ...
    def __mul__(self, other: Numeric) -> LogicalExpr: ...
    def __rmul__(self, other: Numeric) -> LogicalExpr: ...
    def div(self, other: Numeric) -> LogicalExpr: ...
    def __div__(self, other: Numeric) -> LogicalExpr: ...
    def __truediv__(self, other: Numeric) -> LogicalExpr: ...
    def __rdiv__(self, other: Numeric) -> LogicalExpr: ...
    def __rtruediv__(self, other: Numeric) -> LogicalExpr: ...
    # Boolean operators
    def and_(self, other: Boolish) -> LogicalExpr: ...
    def __and__(self, other: Boolish) -> LogicalExpr: ...
    def __rand__(self, other: Boolish) -> LogicalExpr: ...
    def or_(self, other: Boolish) -> LogicalExpr: ...
    def __or__(self, other: Boolish) -> LogicalExpr: ...
    def __ror__(self, other: Boolish) -> LogicalExpr: ...
    # String operators
    def starts_with(self, other: Stringy) -> LogicalExpr: ...
    def contains(self, other: Stringy) -> LogicalExpr: ...

class FunctionExpr: ...

class TextExpr(Enum):
    def __and__(self, other: TextExpr) -> TextExpr: ...
    def __rand__(self, other: TextExpr) -> TextExpr: ...
    def __or__(self, other: TextExpr) -> TextExpr: ...
    def __ror__(self, other: TextExpr) -> TextExpr: ...

class Query:
    def filter(self, expr: LogicalExpr | TextExpr) -> Query: ...
    def topk(
        self, expr: LogicalExpr, k: builtins.int, asc: builtins.bool = False
    ) -> Query: ...
    def rerank(
        self,
        model: typing.Optional[builtins.str] = None,
        query: typing.Optional[builtins.str] = None,
        fields: typing.Sequence[builtins.str] = [],
        topk_multiple: typing.Optional[builtins.int] = None,
    ) -> Query: ...
    def count(self) -> Query: ...

def field(name: builtins.str) -> LogicalExpr: ...
def select(
    *args: builtins.str,
    **kwargs: typing.Union[LogicalExpr, FunctionExpr],
) -> Query: ...
def filter(expr: LogicalExpr | TextExpr) -> Query: ...
def literal(value: typing.Any) -> LogicalExpr: ...
def match(
    token: builtins.str,
    field: builtins.str | None = None,
    weight: builtins.float = 1.0,
    all: builtins.bool = False,
) -> LogicalExpr: ...
def not_(expr: LogicalExpr) -> LogicalExpr: ...

class fn:
    @staticmethod
    def vector_distance(
        field: builtins.str,
        vector: typing.Union[list[int], list[float], topk_sdk.data.Value],
    ) -> FunctionExpr: ...
    @staticmethod
    def semantic_similarity(
        field: builtins.str,
        query: builtins.str,
    ) -> FunctionExpr: ...
    @staticmethod
    def bm25_score() -> FunctionExpr: ...
