"""Polars expression visitor for therismos expressions.

This module provides a visitor that converts therismos expressions into Polars
expressions compatible with both eager DataFrames and lazy LazyFrames.

Requires Polars to be installed. Install with: pip install therismos[polars]

Example
-------
>>> from therismos import F
>>> from therismos.expr.visitors.polars import PolarsExprVisitor
>>>
>>> age = F("age")
>>> status = F("status")
>>> expr = (age > 18) & (status == "active")
>>>
>>> visitor = PolarsExprVisitor()
>>> pl_expr = expr.accept(visitor)
>>>
>>> # Use with eager DataFrame
>>> # df.filter(pl_expr)
>>>
>>> # Use with lazy LazyFrame
>>> # lf.filter(pl_expr)

Notes
-----
- Regex flags are converted to inline modifiers ((?i), (?m), (?s))
- TrueExpr converts to pl.lit(True), FalseExpr to pl.lit(False)
- Empty AllExpr returns pl.lit(True), empty AnyExpr returns pl.lit(False)
"""

from __future__ import annotations

import functools
import re

try:
    import polars as pl
except ImportError as exc:
    msg = "polars is required. Install with: pip install therismos[polars]"
    raise ImportError(msg) from exc

from therismos.expr._expr import (
    AllExpr,
    AnyExpr,
    Between,
    Eq,
    FalseExpr,
    Ge,
    Gt,
    In,
    IsNull,
    Le,
    Lt,
    Ne,
    NotExpr,
    Regex,
    TrueExpr,
)
from therismos.expr.visitors._visitors import _check_value


