"""Pandas expression visitor for therismos expressions.

This module provides a visitor that converts therismos expressions into callable
filter functions compatible with pandas DataFrames.

Requires Pandas to be installed. Install with: pip install therismos[pandas]

Example
-------
>>> import pandas as pd
>>> from therismos import F
>>> from therismos.expr.visitors.pandas import PandasExprVisitor
>>>
>>> age = F("age")
>>> status = F("status")
>>> expr = (age > 18) & (status == "active")
>>>
>>> visitor = PandasExprVisitor()
>>> mask = expr.accept(visitor)
>>>
>>> df = pd.DataFrame({"age": [20, 15, 30], "status": ["active", "active", "inactive"]})
>>> df[mask(df)]

Notes
-----
- Returns a PandasFilter (callable) that accepts a DataFrame and returns a boolean Series
- Closures capture field/value at definition time to avoid loop variable capture issues
- TrueExpr returns Series of True, FalseExpr returns Series of False
- Regex flags are passed directly to pandas str.contains
"""

from __future__ import annotations

from collections.abc import Callable
from typing import TYPE_CHECKING

try:
    import pandas as pd
except ImportError as exc:
    msg = "pandas is required. Install with: pip install therismos[pandas]"
    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

if TYPE_CHECKING:
    pass

PandasFilter = Callable[["pd.DataFrame"], "pd.Series[bool]"]
"""Type alias for a callable that accepts a DataFrame and returns a boolean Series."""


