from __future__ import annotations

import datetime
from dataclasses import dataclass
from typing import Any, Dict, List, Optional, Tuple, Callable

import dask.dataframe as dd
import pandas as pd
from sqlalchemy import func, cast
from sqlalchemy.sql.sqltypes import Date, Time

from sibi_flux.logger import Logger


# -------------------- Deferred filter expression AST --------------------
class Expr:
    def mask(self, df: dd.DataFrame) -> dd.Series:
        raise NotImplementedError

    def to_parquet_filters(
        self,
    ) -> List[Tuple[str, str, Any] | List[Tuple[str, str, Any]]]:
        # By default, nothing to push down
        return []

    def __and__(self, other: "Expr") -> "Expr":
        return And(self, other)

    def __or__(self, other: "Expr") -> "Expr":
        return Or(self, other)

    def __invert__(self) -> "Expr":
        return Not(self)


@dataclass(frozen=True)
class TrueExpr(Expr):
    """Matches all rows; useful as a neutral starting point."""

    def mask(self, df: dd.DataFrame) -> dd.Series:
        return df.map_partitions(
            lambda p: pd.Series(True, index=p.index), meta=pd.Series(dtype=bool)
        )


@dataclass(frozen=True)
class ColOp(Expr):
    field: str
    casting: Optional[str]
    op: str
    value: Any
    handler: "FilterHandler"  # reuse your parsing + Dask ops

    def mask(self, df: dd.DataFrame) -> dd.Series:
        col = self.handler._get_dask_column(df, self.field, self.casting)
        val = self.handler._parse_filter_value(self.casting, self.value)
        return self.handler._apply_operation_dask(col, self.op, val)

    def to_parquet_filters(self):
        # Only basic comparisons can be pushed down
        if self.op not in self.handler._pushdown_ops():
            return []

        val = self.handler._parse_filter_value(self.casting, self.value)

        # Handle date casting rewrite for parquet compatibility if needed
        # (This class assumes simple ops, but we reuse the handler's logic if possible)
        # For simplicity, ColOp usually represents the post-rewrite state in AST usage.

        if self.op == "exact":
            return [(self.field, "=", val)]
        if self.op == "not_exact":
            return [(self.field, "!=", val)]
        if self.op in {"gt", "gte", "lt", "lte"}:
            sym = {"gt": ">", "gte": ">=", "lt": "<", "lte": "<="}[self.op]
            return [(self.field, sym, val)]
        if self.op == "in":
            return [(self.field, "in", list(val) if not isinstance(val, list) else val)]
        if self.op == "not_in":
            return [
                (self.field, "not in", list(val) if not isinstance(val, list) else val)
            ]
        if self.op == "range":
            lo, hi = val
            return [(self.field, ">=", lo), (self.field, "<=", hi)]
        return []


@dataclass(frozen=True)
class And(Expr):
    left: Expr
    right: Expr

    def mask(self, df: dd.DataFrame) -> dd.Series:
        return self.left.mask(df) & self.right.mask(df)

    def to_parquet_filters(self):
        # AND = concatenate both sides' AND-terms
        return [*self.left.to_parquet_filters(), *self.right.to_parquet_filters()]


@dataclass(frozen=True)
class Or(Expr):
    left: Expr
    right: Expr

    def mask(self, df: dd.DataFrame) -> dd.Series:
        return self.left.mask(df) | self.right.mask(df)

    def to_parquet_filters(self):
        # OR must be returned as list-of-lists; if either side has non-pushdown, defer to mask
        lf, rf = self.left.to_parquet_filters(), self.right.to_parquet_filters()
        if not lf or not rf:
            return []
        return [lf, rf]


@dataclass(frozen=True)
class Not(Expr):
    inner: Expr

    def mask(self, df: dd.DataFrame) -> dd.Series:
        return ~self.inner.mask(df)

    def to_parquet_filters(self):
        return []


