from __future__ import annotations
import datetime
from typing import Dict
import polars as pl
from tools.base import BaseTool


class FilterTool(BaseTool):
    def execute(self, inputs: Dict[str, pl.DataFrame]) -> Dict[str, pl.DataFrame]:
        df = inputs.get("Input", pl.DataFrame())
        if df.is_empty():
            return {"True": df, "False": df}

        mode = self._cfg("Mode", "Custom") or "Custom"

        if mode == "Simple":
            mask = self._build_simple_mask(df)
        else:
            expr_text = self._cfg("Expression", "True") or "True"
            try:
                mask = self.ctx.transpiler.eval_mask(df, expr_text)
            except Exception as e:
                raise RuntimeError(
                    f"Filter expression failed: {expr_text!r} → {e}"
                ) from e

        true_df = df.filter(mask)
        false_df = df.filter(~mask)
        return {"True": true_df, "False": false_df}

    def _build_simple_mask(self, df: pl.DataFrame) -> pl.Series:
        simple = self.config.find("Simple") if self.config is not None else None
        if simple is None:
            return pl.Series([True] * len(df))

        operator = (simple.findtext("Operator") or "=").strip()
        field = (simple.findtext("Field") or "").strip()
        operands = simple.find("Operands")
        raw_value = (operands.findtext("Operand") if operands is not None else None) or ""
        date_type = (operands.findtext("DateType") if operands is not None else None) or "fixed"

        value = self._resolve_date_value(raw_value, date_type, operands)

        if field not in df.columns:
            return pl.Series([False] * len(df))

        col = pl.col(field)

        # Coerce string value to match column dtype
        col_dtype = df[field].dtype
        typed_value = self._coerce_value(value, col_dtype)

        op_map = {
            "=": lambda: df.select(col == typed_value).to_series(),
            "Equals": lambda: df.select(col == typed_value).to_series(),
            "!=": lambda: df.select(col != typed_value).to_series(),
            "Does not equal": lambda: df.select(col != typed_value).to_series(),
            ">": lambda: df.select(col > typed_value).to_series(),
            ">=": lambda: df.select(col >= typed_value).to_series(),
            "<": lambda: df.select(col < typed_value).to_series(),
            "<=": lambda: df.select(col <= typed_value).to_series(),
            "IsNull": lambda: df.select(col.is_null()).to_series(),
            "IsNotNull": lambda: df.select(col.is_not_null()).to_series(),
            "Is null": lambda: df.select(col.is_null()).to_series(),
            "Is not null": lambda: df.select(col.is_not_null()).to_series(),
            "Contains": lambda: df.select(col.cast(pl.String).str.contains(str(value))).to_series(),
            "Does not contain": lambda: df.select(~col.cast(pl.String).str.contains(str(value))).to_series(),
            "Is empty": lambda: df.select(col.is_null() | (col.cast(pl.String) == "")).to_series(),
            "Is not empty": lambda: df.select(col.is_not_null() & (col.cast(pl.String) != "")).to_series(),
            "Comes before (<)": lambda: df.select(col < value).to_series(),
            "Comes after (>)": lambda: df.select(col > value).to_series(),
            "Is true": lambda: df.select(col == True).to_series(),  # noqa: E712
            "Is false": lambda: df.select(col == False).to_series(),  # noqa: E712
            "StartsWith": lambda: df.select(col.cast(pl.String).str.starts_with(str(value))).to_series(),
            "EndsWith": lambda: df.select(col.cast(pl.String).str.ends_with(str(value))).to_series(),
        }

        fn = op_map.get(operator)
        if fn is not None:
            return fn()

        if operator == "Range":
            operands_list = operands.findall("Operand") if operands is not None else []
            if len(operands_list) >= 2:
                lo = self._coerce_value(operands_list[0].text or "", col_dtype)
                hi = self._coerce_value(operands_list[1].text or "", col_dtype)
                return df.select((col >= lo) & (col <= hi)).to_series()

        # Fallback
        return pl.Series([True] * len(df))

    def _coerce_value(self, value: str, dtype: pl.PolarsDataType) -> object:
        if dtype in (pl.Int8, pl.Int16, pl.Int32, pl.Int64,
                     pl.UInt8, pl.UInt16, pl.UInt32, pl.UInt64):
            try:
                return int(value)
            except (ValueError, TypeError):
                return value
        if dtype in (pl.Float32, pl.Float64):
            try:
                return float(value)
            except (ValueError, TypeError):
                return value
        if dtype == pl.Boolean:
            return value.lower() in ("true", "1", "yes")
        return value

    def _resolve_date_value(self, raw: str, date_type: str, operands) -> object:
        today = datetime.date.today()
        if date_type == "today":
            return str(today)
        if date_type == "tomorrow":
            return str(today + datetime.timedelta(days=1))
        if date_type == "yesterday":
            return str(today - datetime.timedelta(days=1))

        period_type = (operands.findtext("PeriodType") if operands is not None else None) or "days"
        period_count_str = (operands.findtext("PeriodCount") if operands is not None else None) or "0"
        try:
            period_count = int(period_count_str)
        except ValueError:
            period_count = 0

        if date_type in ("PeriodAfter", "PeriodBefore"):
            delta = datetime.timedelta(**{period_type: period_count})
            anchor = today + delta if date_type == "PeriodAfter" else today - delta
            return str(anchor)

        return raw