class PolarsExprVisitor:
    """Visitor that converts therismos expressions to Polars expressions.

    This visitor traverses expression trees and generates Polars expressions
    that can be used to filter both eager DataFrames and lazy LazyFrames.

    Example
    -------
    >>> import polars as pl
    >>> from therismos import F
    >>> from therismos.expr.visitors.polars import PolarsExprVisitor
    >>>
    >>> age = F("age")
    >>> name = F("name")
    >>> expr = (age >= 21) & (name == "Alice")
    >>>
    >>> visitor = PolarsExprVisitor()
    >>> pl_expr = expr.accept(visitor)
    >>> df = pl.DataFrame({"age": [21, 30, 15], "name": ["Alice", "Bob", "Alice"]})
    >>> df.filter(pl_expr)
    shape: (1, 2)
    """

    def visit_eq(self, expr: Eq) -> pl.Expr:
        """Convert equality expression to Polars expression.

        :param expr: Equality expression.
        :type expr: Eq
        :returns: Polars expression like pl.col(field) == value.
        :rtype: pl.Expr
        """
        val = _check_value(expr.casted_value(), f"Eq({expr.field.name})")
        return pl.col(expr.field.name) == val  # type: ignore[no-any-return]

    def visit_ne(self, expr: Ne) -> pl.Expr:
        """Convert inequality expression to Polars expression.

        :param expr: Inequality expression.
        :type expr: Ne
        :returns: Polars expression like pl.col(field) != value.
        :rtype: pl.Expr
        """
        val = _check_value(expr.casted_value(), f"Ne({expr.field.name})")
        return pl.col(expr.field.name) != val  # type: ignore[no-any-return]

    def visit_lt(self, expr: Lt) -> pl.Expr:
        """Convert less-than expression to Polars expression.

        :param expr: Less-than expression.
        :type expr: Lt
        :returns: Polars expression like pl.col(field) < value.
        :rtype: pl.Expr
        """
        return pl.col(expr.field.name) < _check_value(  # type: ignore[no-any-return]
            expr.casted_value(), f"Lt({expr.field.name})"
        )

    def visit_le(self, expr: Le) -> pl.Expr:
        """Convert less-than-or-equal expression to Polars expression.

        :param expr: Less-than-or-equal expression.
        :type expr: Le
        :returns: Polars expression like pl.col(field) <= value.
        :rtype: pl.Expr
        """
        val = _check_value(expr.casted_value(), f"Le({expr.field.name})")
        return pl.col(expr.field.name) <= val  # type: ignore[no-any-return]

    def visit_gt(self, expr: Gt) -> pl.Expr:
        """Convert greater-than expression to Polars expression.

        :param expr: Greater-than expression.
        :type expr: Gt
        :returns: Polars expression like pl.col(field) > value.
        :rtype: pl.Expr
        """
        return pl.col(expr.field.name) > _check_value(  # type: ignore[no-any-return]
            expr.casted_value(), f"Gt({expr.field.name})"
        )

    def visit_ge(self, expr: Ge) -> pl.Expr:
        """Convert greater-than-or-equal expression to Polars expression.

        :param expr: Greater-than-or-equal expression.
        :type expr: Ge
        :returns: Polars expression like pl.col(field) >= value.
        :rtype: pl.Expr
        """
        val = _check_value(expr.casted_value(), f"Ge({expr.field.name})")
        return pl.col(expr.field.name) >= val  # type: ignore[no-any-return]

    def visit_between(self, expr: Between) -> pl.Expr:
        """Convert between expression to a Polars half-open range filter.

        :param expr: Between expression (lower <= field < upper).
        :type expr: Between
        :returns: Polars expression ``(col >= lower) & (col < upper)``.
        :rtype: pl.Expr
        """
        col = pl.col(expr.field.name)
        lower = _check_value(expr.casted_lower(), f"Between({expr.field.name}).lower")
        upper = _check_value(expr.casted_upper(), f"Between({expr.field.name}).upper")
        return (col >= lower) & (col < upper)  # type: ignore[no-any-return]

    def visit_regex(self, expr: Regex) -> pl.Expr:
        """Convert regex expression to Polars string contains expression.

        Regex flags are converted to inline modifiers prepended to the pattern:
        - re.IGNORECASE -> (?i)
        - re.MULTILINE -> (?m)
        - re.DOTALL -> (?s)

        :param expr: Regex expression.
        :type expr: Regex
        :returns: Polars expression using str.contains.
        :rtype: pl.Expr
        """
        pattern = expr.value

        if expr.flags is not None:
            modifiers = ""
            if expr.flags & re.IGNORECASE:
                modifiers += "i"
            if expr.flags & re.MULTILINE:
                modifiers += "m"
            if expr.flags & re.DOTALL:
                modifiers += "s"
            if modifiers:
                pattern = f"(?{modifiers}){pattern}"

        return pl.col(expr.field.name).str.contains(pattern)

    def visit_in(self, expr: In) -> pl.Expr:
        """Convert IN expression to Polars is_in expression.

        :param expr: IN expression.
        :type expr: In
        :returns: Polars expression using is_in.
        :rtype: pl.Expr
        """
        casted_values = [
            _check_value(expr.field.cast(v), f"In({expr.field.name})") for v in expr.values
        ]
        return pl.col(expr.field.name).is_in(casted_values)

    def visit_is_null(self, expr: IsNull) -> pl.Expr:
        """Convert null-check expression to Polars null check.

        :param expr: Null-check expression.
        :type expr: IsNull
        :returns: Polars expression using is_null or is_not_null.
        :rtype: pl.Expr
        """
        if expr.is_null:
            return pl.col(expr.field.name).is_null()
        return pl.col(expr.field.name).is_not_null()

    def visit_true(self, expr: TrueExpr) -> pl.Expr:
        """Convert TRUE constant to Polars literal True.

        :param expr: TRUE expression.
        :type expr: TrueExpr
        :returns: pl.lit(True).
        :rtype: pl.Expr
        """
        return pl.lit(True)

    def visit_false(self, expr: FalseExpr) -> pl.Expr:
        """Convert FALSE constant to Polars literal False.

        :param expr: FALSE expression.
        :type expr: FalseExpr
        :returns: pl.lit(False).
        :rtype: pl.Expr
        """
        return pl.lit(False)

    def visit_all(self, expr: AllExpr) -> pl.Expr:
        """Convert AND expression to Polars conjunction.

        An empty AllExpr returns pl.lit(True).

        :param expr: AND expression.
        :type expr: AllExpr
        :returns: Polars expression using & operator.
        :rtype: pl.Expr
        """
        sub_exprs = [child.accept(self) for child in expr.exprs]
        if not sub_exprs:
            return pl.lit(True)
        return functools.reduce(lambda a, b: a & b, sub_exprs)

    def visit_any(self, expr: AnyExpr) -> pl.Expr:
        """Convert OR expression to Polars disjunction.

        An empty AnyExpr returns pl.lit(False).

        :param expr: OR expression.
        :type expr: AnyExpr
        :returns: Polars expression using | operator.
        :rtype: pl.Expr
        """
        sub_exprs = [child.accept(self) for child in expr.exprs]
        if not sub_exprs:
            return pl.lit(False)
        return functools.reduce(lambda a, b: a | b, sub_exprs)

    def visit_not(self, expr: NotExpr) -> pl.Expr:
        """Convert NOT expression to Polars negation.

        :param expr: NOT expression.
        :type expr: NotExpr
        :returns: Polars expression using ~ operator.
        :rtype: pl.Expr
        """
        return ~expr.expr.accept(self)


__all__ = [
    "PolarsExprVisitor",
]
