"""
Note: This is mostly generated by GPT5, and checked via manual tests.
The logic has not been manually checked in detail.

Example:
    >>> from kwdagger.mlops.query_plan import *  # NOQA
    >>> import pandas as pd
    >>> import ubelt as ub
    >>> #
    >>> table1 = pd.DataFrame({
    >>>     "a": [1, 2, 3, 4],
    >>>     "b": ["foo", "bar", "foo", "baz"],
    >>>     "c": [10, 20, 30, 40],
    >>> })
    >>> plan = QueryPlan.parse('(df.b.str.contains("foo")) & (df.a > 2)', strict=True)
    >>> print(plan.describe())
    >>> assert len(plan.apply(table1)) == 1
    >>> assert len(plan.apply_all({'table1': table1})['table1']) == 1
    >>> #
    >>> table2 = pd.DataFrame({
    >>>     "a": [1, 2, 3, 4, 5, 6],
    >>>     "b": ["foo", "bar", "foo", "baz", "foo", "foo"],
    >>>     "c": [10, 20, 30, 40, 50, 60],
    >>>     "d": [10, 20, 30, 40, 60, 70],
    >>> })
    >>> # YAML list query (implicit AND)
    >>> query_str = '''
    >>>     - df['a'] > 1
    >>>     - df['b'] == 'foo'
    >>> '''
    >>> #
    >>> plan = QueryPlan.parse(query_str, strict=False)
    >>> #
    >>> # 2) Apply to a single node df:
    >>> filtered_df = plan.apply(table1, "table1")
    >>> assert len(filtered_df) == 1
    >>> #
    >>> # 3) Or apply to many:
    >>> node_to_df = {"table1": table1, "table2": table2}
    >>> filtered = plan.apply_all(node_to_df)
    >>> assert ub.udict.map_values(filtered, len) == {'table1': 1, 'table2': 3}

    >>> # Check that comments work
    >>> query_str = '''
    >>>     - df['a'] > 1
    >>>     - df['b'] == 'foo'
    >>>     # - df['a'] < 1
    >>> '''
    >>> #
    >>> plan = QueryPlan.parse(query_str, strict=False)
    >>> #
    >>> # 2) Apply to a single node df:
    >>> filtered_df = plan.apply(table1, "table1")
    >>> assert len(filtered_df) == 1

    >>> # More complex example
    >>> query_str = '''
    >>> __all__:
    >>>   - df['c'] >= 20
    >>> table2:
    >>>   - and:
    >>>       - df['c'] != 20
    >>>   - or:
    >>>       - df['c'] == df['d']
    >>>       - df['c'] == 60
    >>> '''
    >>> plan = QueryPlan.parse(query_str, strict=False)
    >>> print(plan.describe())
    >>> node_to_df = {"table1": table1, "table2": table2}
    >>> filtered = plan.apply_all(node_to_df)
    >>> print(f'filtered = {ub.urepr(filtered, nl=1)}')
    >>> assert ub.udict.map_values(filtered, len) == {'table1': 3, 'table2': 3}

    >>> # Bad expression is skipped in non-strict mode
    >>> yaml_str = '''
    ... __all__:
    ...   - df['does_not_exist'] > 0
    ...   - df['a'] > 2
    ... '''
    >>> df = pd.DataFrame({
    ...     "a": [1, 2, 3, 4],
    ...     "b": ["foo", "bar", "foo", "baz"],
    ...     "c": [10, 20, 30, 40],
    ... })
    >>> plan = QueryPlan.parse(yaml_str, strict=False)
    >>> plan.apply(df, "nodeY")["a"].tolist()

Example:
    >>> from kwdagger.mlops.query_plan import *  # NOQA
    >>> import pandas as pd
    >>> # Test case with lists and str maniplation
    >>> yaml_str = '''
    ... __all__:
    ...   - df["a"].apply(str).str.contains("[3, 4, 5]", regex=False)
    ... '''
    >>> df = pd.DataFrame({
    ...     "a": [[1, 2, 3], [3, 4, 5]],
    ...     "b": ["foo", "bar"],
    ...     "c": [10, 20],
    ... })
    >>> plan = QueryPlan.parse(yaml_str, strict=False)
    >>> plan.apply(df)
"""

from __future__ import annotations
from dataclasses import dataclass, field
from types import SimpleNamespace
from typing import Dict, List, Literal, Optional, Sequence, Tuple, Union
import warnings
import pandas as pd
import numpy as np

try:
    import yaml  # PyYAML
except Exception:  # pragma: no cover
    yaml = None