# -------------------- Filter handler --------------------
class FilterHandler:
    """
    Handles the application of filters to SQLAlchemy and Dask backends.
    """

    def __init__(self, backend, logger=None, debug=False):
        self.backend = backend
        self.logger = logger or Logger.default_logger(
            logger_name=self.__class__.__name__
        )
        self.logger.set_level(Logger.DEBUG if debug else Logger.INFO)
        self.backend_methods = self._get_backend_methods(backend)
 
    # Constants for faster lookup
    COMPARISON_OPERATORS: ClassVar[List[str]] = [
        "gte", "lte", "gt", "lt", "exact", "in", "range", "contains",
        "startswith", "endswith", "isnull", "not_exact", "not_contains",
        "not_in", "regex", "icontains", "istartswith", "iendswith",
        "iexact", "iregex",
    ]
    DT_OPERATORS: ClassVar[List[str]] = ["date", "time"]
    DATE_OPERATORS: ClassVar[List[str]] = [
        "year", "month", "day", "hour", "minute", "second", "week_day"
    ]

    # --------- Pushdown helpers ---------
    def _pushdown_ops(self) -> set[str]:
        """Ops that can be translated to PyArrow parquet filters."""
        # Added not_exact and not_in for negation support
        return {"exact", "gt", "gte", "lt", "lte", "in", "range", "not_exact", "not_in"}

    def split_pushdown_and_residual(
        self, filters: Dict[str, Any]
    ) -> Tuple[List[Tuple[str, str, Any]], Dict[str, Any]]:
        """
        Split input filter dict into:
          - parquet_filters: list of (col, op, val) tuples for pushdown
          - residual_filters: dict of filters to be applied via Dask mask
        """
        pushdown_candidates = {}
        residual_filters = {}

        for key, value in filters.items():
            field, casting, op = self._parse_filter_key(key)

            # --- OPTIMIZATION: Date Casting Rewrite ---
            # Rewrite col__date__gte to col__gte to enable pushdown
            if casting == "date" and op in {"gt", "gte", "lt", "lte", "range"}:
                try:
                    rewritten = self._rewrite_date_logic(op, value)
                    for new_op, new_val in rewritten.items():
                        pushdown_candidates[f"{field}__{new_op}"] = new_val
                    continue
                except (ValueError, TypeError):
                    pass

            # --- Standard Logic ---
            if casting is not None:
                residual_filters[key] = value
                continue

            if op in self._pushdown_ops():
                pushdown_candidates[key] = value
            else:
                residual_filters[key] = value

        parquet_filters = self.to_parquet_filters(pushdown_candidates)
        return parquet_filters, residual_filters

    def to_parquet_filters(
        self, filters: Optional[Dict[str, Any]] = None
    ) -> List[Tuple[str, str, Any]]:
        """
        Convert a simple dict of {key: val} into PyArrow DNF filters.
        """
        filters = filters or {}
        out: List[Tuple[str, str, Any]] = []

        for key, value in filters.items():
            field, casting, op = self._parse_filter_key(key)

            if casting is not None:
                continue

            if op not in self._pushdown_ops():
                continue

            val = self._parse_filter_value(casting, value)

            if op == "range":
                if isinstance(val, (list, tuple)) and len(val) == 2:
                    lo, hi = val
                    out.extend([(field, ">=", lo), (field, "<=", hi)])
                continue

            # Standard Ops
            if op == "exact":
                out.append((field, "=", val))
            elif op == "gt":
                out.append((field, ">", val))
            elif op == "gte":
                out.append((field, ">=", val))
            elif op == "lt":
                out.append((field, "<", val))
            elif op == "lte":
                out.append((field, "<=", val))

            # Negation Ops
            elif op == "not_exact":
                out.append((field, "!=", val))

            elif op == "in":
                safe_val = list(val) if not isinstance(val, list) else val
                if safe_val:
                    out.append((field, "in", safe_val))

            elif op == "not_in":
                safe_val = list(val) if not isinstance(val, list) else val
                if safe_val:
                    out.append((field, "not in", safe_val))

        return out

    def _rewrite_date_logic(self, op: str, value: Any) -> Dict[str, Any]:
        """Converts date-based logic to timestamp-based logic for pushdown."""

        # Helper to normalize and enforce UTC
        def _to_utc_ts(v):
            ts = pd.Timestamp(v)
            if ts.tz is None:
                # Assume input is meant to be UTC if naive
                ts = ts.tz_localize("UTC")
            else:
                ts = ts.tz_convert("UTC")
            # Normalize to midnight if it's strictly a date comparison
            return ts.normalize()

        if isinstance(value, (list, tuple)):
            dt_val = [_to_utc_ts(v) for v in value]
        else:
            dt_val = _to_utc_ts(value)

        # Logic remains the same, but now operates on UTC-aware timestamps
        if op == "gte":
            return {"gte": dt_val}
        elif op == "gt":
            # Add 1 day for strict greater than on a date casting
            return {"gte": dt_val + pd.Timedelta(days=1)}
        elif op == "lt":
            return {"lt": dt_val}
        elif op == "lte":
            return {"lt": dt_val + pd.Timedelta(days=1)}
        elif op == "range":
            start, end = dt_val
            return {"range": [start, end + pd.Timedelta(days=1)]}

        raise ValueError(f"Cannot rewrite op {op} for pushdown")

    # --------- Expression compiler / mask builder ---------
    def compile_filters(self, filters: Optional[Dict[str, Any]] = None) -> Expr:
        filters = filters or {}
        if not filters:
            return TrueExpr()

        if "$and" in filters:
            expr_and: Expr = TrueExpr()
            for sub in filters["$and"]:
                expr_and = expr_and & self.compile_filters(sub)
            return expr_and

        if "$or" in filters:
            subs = [self.compile_filters(sub) for sub in filters["$or"]]
            if not subs:
                return TrueExpr()
            expr = subs[0]
            for s in subs[1:]:
                expr = expr | s
            return expr

        if "$not" in filters:
            return ~self.compile_filters(filters["$not"])

        expr: Expr = TrueExpr()
        for key, value in filters.items():
            field, casting, op = self._parse_filter_key(key)
            expr = expr & ColOp(
                field=field, casting=casting, op=op, value=value, handler=self
            )
        return expr

    def build_mask_fn(
        self, filters: Optional[Dict[str, Any]] = None
    ) -> Callable[[dd.DataFrame], dd.Series]:
        expr = self.compile_filters(filters)

        def _fn(df: dd.DataFrame) -> dd.Series:
            return expr.mask(df)

        return _fn

    def apply_filters(self, query_or_df, model=None, filters=None):
        filters = filters or {}
        for key, value in filters.items():
            field_name, casting, operation = self._parse_filter_key(key)
            parsed_value = self._parse_filter_value(casting, value)
            if self.backend == "sqlalchemy":
                column = self.backend_methods["get_column"](field_name, model, casting)
                condition = self.backend_methods["apply_operation"](
                    column, operation, parsed_value
                )
                query_or_df = self.backend_methods["apply_condition"](
                    query_or_df, condition
                )
            elif self.backend == "dask":
                column = self.backend_methods["get_column"](
                    query_or_df, field_name, casting
                )
                condition = self.backend_methods["apply_operation"](
                    column, operation, parsed_value
                )
                query_or_df = self.backend_methods["apply_condition"](
                    query_or_df, condition
                )
            else:
                raise ValueError(f"Unsupported backend: {self.backend}")
        return query_or_df

    # --------- Parsing & backend plumbing ---------
    @staticmethod
    def _parse_filter_key(key):
        parts = key.split("__")
        field_name = parts[0]
        casting = None
        operation = "exact"

        if len(parts) == 3:
            # Handle special 'not' modifiers in the middle
            if parts[1] == "not":
                casting = None
                operation = f"not_{parts[2]}"  # e.g. not_in
            else:
                casting = parts[1]
                operation = parts[2]

            # ALIASES
            if operation == "ne":
                operation = "not_exact"
            if operation == "nin":
                operation = "not_in"

        elif len(parts) == 2:
            second = parts[1]

            # ALIASES
            if second == "ne":
                operation = "not_exact"
            elif second == "nin":
                operation = "not_in"
            elif second in FilterHandler.COMPARISON_OPERATORS:
                operation = second
            elif (
                second
                in FilterHandler.DT_OPERATORS + FilterHandler.DATE_OPERATORS
            ):
                casting = second
            else:
                operation = second

        return field_name, casting, operation

    @staticmethod
    def _parse_filter_value(casting, value):
        if casting == "date":
            if isinstance(value, str):
                return pd.Timestamp(value)
            if isinstance(value, list):
                return [pd.Timestamp(v) for v in value]
        elif casting == "time":
            if isinstance(value, list):
                return [FilterHandler._time_to_seconds(v) for v in value]
            return FilterHandler._time_to_seconds(value)
        return value

    @staticmethod
    def _get_backend_methods(backend):
        if backend == "sqlalchemy":
            return {
                "get_column": FilterHandler._get_sqlalchemy_column,
                "apply_operation": FilterHandler._apply_operation_sqlalchemy,
                "apply_condition": lambda query, condition: query.filter(condition),
            }
        elif backend == "dask":
            return {
                "get_column": FilterHandler._get_dask_column,
                "apply_operation": FilterHandler._apply_operation_dask,
                "apply_condition": lambda df, condition: df[condition],
            }
        else:
            raise ValueError(f"Unsupported backend: {backend}")

    @staticmethod
    def _get_sqlalchemy_column(field_name, model, casting):
        column = getattr(model, field_name, None)
        if not column:
            raise AttributeError(
                f"Field '{field_name}' not found in model '{model.__name__}'"
            )
        if casting == "date":
            column = cast(column, Date)
        elif casting == "time":
            column = cast(column, Time)
        elif casting in FilterHandler._date_operators():
            column = func.extract(casting, column)
        return column

    @staticmethod
    def _get_dask_column(df, field_name, casting):
        needs_dt = casting in (
            FilterHandler.DT_OPERATORS + FilterHandler.DATE_OPERATORS
        )
        column = (
            dd.to_datetime(df[field_name], errors="coerce")
            if needs_dt
            else df[field_name]
        )
        if needs_dt:
            column = FilterHandler._strip_tz(column)
        if casting == "date":
            column = column.dt.floor("D")
        elif casting == "time":
            column = column.dt.hour * 3600 + column.dt.minute * 60 + column.dt.second
        elif casting in FilterHandler._date_operators():
            attr = "weekday" if casting == "week_day" else casting
            column = getattr(column.dt, attr)
        return column

    @staticmethod
    def _apply_operation_sqlalchemy(column, operation, value):
        operation_map = FilterHandler._operation_map_sqlalchemy()
        if operation not in operation_map:
            raise ValueError(f"Unsupported operation: {operation}")
        return operation_map[operation](column, value)

    @staticmethod
    def _apply_operation_dask(column, operation, value):
        operation_map = FilterHandler._operation_map_dask()
        if operation not in operation_map:
            raise ValueError(f"Unsupported operation: {operation}")
        return operation_map[operation](column, value)

    @staticmethod
    def _operation_map_sqlalchemy():
        return {
            "exact": lambda col, val: col == val,
            "gt": lambda col, val: col > val,
            "gte": lambda col, val: col >= val,
            "lt": lambda col, val: col < val,
            "lte": lambda col, val: col <= val,
            "in": lambda col, val: col.in_(val),
            "range": lambda col, val: col.between(val[0], val[1]),
            "contains": lambda col, val: col.like(f"%{val}%"),
            "startswith": lambda col, val: col.like(f"{val}%"),
            "endswith": lambda col, val: col.like(f"%{val}"),
            "isnull": lambda col, val: col.is_(None) if val else col.isnot(None),
            "not_exact": lambda col, val: col != val,
            "not_contains": lambda col, val: ~col.like(f"%{val}%"),
            "not_in": lambda col, val: ~col.in_(val),
            "regex": lambda col, val: col.op("~")(val),
            "icontains": lambda col, val: col.ilike(f"%{val}%"),
            "istartswith": lambda col, val: col.ilike(f"{val}%"),
            "iendswith": lambda col, val: col.ilike(f"%{val}"),
            "iexact": lambda col, val: col.ilike(val),
            "iregex": lambda col, val: col.op("~*")(val),
        }

    @staticmethod
    def _operation_map_dask():
        return {
            "exact": lambda col, val: col == val,
            "gt": lambda col, val: col > val,
            "gte": lambda col, val: col >= val,
            "lt": lambda col, val: col < val,
            "lte": lambda col, val: col <= val,
            "in": lambda col, val: FilterHandler._apply_isin(col, val, negated=False),
            "not_in": lambda col, val: FilterHandler._apply_isin(col, val, negated=True),
            "range": lambda col, val: (col >= val[0]) & (col <= val[1]),
            "contains": lambda col, val: FilterHandler._as_str(col).str.contains(
                val, regex=True, na=False
            ),
            "startswith": lambda col, val: FilterHandler._as_str(col).str.startswith(
                val, na=False
            ),
            "endswith": lambda col, val: FilterHandler._as_str(col).str.endswith(
                val, na=False
            ),
            "not_contains": lambda col, val: ~FilterHandler._as_str(col).str.contains(
                val, regex=True, na=False
            ),
            "regex": lambda col, val: FilterHandler._as_str(col).str.contains(
                val, regex=True, na=False
            ),
            "icontains": lambda col, val: FilterHandler._as_str(col).str.contains(
                val, case=False, regex=True, na=False
            ),
            "istartswith": lambda col, val: FilterHandler._as_str(col)
            .str.lower()
            .str.startswith(str(val).lower(), na=False),
            "iendswith": lambda col, val: FilterHandler._as_str(col)
            .str.lower()
            .str.endswith(str(val).lower(), na=False),
            "iexact": lambda col, val: FilterHandler._as_str(col).str.lower()
            == str(val).lower(),
            "iregex": lambda col, val: FilterHandler._as_str(col).str.contains(
                val, case=False, regex=True, na=False
            ),
            "isnull": lambda col, val: col.isnull() if val else col.notnull(),
            "not_exact": lambda col, val: col != val,
        }

    @staticmethod
    def _as_str(col):
        return col.astype("string").fillna("")

    @staticmethod
    def _strip_tz(col):

        def _part(s: pd.Series) -> pd.Series:
            try:
                return s.dt.tz_convert("UTC").dt.tz_localize(None)
            except Exception:
                try:
                    return s.dt.tz_localize(None)
                except Exception:
                    return s

        return col.map_partitions(_part, meta=col._meta)

    @staticmethod
    def _apply_isin(col, val, negated: bool = False):
        aligned_col, aligned_vals = FilterHandler._align_in_types(col, val)
        mask = aligned_col.isin(aligned_vals)
        return ~mask if negated else mask

    @staticmethod
    def _time_to_seconds(t):
        if isinstance(t, str):
            t = datetime.time.fromisoformat(t)
        return t.hour * 3600 + t.minute * 60 + t.second

    @staticmethod
    def _dt_operators():
        return ["date", "time"]

    @staticmethod
    def _date_operators():
        return ["year", "month", "day", "hour", "minute", "second", "week_day"]

    @staticmethod
    def _comparison_operators():
        return [
            "gte",
            "lte",
            "gt",
            "lt",
            "exact",
            "in",
            "range",
            "contains",
            "startswith",
            "endswith",
            "isnull",
            "not_exact",
            "not_contains",
            "not_in",
            "regex",
            "icontains",
            "istartswith",
            "iendswith",
            "iexact",
            "iregex",
        ]

    @staticmethod
    def _align_in_types(col, val):
        if isinstance(val, (set, tuple)):
            vals = list(val)
        elif isinstance(val, list):
            vals = val
        else:
            vals = [val]
        kind = getattr(getattr(col, "dtype", None), "kind", None)
        if kind in ("i", "u"):
            try:
                return col.astype("Int64"), [int(x) for x in vals]
            except Exception:
                pass
        if kind in ("f",):
            try:
                return col.astype("float64"), [float(x) for x in vals]
            except Exception:
                pass
        return FilterHandler._as_str(col), [str(x) for x in vals]