class PandasExprVisitor:
    """Visitor that converts therismos expressions to pandas filter callables.

    This visitor traverses expression trees and generates callables that accept
    a pandas DataFrame and return a boolean Series for filtering.

    Example
    -------
    >>> import pandas as pd
    >>> from therismos import F
    >>> from therismos.expr.visitors.pandas import PandasExprVisitor
    >>>
    >>> age = F("age")
    >>> name = F("name")
    >>> expr = (age >= 21) & (name == "Alice")
    >>>
    >>> visitor = PandasExprVisitor()
    >>> mask = expr.accept(visitor)
    >>> df = pd.DataFrame({"age": [21, 30, 15], "name": ["Alice", "Bob", "Alice"]})
    >>> df[mask(df)]
       age   name
    0   21  Alice
    """

    def visit_eq(self, expr: Eq) -> PandasFilter:
        """Convert equality expression to pandas filter callable.

        :param expr: Equality expression.
        :type expr: Eq
        :returns: Callable that filters rows where field equals value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Eq({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] == value  # type: ignore[no-any-return]

        return _filter

    def visit_ne(self, expr: Ne) -> PandasFilter:
        """Convert inequality expression to pandas filter callable.

        :param expr: Inequality expression.
        :type expr: Ne
        :returns: Callable that filters rows where field does not equal value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Ne({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] != value  # type: ignore[no-any-return]

        return _filter

    def visit_lt(self, expr: Lt) -> PandasFilter:
        """Convert less-than expression to pandas filter callable.

        :param expr: Less-than expression.
        :type expr: Lt
        :returns: Callable that filters rows where field is less than value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Lt({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] < value  # type: ignore[no-any-return]

        return _filter

    def visit_le(self, expr: Le) -> PandasFilter:
        """Convert less-than-or-equal expression to pandas filter callable.

        :param expr: Less-than-or-equal expression.
        :type expr: Le
        :returns: Callable that filters rows where field is less than or equal to value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Le({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] <= value  # type: ignore[no-any-return]

        return _filter

    def visit_gt(self, expr: Gt) -> PandasFilter:
        """Convert greater-than expression to pandas filter callable.

        :param expr: Greater-than expression.
        :type expr: Gt
        :returns: Callable that filters rows where field is greater than value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Gt({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] > value  # type: ignore[no-any-return]

        return _filter

    def visit_ge(self, expr: Ge) -> PandasFilter:
        """Convert greater-than-or-equal expression to pandas filter callable.

        :param expr: Greater-than-or-equal expression.
        :type expr: Ge
        :returns: Callable that filters rows where field is greater than or equal to value.
        :rtype: PandasFilter
        """
        field = expr.field.name
        value = _check_value(expr.casted_value(), f"Ge({expr.field.name})")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field] >= value  # type: ignore[no-any-return]

        return _filter

    def visit_between(self, expr: Between) -> PandasFilter:
        """Convert between expression to a pandas half-open range filter callable.

        :param expr: Between expression (lower <= field < upper).
        :type expr: Between
        :returns: Callable that filters rows where lower <= field < upper.
        :rtype: PandasFilter
        """
        field = 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")

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return (df[field] >= lower) & (df[field] < upper)  # type: ignore[no-any-return]

        return _filter

    def visit_regex(self, expr: Regex) -> PandasFilter:
        """Convert regex expression to pandas filter callable.

        :param expr: Regex expression.
        :type expr: Regex
        :returns: Callable that filters rows where field matches the regex pattern.
        :rtype: PandasFilter
        """
        field = expr.field.name
        pattern = expr.value
        flags = expr.flags if expr.flags is not None else 0

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field].str.contains(pattern, flags=flags, regex=True, na=False)

        return _filter

    def visit_in(self, expr: In) -> PandasFilter:
        """Convert IN expression to pandas filter callable.

        :param expr: IN expression.
        :type expr: In
        :returns: Callable that filters rows where field value is in the set of values.
        :rtype: PandasFilter
        """
        field = expr.field.name
        casted_values = [
            _check_value(expr.field.cast(v), f"In({expr.field.name})") for v in expr.values
        ]

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return df[field].isin(casted_values)

        return _filter

    def visit_is_null(self, expr: IsNull) -> PandasFilter:
        """Convert null-check expression to pandas filter callable.

        :param expr: Null-check expression.
        :type expr: IsNull
        :returns: Callable that filters rows where field is null or not null.
        :rtype: PandasFilter
        """
        field = expr.field.name
        is_null = expr.is_null

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            if is_null:
                return df[field].isna()
            return df[field].notna()

        return _filter

    def visit_true(self, expr: TrueExpr) -> PandasFilter:
        """Convert TRUE constant to pandas filter callable returning all True.

        :param expr: TRUE expression.
        :type expr: TrueExpr
        :returns: Callable that returns a Series of True for all rows.
        :rtype: PandasFilter
        """

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return pd.Series(True, index=df.index)

        return _filter

    def visit_false(self, expr: FalseExpr) -> PandasFilter:
        """Convert FALSE constant to pandas filter callable returning all False.

        :param expr: FALSE expression.
        :type expr: FalseExpr
        :returns: Callable that returns a Series of False for all rows.
        :rtype: PandasFilter
        """

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return pd.Series(False, index=df.index)

        return _filter

    def visit_all(self, expr: AllExpr) -> PandasFilter:
        """Convert AND expression to combined pandas filter callable.

        An empty AllExpr returns a filter that matches all rows.

        :param expr: AND expression.
        :type expr: AllExpr
        :returns: Callable that filters rows matching all sub-expressions.
        :rtype: PandasFilter
        """
        sub_filters = [child.accept(self) for child in expr.exprs]

        if not sub_filters:

            def _true_filter(df: pd.DataFrame) -> pd.Series[bool]:
                return pd.Series(True, index=df.index)

            return _true_filter

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            result = sub_filters[0](df)
            for f in sub_filters[1:]:
                result = result & f(df)
            return result

        return _filter

    def visit_any(self, expr: AnyExpr) -> PandasFilter:
        """Convert OR expression to combined pandas filter callable.

        An empty AnyExpr returns a filter that matches no rows.

        :param expr: OR expression.
        :type expr: AnyExpr
        :returns: Callable that filters rows matching any sub-expression.
        :rtype: PandasFilter
        """
        sub_filters = [child.accept(self) for child in expr.exprs]

        if not sub_filters:

            def _false_filter(df: pd.DataFrame) -> pd.Series[bool]:
                return pd.Series(False, index=df.index)

            return _false_filter

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            result = sub_filters[0](df)
            for f in sub_filters[1:]:
                result = result | f(df)
            return result

        return _filter

    def visit_not(self, expr: NotExpr) -> PandasFilter:
        """Convert NOT expression to negated pandas filter callable.

        :param expr: NOT expression.
        :type expr: NotExpr
        :returns: Callable that negates the child filter.
        :rtype: PandasFilter
        """
        child_filter = expr.expr.accept(self)

        def _filter(df: pd.DataFrame) -> pd.Series[bool]:
            return ~child_filter(df)

        return _filter


__all__ = [
    "PandasExprVisitor",
    "PandasFilter",
]