GroupType = Literal["and", "or"]
Expr = str
Group = Tuple[GroupType, List[Expr]]


@dataclass
class QueryPlan:
    """
    Holds a normalized, per-node query plan and can apply it to DataFrames.

    Plan format:
      plan["__all__"] -> list[Group]
      plan["node_name"] -> list[Group]

    Group is ("and" | "or", [expr1, expr2, ...])
    - "and" groups are applied progressively (AND semantics).
    - "or"  groups are evaluated against the current df, combined via OR, then applied.

    Evaluation:
    - Full Python eval, with aliases: df, table, t, agg.table
    - Expressions may return:
        * boolean mask aligned to df.index,
        * index/labels (iterable) selecting rows,
        * a filtered DataFrame (its index will be used),
        * truthy/falsey scalar is INVALID (raises).
    - On error (e.g., missing column), behavior:
        * strict=False: warn and skip that expression
        * strict=True: raise
    """
    plan: Dict[str, List[Group]] = field(default_factory=dict)
    strict: bool = True  # set True to fail fast on bad expressions

    @classmethod
    def parse(QueryPlan, cli_arg: Optional[str], *, strict: bool = True) -> QueryPlan:
        """
        Parse the CLI --query argument (YAML or raw) into a QueryPlan.

        Accepted shapes:
        1) None / empty -> no-op plan
        2) RAW STRING   -> apply to __all__ as a single AND group
        3) YAML scalar  -> same as RAW STRING
        4) YAML list    -> __all__: AND of each element (strings)
        5) YAML mapping -> keys: '__all__' or explicit node names
                           values can be:
                             - list[str]            -> AND group
                             - {'and': [str,...]}   -> AND group
                             - {'or':  [str,...]}   -> OR group
                             - {'and': [...], 'or': [...]} -> both groups (AND applied first, then OR)
                             - list[dict|str]       -> mixed; each item is either str (AND) or dict('and'|'or')

        Examples:
          # Raw string (back-compat)
          --query="df['params.key'] == 3"

          # YAML list (global AND chain)
          --query='
          - df["a"] > 0
          - df["b"].str.contains("foo")
          '

          # Per-node with OR group
          --query='
          __all__:
            - df["common"].notna()
          eval1:
            and:
              - df["params.key1"] == 3
            or:
              - df["params.key1"] == 7
              - df["params.key1"] == 9
          eval2:
            - df["params.key2"] == 1
          '
        """
        if not cli_arg or not str(cli_arg).strip():
            return QueryPlan(plan={}, strict=strict)

        data = None
        if yaml is not None:
            try:
                data = yaml.safe_load(cli_arg)
            except Exception:
                # Not valid YAML; treat as raw string
                data = None

        if data is None:
            # Raw string → __all__ AND group
            plan = {"__all__": [("and", [str(cli_arg)])]}
            return QueryPlan(plan=plan, strict=strict)

        # YAML scalar -> same as raw string
        if isinstance(data, str):
            plan = {"__all__": [("and", [data])]}
            return QueryPlan(plan=plan, strict=strict)

        # YAML list -> __all__ AND chain
        if isinstance(data, list):
            exprs = [str(x) for x in data]
            plan = {"__all__": [("and", exprs)]}
            return QueryPlan(plan=plan, strict=strict)

        # YAML mapping -> per-node
        if isinstance(data, dict):
            plan: Dict[str, List[Group]] = {}
            for key, val in data.items():
                if key == "options":
                    # currently ignored; you can extend (e.g., strict/error modes) here
                    continue
                node = str(key)
                groups: List[Group] = _coerce_value_to_groups(val)
                if groups:
                    plan[node] = groups
            return QueryPlan(plan=plan, strict=strict)

        raise TypeError(f"Unsupported --query type: {type(data).__name__}")

    # ----------------------------
    # Public API
    # ----------------------------
    def apply(self, df: pd.DataFrame, node: str = '__all__') -> pd.DataFrame:
        """Apply __all__ + node-specific groups to df and return a filtered view."""
        groups = self._groups_for(node)
        if not groups:
            return df

        current = df
        for gtype, exprs in groups:
            if gtype == "and":
                for expr in exprs:
                    mask = self._eval_expr_as_mask(current, expr, node=node)
                    if mask is None:  # skipped due to warning in non-strict mode
                        continue
                    current = current.loc[mask]
            elif gtype == "or":
                masks = []
                for expr in exprs:
                    mask = self._eval_expr_as_mask(current, expr, node=node)
                    if mask is None:  # skip this branch of the OR
                        continue
                    masks.append(mask)
                if masks:
                    combo = masks[0]
                    for m in masks[1:]:
                        combo = combo | m.reindex(combo.index, fill_value=False)
                    current = current.loc[combo]
                else:
                    # all exprs failed/skipped; keep df unchanged
                    pass
            else:
                raise ValueError(f"Unknown group type: {gtype!r}")
        return current

    def apply_all(self, node_to_df: Dict[str, pd.DataFrame]) -> Dict[str, pd.DataFrame]:
        """Apply the plan to many nodes in one go."""
        out: Dict[str, pd.DataFrame] = {}
        for node, df in node_to_df.items():
            out[node] = self.apply(df, node=node)
        return out

    # NEW: index-returning variants
    def apply_index(self, df: pd.DataFrame, node: str = '__all__') -> pd.Index:
        """
        Efficiently compute the filtered index for `node` without repeatedly slicing `df`.
        Builds a cumulative boolean mask over the ORIGINAL df:
          - AND groups: cumulative &= mask
          - OR  groups: cumulative &= (mask1 | mask2 | ...)
        Notes:
          * Identical semantics for pure row-filter expressions.
          * If an expression is skipped (non-strict error), it is simply omitted.
        """
        groups = self._groups_for(node)
        if not groups:
            return df.index

        # Start with "all True"
        cumulative = pd.Series(True, index=df.index)

        for gtype, exprs in groups:
            if gtype == "and":
                for expr in exprs:
                    mask = self._eval_expr_as_mask(df, expr, node=node)  # eval on full df
                    if mask is None:
                        continue
                    # align and AND into cumulative
                    mask = mask.reindex(df.index, fill_value=False)
                    cumulative &= mask
            elif gtype == "or":
                or_mask = None
                for expr in exprs:
                    mask = self._eval_expr_as_mask(df, expr, node=node)  # eval on full df
                    if mask is None:
                        continue
                    mask = mask.reindex(df.index, fill_value=False)
                    or_mask = mask if or_mask is None else (or_mask | mask)
                if or_mask is not None:
                    cumulative &= or_mask
                else:
                    # all OR branches skipped → no change to cumulative
                    pass
            else:
                raise ValueError(f"Unknown group type: {gtype!r}")

            # Early exit: if cumulative already all False, nothing else can bring rows back
            if not cumulative.any():
                return df.index[:0]

        return df.index[cumulative.to_numpy()]

    # ----------------------------
    # Internals
    # ----------------------------
    def _groups_for(self, node: str) -> List[Group]:
        groups: List[Group] = []
        if "__all__" in self.plan:
            groups.extend(self.plan["__all__"])
        if node in self.plan:
            groups.extend(self.plan[node])
        return groups

    def _eval_expr_as_mask(self, df: pd.DataFrame, expr: str, *, node: str) -> Optional[pd.Series]:
        """
        Evaluate an expression into a boolean mask aligned to df.index.
        Returns None if skipped (strict=False and an error occurs).
        """
        try:
            # Build a minimal, explicit eval environment
            agg = SimpleNamespace(table=df)
            ns = globals()
            ns = locals()
            local_env = {
                "df": df,
                "table": df,
                "t": df,
                "agg": agg,
                "pd": pd,
                "np": np,
                "True": True,
                "False": False,
                "None": None,
                'str': str,
            }
            # Fully-python eval by requirement; Unsafe, could tighten up,
            # but we usually will need access to above data.
            val = eval(expr, ns | local_env, ns | local_env)
        except Exception as ex:
            msg = f"[query:{node}] skipped expr due to error: {expr!r} -> {type(ex).__name__}: {ex}"
            if self.strict:
                raise RuntimeError(msg) from ex
            warnings.warn(msg)
            return None

        # Normalize to mask
        try:
            if isinstance(val, pd.DataFrame):
                idx = val.index
                mask = df.index.isin(idx)
            elif isinstance(val, pd.Series) and val.dtype == bool and val.index.equals(df.index):
                mask = val
            elif isinstance(val, (pd.Index, list, tuple, np.ndarray)):
                mask = df.index.isin(list(val))
            else:
                raise TypeError(
                    "Expression must return a filtered DataFrame, a boolean Series aligned to df.index, "
                    "or an index/iterable of row labels."
                )
            # Ensure boolean Series aligned to df
            mask = pd.Series(mask, index=df.index)
            mask = mask.fillna(False).astype(bool)
            return mask
        except Exception as ex:
            msg = f"[query:{node}] invalid expr result for {expr!r}: {type(val).__name__} ({ex})"
            if self.strict:
                raise RuntimeError(msg) from ex
            warnings.warn(msg)
            return None

    # ----------------------------
    # Introspection / Presentation
    # ----------------------------
    def groups_for_node(self, node: str) -> List[Group]:
        """
        Return the merged list of groups that will be applied to `node`
        (i.e., __all__ groups + node-specific groups).
        """
        return self._groups_for(node)

    def describe(self, nodes: Optional[Sequence[str]] = None, *, highlight_cols: bool = True) -> str:
        """
        Pretty-print the effective plan per explicit node (i.e., keys other than __all__).
        If `nodes` is None, all explicit nodes in the plan are included.
        """
        import re
        # choose nodes to show: explicit ones, not '__all__'
        if nodes is None:
            nodes = [k for k in self.plan.keys() if k != "__all__"]
            # if user only supplied __all__, show a synthetic target '*'
            if not nodes and "__all__" in self.plan:
                nodes = ["*"]  # represents "all nodes"

        def _hl_cols(expr: str) -> str:
            if not highlight_cols:
                return expr
            # Highlight df['col'] / df["col"]
            pattern = r"""df\[\s*(['"])(?P<col>.+?)\1\s*\]"""
            def repl(m):
                col = m.group("col")
                return f"df['**{col}**']"
            return re.sub(pattern, repl, expr)
        # Build text
        lines: List[str] = []
        # Header
        if "__all__" in self.plan:
            ga = self.plan["__all__"]
            lines.append("QueryPlan (merged per node; __all__ applies to every node)")
            lines.append("")
            lines.append("  __all__:")
            for gi, (gtype, exprs) in enumerate(ga, start=1):
                lines.append(f"    - {gtype.upper()} group #{gi}:")
                for ei, e in enumerate(exprs, start=1):
                    lines.append(f"        [{ei}] { _hl_cols(e) }")
            lines.append("")
        else:
            lines.append("QueryPlan (no __all__ groups)")
            lines.append("")

        for node in nodes:
            lines.append(f"  node: {node}")
            groups = self._groups_for(node if node != "*" else "__all__")
            if not groups:
                lines.append("    (no groups)")
                continue
            for gi, (gtype, exprs) in enumerate(groups, start=1):
                lines.append(f"    - {gtype.upper()} group #{gi}:")
                for ei, e in enumerate(exprs, start=1):
                    lines.append(f"        [{ei}] { _hl_cols(e) }")
            lines.append("")
        return "\n".join(lines)


# ----------------------------
# Parsing / Coercion
# ----------------------------

def _coerce_value_to_groups(val: Union[str, Sequence, Dict]) -> List[Group]:
    """
    Normalize a node's YAML value into a list of Groups.
    - str               -> [("and", [str])]
    - list[str|dict]    -> flatten: str -> AND; dict -> and/or keys
    - dict              -> keys may be 'and' and/or 'or'
    """
    groups: List[Group] = []

    def add_and(exprs: Sequence[str]):
        exprs = [str(e) for e in exprs if str(e).strip()]
        if exprs:
            groups.append(("and", list(exprs)))

    def add_or(exprs: Sequence[str]):
        exprs = [str(e) for e in exprs if str(e).strip()]
        if exprs:
            groups.append(("or", list(exprs)))

    if isinstance(val, str):
        add_and([val])
        return groups

    if isinstance(val, list):
        # Each element can be a string (AND) or a dict with 'and'/'or'
        and_bucket: List[str] = []
        for item in val:
            if isinstance(item, str):
                and_bucket.append(item)
            elif isinstance(item, dict):
                # flush any pending ANDs before appending a new group
                if and_bucket:
                    add_and(and_bucket)
                    and_bucket = []
                if "and" in item:
                    add_and(item.get("and", []))
                if "or" in item:
                    add_or(item.get("or", []))
            else:
                raise TypeError(f"Invalid list item type in query: {type(item).__name__}")
        if and_bucket:
            add_and(and_bucket)
        return groups

    if isinstance(val, dict):
        # explicit 'and' and/or 'or'
        if "and" in val:
            and_val = val.get("and", [])
            if isinstance(and_val, str):
                add_and([and_val])
            else:
                add_and(and_val)
        if "or" in val:
            or_val = val.get("or", [])
            if isinstance(or_val, str):
                add_or([or_val])
            else:
                add_or(or_val)
        # if no and/or keys, treat other keys as error
        if not groups and val:
            raise ValueError(f"Expected 'and' and/or 'or' keys, got: {list(val.keys())}")
        return groups

    raise TypeError(f"Invalid node value type in query: {type(val).__name__}")
