RULE/utils.py

# utils.py


import ast
import csv
import datetime
import io
import json
import math
import os
import re
import tempfile
import time
import uuid
from contextlib import redirect_stderr, redirect_stdout
from typing import Any, Dict, List, Optional, Tuple

import numpy as np
import pandas as pd
import requests
# Mongo
from pymongo import ASCENDING, DESCENDING, MongoClient

# Core rule-engine helpers (from your pandas_rule.py)
from RULE.pandas_rule import (  # NL operation helpers (used in Part 2); LLM-based routing functions
    classify_query_intent, generate_aggregate_code, generate_transform_code,
    load_file, nl_to_condition, parse_column_ops, parse_concat_tables,
    parse_create_column, parse_lambda_apply, parse_nl_operation,
    parse_transform_column, safe_apply)

# ---------------- Constants and Globals ----------------

APP_VERSION = "build-2025-09-25-hf-storage-safe-v1"

# In-memory datasets (used temporarily for workflow execution, then cleared)
STORED_DFS: Dict[str, pd.DataFrame] = {}
STORED_META: Dict[str, Dict[str, Any]] = {}  # in-memory file metadata

# ===== Mongo Configuration =====
MONGODB_URI = os.environ.get("MONGODB_URI")
# Support both MONGODB_DB and MONGODB_DATABASE for compatibility
MONGODB_DB = os.environ.get("MONGODB_DB", "rulzAI")

# Collections (as requested)
RULES_COLL_NAME = os.environ.get("RULES_COLL_NAME", "Rules")
DATASET_COLL_NAME = os.environ.get("DATASET_COLL_NAME", "dataset")
FILTERED_COLL_NAME = os.environ.get("FILTERED_COLL_NAME", "filtered_dataset")
WORKFLOW_COLL_NAME = os.environ.get("WORKFLOW_COLL_NAME", "RulzAI workflow")
WORKFLOW_EXECUTE_COLL_NAME = os.environ.get("WORKFLOW_EXECUTE_COLL_NAME", "workflow_execute")

MONGO_CLIENT = None
MONGO_DB = None
MONGO_RULES = None
MONGO_DATASET = None
MONGO_FILTERED = None
MONGO_WORKFLOWS = None
MONGO_WORKFLOW_EXECUTE = None

try:
    if not MONGODB_URI:
        raise ValueError(
            "MONGODB_URI environment variable is not set. "
            "Please configure MongoDB connection in your .env file. "
            "For cloud: MONGODB_URI=mongodb+srv://user:pass@host/?appName=app "
            "For local: MONGODB_URI=mongodb://localhost:27017/"
        )
    
    MONGO_CLIENT = MongoClient(
        MONGODB_URI,
        serverSelectionTimeoutMS=6000,
        connectTimeoutMS=6000,
        retryWrites=True,
        w="majority",
    )
    MONGO_CLIENT.admin.command("ping")
    MONGO_DB = MONGO_CLIENT[MONGODB_DB]
    MONGO_RULES = MONGO_DB[RULES_COLL_NAME]
    MONGO_DATASET = MONGO_DB[DATASET_COLL_NAME]
    MONGO_FILTERED = MONGO_DB[FILTERED_COLL_NAME]
    MONGO_WORKFLOWS = MONGO_DB[WORKFLOW_COLL_NAME]
    MONGO_WORKFLOW_EXECUTE = MONGO_DB[WORKFLOW_EXECUTE_COLL_NAME]

    # Indexes
    # Rules: unique name per type=rule
    MONGO_RULES.create_index(
        [("type", ASCENDING), ("name", ASCENDING)],
        unique=True,
        name="uniq_rule_name",
        partialFilterExpression={"type": "rule", "name": {"$type": "string"}}
    )
    # Timestamps and run listing
    MONGO_RULES.create_index([("type", ASCENDING), ("created_at", DESCENDING)], name="rules_created_at")
    MONGO_RULES.create_index([("type", ASCENDING), ("last_run_at", DESCENDING)], name="rules_last_run_at")
    MONGO_RULES.create_index([("type", ASCENDING), ("rule_id", ASCENDING), ("ts", DESCENDING)], name="runs_rule_ts")
    MONGO_RULES.create_index([("type", ASCENDING), ("status", ASCENDING), ("ts", DESCENDING)], name="runs_status_ts")

    # Dataset: unique name for meta docs
    MONGO_DATASET.create_index(
        [("type", ASCENDING), ("name", ASCENDING)],
        unique=True,
        name="uniq_dataset_name",
        partialFilterExpression={"type": "dataset_meta", "name": {"$type": "string"}}
    )
    MONGO_DATASET.create_index([("type", ASCENDING), ("created_at", DESCENDING)], name="dataset_created_at")
    MONGO_DATASET.create_index(
        [("type", ASCENDING), ("dataset_id", ASCENDING), ("chunk_index", ASCENDING)],
        name="dataset_chunks_order",
        partialFilterExpression={"type": "dataset_chunk"}
    )

    # Filtered: unique name for meta docs
    MONGO_FILTERED.create_index(
        [("type", ASCENDING), ("name", ASCENDING)],
        unique=True,
        name="uniq_filtered_name",
        partialFilterExpression={"type": "filtered_meta", "name": {"$type": "string"}}
    )
    MONGO_FILTERED.create_index([("type", ASCENDING), ("created_at", DESCENDING)], name="filtered_created_at")
    MONGO_FILTERED.create_index(
        [("type", ASCENDING), ("filtered_id", ASCENDING), ("chunk_index", ASCENDING)],
        name="filtered_chunks_order",
        partialFilterExpression={"type": "filtered_chunk"}
    )

    # Workflow Execute indexes
    MONGO_WORKFLOW_EXECUTE.create_index(
        [("workflow_id", ASCENDING), ("created_at", DESCENDING)],
        name="workflow_exec_wf_created"
    )
    MONGO_WORKFLOW_EXECUTE.create_index(
        [("process_id", ASCENDING)],
        name="workflow_exec_pid"
    )
    MONGO_WORKFLOW_EXECUTE.create_index(
        [("status", ASCENDING), ("created_at", DESCENDING)],
        name="workflow_exec_status_created"
    )
    MONGO_WORKFLOW_EXECUTE.create_index(
        [("type", ASCENDING), ("execution_id", ASCENDING), ("result_type", ASCENDING), ("chunk_index", ASCENDING)],
        name="workflow_result_lookup"
    )

    print(f"MongoDB connected successfully to database '{MONGODB_DB}'.")
except ValueError as ve:
    print(f"MongoDB configuration error: {ve}")
    MONGO_DB = None
except Exception as e:
    print(f"MongoDB connection failed: {e}")
    MONGO_DB = None


# ===== Process timers (for durations) =====
PROCESS_TIMERS: Dict[str, float] = {}


# ---------------- Utilities ----------------

def now_utc_iso() -> str:
    return datetime.datetime.now(datetime.timezone.utc).isoformat()


def gen_id(prefix: str) -> str:
    return f"{prefix}_{uuid.uuid4().hex}"


def slugify(s: str, max_len: int = 64) -> str:
    s = re.sub(r'[^A-Za-z0-9._-]+', '_', str(s).strip())
    s = re.sub(r'_{2,}', '_', s).strip('_')
    return s[:max_len] if len(s) > max_len else s


def ensure_outcome_labels(df: pd.DataFrame) -> pd.DataFrame:
    if "Outcome_Labels" not in df.columns:
        df["Outcome_Labels"] = [[] for _ in range(len(df))]
        return df

    def to_list(x):
        if isinstance(x, list):
            return x
        if pd.isna(x):
            return []
        try:
            val = ast.literal_eval(str(x))
            if isinstance(val, list):
                return val
        except Exception:
            pass
        return [x]

    df["Outcome_Labels"] = df["Outcome_Labels"].apply(to_list)
    return df


def df_preview(df: pd.DataFrame, max_rows: int = 20):
    sub = df.head(max_rows).copy()
    sub = sub.replace([np.inf, -np.inf], np.nan)
    sub = sub.where(pd.notnull(sub), None)
    return sub.to_dict(orient="records")


def to_json_safe(obj: Any) -> Any:
    if obj is None:
        return None
    if isinstance(obj, (str, bool, int)):
        return obj
    if isinstance(obj, float):
        return None if (math.isnan(obj) or math.isinf(obj)) else obj
    if isinstance(obj, (np.integer,)):
        return int(obj)
    if isinstance(obj, (np.floating,)):
        val = float(obj)
        return None if (math.isnan(val) or math.isinf(val)) else val
    if isinstance(obj, dict):
        return {to_json_safe(k): to_json_safe(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple, set)):
        return [to_json_safe(x) for x in obj]
    if isinstance(obj, pd.DataFrame):
        return df_preview(obj)
    if isinstance(obj, pd.Series):
        return to_json_safe(obj.to_dict())
    return str(obj)


def capture_output(fn, *args, **kwargs):
    buf = io.StringIO()
    with redirect_stdout(buf), redirect_stderr(buf):
        result = fn(*args, **kwargs)
    return result, buf.getvalue()


def normalize_dataframe(df: pd.DataFrame) -> pd.DataFrame:
    df.columns = (
        df.columns.astype(str)
        .str.strip()
        .str.replace(r"[^0-9a-zA-Z]+", "_", regex=True)
        .str.strip("_")
    )
    return df


def coerce_numeric(df: pd.DataFrame) -> pd.DataFrame:
    for col in df.columns:
        if df[col].dtype == "object":
            conv = pd.to_numeric(df[col], errors="coerce")
            if len(conv) == 0:
                continue
            non_na_ratio = conv.notna().mean()
            if non_na_ratio == 1.0 or (non_na_ratio >= 0.8 and conv.notna().sum() > 0):
                df[col] = conv
    return df


def load_uploaded_to_df(upload: Any, max_mb: int = int(os.environ.get("MAX_UPLOAD_MB", "512"))) -> pd.DataFrame:
    """
    Storage-safe upload: stream to temp file with size guard, load, and delete temp immediately.
    Prevents /tmp growth and limits oversized uploads.
    'upload' can be FastAPI's UploadFile (recommended) or any object with file-like .file and .filename.
    """
    suffix = os.path.splitext(getattr(upload, "filename", "upload.csv"))[1].lower() or ".csv"
    fd, tmp_path = tempfile.mkstemp(suffix=suffix, prefix="upload_")
    max_bytes = max_mb * 1024 * 1024
    written = 0
    try:
        with os.fdopen(fd, "wb") as f:
            while True:
                chunk = upload.file.read(1024 * 1024)  # 1 MB
                if not chunk:
                    break
                written += len(chunk)
                if written > max_bytes:
                    raise ValueError(f"Upload too large: {written/1e6:.1f} MB > {max_mb} MB limit")
                f.write(chunk)
        try:
            upload.file.seek(0)
        except Exception:
            pass
        df = load_file(tmp_path)
    finally:
        try:
            os.remove(tmp_path)
        except Exception:
            pass
    return df


# ---------------- Process logging (stored in Rules collection) ----------------

def begin_process(endpoint: str, params: Optional[Dict[str, Any]] = None) -> str:
    pid = gen_id("proc")
    PROCESS_TIMERS[pid] = time.time()
    if MONGO_DB is not None:
        doc = {
            "_id": pid,
            "type": "process",
            "endpoint": endpoint,
            "params": params or {},
            "started_at": now_utc_iso(),
            "status": "running",
            "version": APP_VERSION,
        }
        MONGO_RULES.insert_one(doc)
    return pid


def end_process(process_id: str, status: str = "success", extra: Optional[Dict[str, Any]] = None):
    ended_at = now_utc_iso()
    started_ts = PROCESS_TIMERS.pop(process_id, None)
    duration_ms = int((time.time() - started_ts) * 1000) if started_ts else None
    if MONGO_DB is not None:
        update = {"$set": {"status": status, "ended_at": ended_at}}
        if duration_ms is not None:
            update["$set"]["duration_ms"] = duration_ms
        if extra:
            update["$set"]["result"] = extra
        MONGO_RULES.update_one({"_id": process_id, "type": "process"}, update)


# ---------------- Mongo helpers: dataset and filtered_dataset ----------------

DEFAULT_CHUNK_SIZE = int(os.environ.get("MONGO_CHUNK_SIZE", "5000"))

def _enforce_unique_name(col, type_value: str, name: str, overwrite: bool):
    if name is None:
        return
    existing = col.find_one({"type": type_value, "name": name}, {"_id": 1})
    if existing and not overwrite:
        raise ValueError(f"name '{name}' already exists in collection '{col.name}'. Use overwrite=true to replace.")
    if existing and overwrite:
        # Delete existing meta + chunks
        if type_value == "dataset_meta":
            ds_id = existing["_id"]
            col.delete_many({"dataset_id": ds_id})
        elif type_value == "filtered_meta":
            fd_id = existing["_id"]
            col.delete_many({"filtered_id": fd_id})
        col.delete_one({"_id": existing["_id"], "type": type_value})


def _df_to_chunks(df: pd.DataFrame, chunk_size: int) -> List[pd.DataFrame]:
    if chunk_size <= 0:
        chunk_size = DEFAULT_CHUNK_SIZE
    n = len(df)
    if n == 0:
        return [df]
    return [df.iloc[i:i + chunk_size].copy() for i in range(0, n, chunk_size)]


def save_dataset_to_mongo(
    df: pd.DataFrame,
    name: str,
    *,
    overwrite: bool = False,
    chunk_size: int = DEFAULT_CHUNK_SIZE
) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    _enforce_unique_name(MONGO_DATASET, "dataset_meta", name, overwrite)
    ds_id = gen_id("ds")
    created_at = now_utc_iso()
    df = df.copy()
    df = coerce_numeric(ensure_outcome_labels(df))

    chunks = _df_to_chunks(df, chunk_size)
    for idx, c in enumerate(chunks):
        MONGO_DATASET.insert_one({
            "_id": f"{ds_id}:{idx}",
            "type": "dataset_chunk",
            "dataset_id": ds_id,
            "chunk_index": idx,
            "rows": int(len(c)),
            "data": c.replace([np.inf, -np.inf], np.nan).where(pd.notnull(c), None).to_dict(orient="records"),
            "created_at": created_at,
        })

    meta = {
        "_id": ds_id,
        "type": "dataset_meta",
        "name": name,
        "rows": int(len(df)),
        "columns": list(map(str, df.columns)),
        "chunks": len(chunks),
        "created_at": created_at,
    }
    MONGO_DATASET.insert_one(meta)
    return meta


def load_dataset_df_from_mongo(*, id: Optional[str] = None, name: Optional[str] = None) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    if not id and not name:
        raise ValueError("Provide id or name")
    meta = MONGO_DATASET.find_one(
        {"type": "dataset_meta", **({"_id": id} if id else {}), **({"name": name} if name else {})}
    )
    if not meta:
        raise ValueError("Dataset not found")
    chunks = list(MONGO_DATASET.find({"type": "dataset_chunk", "dataset_id": meta["_id"]}).sort([("chunk_index", ASCENDING)]))
    rows = []
    for ch in chunks:
        rows.extend(ch.get("data", []))
    df = pd.DataFrame(rows) if rows else pd.DataFrame(columns=meta.get("columns", []))
    df = ensure_outcome_labels(df)
    return df, meta


def preview_dataset_from_mongo(*, id: Optional[str] = None, name: Optional[str] = None, max_rows: int = 20) -> Dict[str, Any]:
    df, meta = load_dataset_df_from_mongo(id=id, name=name)
    return {
        "dataset": {"_id": meta["_id"], "name": meta["name"]},
        "rows": meta["rows"],
        "preview": df_preview(df, max_rows=max_rows)
    }


def delete_dataset_from_mongo(*, id: Optional[str] = None, name: Optional[str] = None) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    if not id and not name:
        raise ValueError("Provide id or name")
    meta = MONGO_DATASET.find_one({"type": "dataset_meta", **({"_id": id} if id else {}), **({"name": name} if name else {})})
    if not meta:
        raise ValueError("Dataset not found")
    ds_id = meta["_id"]
    ch_res = MONGO_DATASET.delete_many({"type": "dataset_chunk", "dataset_id": ds_id})
    meta_res = MONGO_DATASET.delete_one({"_id": ds_id, "type": "dataset_meta"})
    return {"deleted_docs": int(ch_res.deleted_count + meta_res.deleted_count), "dataset_id": ds_id}


def save_filtered_to_mongo(
    df: pd.DataFrame,
    name: str,
    *,
    outcome: str,
    matched_rows: int,
    overwrite: bool = False,
    chunk_size: int = DEFAULT_CHUNK_SIZE
) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    _enforce_unique_name(MONGO_FILTERED, "filtered_meta", name, overwrite)
    fd_id = gen_id("fd")
    created_at = now_utc_iso()
    df = df.copy()
    df = coerce_numeric(ensure_outcome_labels(df))

    chunks = _df_to_chunks(df, chunk_size)
    for idx, c in enumerate(chunks):
        MONGO_FILTERED.insert_one({
            "_id": f"{fd_id}:{idx}",
            "type": "filtered_chunk",
            "filtered_id": fd_id,
            "chunk_index": idx,
            "rows": int(len(c)),
            "data": c.replace([np.inf, -np.inf], np.nan).where(pd.notnull(c), None).to_dict(orient="records"),
            "created_at": created_at,
        })

    meta = {
        "_id": fd_id,
        "type": "filtered_meta",
        "name": name,
        "outcome": outcome,
        "rows": int(len(df)),
        "matched_rows": int(matched_rows),
        "columns": list(map(str, df.columns)),
        "chunks": len(chunks),
        "created_at": created_at,
    }
    MONGO_FILTERED.insert_one(meta)
    return meta


def load_filtered_df_from_mongo(*, id: Optional[str] = None, name: Optional[str] = None) -> Tuple[pd.DataFrame, Dict[str, Any]]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    if not id and not name:
        raise ValueError("Provide id or name")
    meta = MONGO_FILTERED.find_one({"type": "filtered_meta", **({"_id": id} if id else {}), **({"name": name} if name else {})})
    if not meta:
        raise ValueError("Filtered dataset not found")
    chunks = list(MONGO_FILTERED.find({"type": "filtered_chunk", "filtered_id": meta["_id"]}).sort([("chunk_index", ASCENDING)]))
    rows = []
    for ch in chunks:
        rows.extend(ch.get("data", []))
    df = pd.DataFrame(rows) if rows else pd.DataFrame(columns=meta.get("columns", []))
    df = ensure_outcome_labels(df)
    return df, meta


def persist_df_to_mongo(
    df: pd.DataFrame,
    *,
    name: Optional[str],
    persist_to: str,  # "dataset" | "filtered"  (pipeline will always use "filtered")
    overwrite: bool = False,
    outcome: Optional[str] = None,
    matched_rows: Optional[int] = None,
    chunk_size: int = DEFAULT_CHUNK_SIZE
) -> Dict[str, Any]:
    """
    Persists df to Mongo. For pipeline outputs, we always use persist_to='filtered'.
    'dataset' path remains available for raw uploads or other flows outside the pipeline.
    """
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")

    name = name or f"{persist_to}_{uuid.uuid4().hex[:8]}"
    name = slugify(name)

    if persist_to == "dataset":
        meta = save_dataset_to_mongo(df, name, overwrite=overwrite, chunk_size=chunk_size)
        result = {
            "type": "dataset",
            "id": meta["_id"],
            "name": meta["name"],
            "rows": meta["rows"],
            "chunks": meta["chunks"],
            "created_at": meta["created_at"]
        }
    elif persist_to == "filtered":
        if outcome is None or matched_rows is None:
            raise ValueError("For persist_to='filtered', provide outcome and matched_rows.")
        meta = save_filtered_to_mongo(
            df, name,
            outcome=outcome,
            matched_rows=matched_rows,
            overwrite=overwrite,
            chunk_size=chunk_size
        )
        result = {
            "type": "filtered",
            "id": meta["_id"],
            "name": meta["name"],
            "outcome": meta["outcome"],
            "matched_rows": meta["matched_rows"],
            "rows": meta["rows"],
            "chunks": meta["chunks"],
            "created_at": meta["created_at"]
        }
    else:
        raise ValueError("persist_to must be 'dataset' or 'filtered'")

    # ── Emit db.write event ───────────────────────────────────
    try:
        from triggers.event_bus import EventBus
        from triggers.trigger_models import EventType, PlatformEvent
        EventBus.get().emit_sync(PlatformEvent(
            event_type=EventType.DB_WRITE,
            source="persist_df_to_mongo",
            payload={
                "persist_to": persist_to,
                "name": name,
                "rows": result.get("rows"),
                "id": result.get("id"),
            },
        ))
    except Exception:
        pass  # trigger failures must never block DB writes

    return result


# ---------------- Rules helpers ----------------

def save_temp_rule_doc(
    *,
    nl_query: str,
    outcome: str,
    description: Optional[str],
    condition: Optional[str] = None
) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    tmp_id = gen_id("tmp_rule")
    now = now_utc_iso()
    doc = {
        "_id": tmp_id,
        "type": "temp_rule",
        "outcome": outcome,
        "nl_query": nl_query,
        "condition": condition,        # may be None if compile skipped
        "description": description,
        "created_at": now,
        "committed": False
    }
    MONGO_RULES.insert_one(doc)
    return doc


def fetch_temp_rule(temp_rule_id: str) -> Optional[Dict[str, Any]]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    return MONGO_RULES.find_one({"_id": temp_rule_id, "type": "temp_rule"})


def save_rule_doc(
    *,
    name: str,
    outcome: str,
    condition: Optional[str],
    nl_query: Optional[str],
    description: Optional[str],
    status: str = "active",
    overwrite: bool = False,
    extra: Optional[Dict[str, Any]] = None   # allow bundle fields
) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")

    existing = MONGO_RULES.find_one({"type": "rule", "name": name}, {"_id": 1})
    if existing and not overwrite:
        raise ValueError(f"Rule name '{name}' already exists. Use overwrite=true to replace.")
    if existing and overwrite:
        MONGO_RULES.delete_one({"_id": existing["_id"], "type": "rule"})

    rule_id = gen_id("rule")
    now = now_utc_iso()
    doc = {
        "_id": rule_id,
        "type": "rule",
        "name": name,
        "outcome": outcome,
        "condition": condition,
        "nl_query": nl_query,
        "description": description,
        "status": status,
        "counters": {
            "run_count": 0,
            "success_count": 0,
            "zero_match_count": 0,
            "failure_count": 0
        },
        "created_at": now,
        "last_run_at": None
    }
    if extra:
        doc.update(extra)
    MONGO_RULES.insert_one(doc)
    return doc


def fetch_rule(id: Optional[str] = None, name: Optional[str] = None) -> Optional[Dict[str, Any]]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    if not id and not name:
        raise ValueError("Provide id or name")
    return MONGO_RULES.find_one(
        {"type": "rule", **({"_id": id} if id else {}), **({"name": name} if name else {})}
    )


# ---- Temp bundle + validation helpers ----

def save_temp_bundle_doc(
    *,
    name: Optional[str],
    description: Optional[str],
    rules: List[Dict[str, Any]],
    file_info: Optional[Dict[str, Any]] = None,
    validation: Optional[Dict[str, Any]] = None
) -> Dict[str, Any]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    tmp_id = gen_id("tmp_bundle")
    now = now_utc_iso()
    doc = {
        "_id": tmp_id,
        "type": "temp_bundle",
        "name": name,
        "description": description,
        "rules": rules,
        "bundle_count": int(len(rules)),
        "file": file_info or {},
        "validation": validation or {},
        "created_at": now,
        "committed": False
    }
    MONGO_RULES.insert_one(doc)
    return doc


def fetch_temp_bundle(temp_bundle_id: str) -> Optional[Dict[str, Any]]:
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    return MONGO_RULES.find_one({"_id": temp_bundle_id, "type": "temp_bundle"})


def compute_validation_for_rules(rules: List[Dict[str, Any]]) -> Dict[str, Any]:
    """
    Static validation of rules without needing a dataset.
    Criteria:
      - outcome present
      - has at least one of condition or nl_query
      - duplicate names flagged
      - nl_query compilation attempt (with empty df) tracked
      - mild penalties for poor names (too long / unsafe)
    Scoring (heuristic, 0-100):
      - Start 100
      - -10 for each missing outcome
      - -15 for each missing both condition and nl_query
      - -5 for each duplicate name occurrence (beyond first)
      - -5 for each compile error (when nl_query provided)
      - -2 if name > 64 chars
      - -1 if name contains unsafe chars (differs from slugified)
    """
    issues: List[str] = []
    per_rule: List[Dict[str, Any]] = []
    score = 100

    # Duplicate name detection
    names = [str(r.get("name") or "").strip() for r in rules]
    name_counts: Dict[str, int] = {}
    for n in names:
        if n:
            name_counts[n] = name_counts.get(n, 0) + 1
    dup_names = {n: c for n, c in name_counts.items() if c > 1}
    for n, c in dup_names.items():
        occ_penalty = (c - 1) * 5
        score -= occ_penalty
        issues.append(f"Duplicate rule name '{n}' occurs {c} times (-{occ_penalty})")

    empty_df = pd.DataFrame()

    for idx, r in enumerate(rules):
        nm = str(r.get("name") or f"rule_{idx+1}")
        out = r.get("outcome")
        cond = r.get("condition")
        nlq = r.get("nl_query") or r.get("query")

        has_outcome = bool(out and str(out).strip())
        has_logic = bool((cond and str(cond).strip()) or (nlq and str(nlq).strip()))
        pr = {
            "index": idx,
            "name": nm,
            "has_outcome": has_outcome,
            "has_condition": bool(cond and str(cond).strip()),
            "has_nl_query": bool(nlq and str(nlq).strip()),
            "compile_ok": None,
            "compile_error": None
        }

        if not has_outcome:
            score -= 10
            pr["error"] = (pr.get("error") or "") + "missing outcome; "
            issues.append(f"Rule[{idx}] '{nm}': missing outcome (-10)")

        if not has_logic:
            score -= 15
            pr["error"] = (pr.get("error") or "") + "missing condition and nl_query; "
            issues.append(f"Rule[{idx}] '{nm}': no condition or nl_query (-15)")

        # Minor name checks
        if len(nm) > 64:
            score -= 2
            issues.append(f"Rule[{idx}] '{nm[:24]}...': very long name (-2)")
        safe_nm = slugify(nm)
        if safe_nm != nm:
            score -= 1
            issues.append(f"Rule[{idx}] '{nm}': unsafe chars in name (-1); suggested '{safe_nm}'")

        # Try compiling nl_query if present (non-fatal)
        if nlq and str(nlq).strip():
            try:
                _cond, cap = capture_output(nl_to_condition, nlq, empty_df)
                pr["compile_ok"] = True
            except Exception as ce:
                score -= 5
                pr["compile_ok"] = False
                pr["compile_error"] = str(ce)
                issues.append(f"Rule[{idx}] '{nm}': nl_query compile error (-5): {ce}")

        per_rule.append(pr)

    score = max(0, min(100, score))
    summary = {
        "total_rules": len(rules),
        "score": int(score),
        "issues": issues,
        "per_rule": per_rule,
        "duplicates": dup_names
    }
    return summary


def inc_rule_stats(rule_id: str, status: str):
    if MONGO_DB is None:
        return
    inc = {"counters.run_count": 1}
    if status == "success":
        inc["counters.success_count"] = 1
    elif status == "zero_match":
        inc["counters.zero_match_count"] = 1
    elif status == "failure":
        inc["counters.failure_count"] = 1
    MONGO_RULES.update_one(
        {"_id": rule_id, "type": "rule"},
        {"$inc": inc, "$set": {"last_run_at": now_utc_iso()}}
    )


def record_rule_run(
    *,
    rule_id: str,
    rule_name: str,
    status: str,
    matched_rows: int,
    total_rows: int,
    data_source: str,
    dataset_ref: Optional[Dict[str, Any]],
    process_id: str,
    extra_fields: Optional[Dict[str, Any]] = None
) -> str:
    if MONGO_DB is None:
        return ""
    run_id = gen_id("run")
    doc = {
        "_id": run_id,
        "type": "rule_run",
        "rule_id": rule_id,
        "rule_name": rule_name,
        "status": status,  # success | zero_match | failure
        "matched_rows": int(matched_rows),
        "total_rows": int(total_rows),
        "data_source": data_source,
        "dataset_ref": dataset_ref or {},
        "ts": now_utc_iso(),
        "process_id": process_id
    }
    if extra_fields:
        doc.update(extra_fields)
    MONGO_RULES.insert_one(doc)

    # ── Emit rule.executed event ───────────────────────────────────────
    try:
        from triggers.event_bus import EventBus
        from triggers.trigger_models import EventType, PlatformEvent
        EventBus.get().emit_sync(PlatformEvent(
            event_type=EventType.RULE_EXECUTED,
            source="record_rule_run",
            correlation_id=process_id,
            payload={
                "rule_id": rule_id,
                "rule_name": rule_name,
                "status": status,
                "matched_rows": matched_rows,
                "total_rows": total_rows,
                "data_source": data_source,
                "run_id": run_id,
            },
        ))
    except Exception:
        pass  # trigger failures must never block rule execution

    return run_id


# ---------------- Data source loader for execution ----------------

_DB_CONNECTOR_BASE = "https://api.prediqai.com/db-connector"

def _fetch_from_db_connector(instance_id: str, collection_name: str) -> List[Dict[str, Any]]:
    """
    Fetch ALL records from a DB Connector collection (unlimited=true).
    Returns a list of row dicts ready to be loaded into a DataFrame.
    Raises RuntimeError on any HTTP / parsing failure.
    """
    url = f"{_DB_CONNECTOR_BASE}/instances/{instance_id}/resources/{collection_name}/data"
    payload = {"unlimited": True}
    try:
        resp = requests.post(url, json=payload, timeout=120)
        resp.raise_for_status()
    except requests.exceptions.HTTPError as exc:
        raise RuntimeError(
            f"DB Connector returned HTTP {exc.response.status_code} for "
            f"{instance_id}/{collection_name}: {exc.response.text[:300]}"
        ) from exc
    except Exception as exc:
        raise RuntimeError(
            f"DB Connector request failed for {instance_id}/{collection_name}: {exc}"
        ) from exc

    body = resp.json()
    if not body.get("success"):
        raise RuntimeError(
            f"DB Connector error for {instance_id}/{collection_name}: {body}"
        )
    return body.get("data") or []

def get_df_from_source(
    *,
    data_source: str,
    data_name: Optional[str],
    data_file: Optional[Any],
    normalize_cols: bool = False
) -> Tuple[pd.DataFrame, str, Optional[Dict[str, Any]]]:
    """
    Returns (df, source_descriptor, dataset_ref_dict)
    dataset_ref_dict can be {"dataset_id": "...", "dataset_name": "..."} or {"filtered_id": "...", "filtered_name": "..."}
    """
    src = data_source.lower().strip()
    if src == "upload":
        if data_file is None:
            raise ValueError("data file required for data_source=upload")
        df = load_uploaded_to_df(data_file)
        if normalize_cols:
            df = normalize_dataframe(df)
        df = coerce_numeric(ensure_outcome_labels(df))
        return df, f"upload:{getattr(data_file, 'filename', 'upload')}", None

    if src == "dataset":
        if not data_name:
            raise ValueError("data_name (dataset id or name) required for data_source=dataset")
        df, meta = load_dataset_df_from_mongo(id=data_name if data_name.startswith("ds_") else None,
                                              name=None if data_name.startswith("ds_") else data_name)
        if normalize_cols:
            df = normalize_dataframe(df)
        df = coerce_numeric(ensure_outcome_labels(df))
        return df, f"dataset:{meta['name']}", {"dataset_id": meta["_id"], "dataset_name": meta["name"]}

    if src == "filtered":
        if not data_name:
            raise ValueError("data_name (filtered id or name) required for data_source=filtered")
        df, meta = load_filtered_df_from_mongo(id=data_name if data_name.startswith("fd_") else None,
                                               name=None if data_name.startswith("fd_") else data_name)
        if normalize_cols:
            df = normalize_dataframe(df)
        df = coerce_numeric(ensure_outcome_labels(df))
        return df, f"filtered:{meta['name']}", {"filtered_id": meta["_id"], "filtered_name": meta["name"]}

    if src == "db_connector":
        if not data_name:
            raise ValueError(
                "data_name must be '<instance_id>/<collection_name>' "
                "for data_source=db_connector"
            )
        parts = data_name.strip().split("/", 1)
        if len(parts) != 2 or not parts[0].strip() or not parts[1].strip():
            raise ValueError(
                "data_name must be '<instance_id>/<collection_name>' "
                f"(e.g. '904342c0-6c7c-48f3-a6f2-b86fe5dcd6b3/movies'), got: {data_name!r}"
            )
        instance_id, collection_name = parts[0].strip(), parts[1].strip()
        records = _fetch_from_db_connector(instance_id, collection_name)
        df = pd.DataFrame(records) if records else pd.DataFrame()
        if normalize_cols:
            df = normalize_dataframe(df)
        df = coerce_numeric(ensure_outcome_labels(df))
        return (
            df,
            f"db_connector:{instance_id}/{collection_name}",
            {"instance_id": instance_id, "collection": collection_name},
        )

    # Memory source intentionally disabled (avoid RAM usage)
    raise ValueError("data_source must be one of: upload|dataset|filtered|db_connector")


# ============ Workflow Execution Helpers ============

def record_in_memory_file(stored_name: str, df: pd.DataFrame, *, process_id: Optional[str], action: str):
    """
    Store DataFrame in memory temporarily with metadata.
    Used ONLY for intermediate workflow results. 
    Memory is cleared after workflow completes and saves to MongoDB.
    """
    STORED_DFS[stored_name] = df.copy()
    STORED_META[stored_name] = {
        "name": stored_name,
        "rows": len(df),
        "columns": list(df.columns),
        "process_id": process_id,
        "action": action,
        "stored_at": now_utc_iso(),
        "storage": "memory_temporary"
    }


def save_workflow_execution_to_mongo(
    *,
    workflow_id: str,
    workflow_name: str,
    process_id: str,
    source: str,
    dataset_ref: Optional[Dict[str, Any]],
    sequence: List[Dict[str, Any]],
    execution_results: List[Dict[str, Any]],
    initial_rows: int,
    final_rows: int,
    final_df: pd.DataFrame,
    intermediate_results: Dict[str, Dict[str, Any]],
    stored_final_as: Optional[str] = None,
    status: str = "success",
    chunk_size: int = DEFAULT_CHUNK_SIZE
) -> Dict[str, Any]:
    """
    Save complete workflow execution to MongoDB workflow_execute collection.
    Automatically clears in-memory data after saving.
    Returns metadata about saved execution.
    """
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    
    exec_id = gen_id("wf_exec")
    created_at = now_utc_iso()
    
    # Save final result chunks if available
    final_chunks_saved = 0
    if final_df is not None and len(final_df) > 0:
        final_df = coerce_numeric(ensure_outcome_labels(final_df))
        chunks = _df_to_chunks(final_df, chunk_size)
        
        for idx, chunk_df in enumerate(chunks):
            MONGO_WORKFLOW_EXECUTE.insert_one({
                "_id": f"{exec_id}:final:{idx}",
                "type": "workflow_result_chunk",
                "execution_id": exec_id,
                "workflow_id": workflow_id,
                "result_type": "final",
                "chunk_index": idx,
                "rows": len(chunk_df),
                "data": chunk_df.replace([np.inf, -np.inf], np.nan).where(pd.notnull(chunk_df), None).to_dict(orient="records"),
                "created_at": created_at
            })
            final_chunks_saved += 1
    
    # Save intermediate results from memory
    intermediate_chunks_saved = {}
    for store_key, meta in intermediate_results.items():
        if store_key in STORED_DFS:
            inter_df = STORED_DFS[store_key]
            inter_df = coerce_numeric(ensure_outcome_labels(inter_df))
            chunks = _df_to_chunks(inter_df, chunk_size)
            
            chunks_count = 0
            for idx, chunk_df in enumerate(chunks):
                MONGO_WORKFLOW_EXECUTE.insert_one({
                    "_id": f"{exec_id}:inter:{slugify(store_key)}:{idx}",
                    "type": "workflow_result_chunk",
                    "execution_id": exec_id,
                    "workflow_id": workflow_id,
                    "result_type": "intermediate",
                    "result_name": store_key,
                    "step": meta.get("step"),
                    "operation_type": meta.get("type"),
                    "operation_name": meta.get("name"),
                    "chunk_index": idx,
                    "rows": len(chunk_df),
                    "data": chunk_df.replace([np.inf, -np.inf], np.nan).where(pd.notnull(chunk_df), None).to_dict(orient="records"),
                    "created_at": created_at
                })
                chunks_count += 1
            
            intermediate_chunks_saved[store_key] = {
                "chunks": chunks_count,
                "rows": len(inter_df),
                "columns": list(inter_df.columns),
                **meta
            }
    
    # Create execution metadata document
    exec_meta = {
        "_id": exec_id,
        "type": "workflow_execution",
        "workflow_id": workflow_id,
        "workflow_name": workflow_name,
        "process_id": process_id,
        "status": status,
        "source": source,
        "dataset_ref": dataset_ref or {},
        "sequence": sequence,
        "execution_results": execution_results,
        "initial_rows": initial_rows,
        "final_rows": final_rows,
        "final_chunks": final_chunks_saved,
        "final_columns": list(final_df.columns) if final_df is not None else [],
        "stored_final_as": stored_final_as,
        "intermediate_results": intermediate_chunks_saved,
        "successful_steps": len([r for r in execution_results if r.get("status") in ["success", "skipped"]]),
        "failed_steps": len([r for r in execution_results if r.get("status") == "error"]),
        "total_steps": len(sequence),
        "created_at": created_at
    }
    
    MONGO_WORKFLOW_EXECUTE.insert_one(exec_meta)
    
    # IMPORTANT: Clear in-memory data after saving to MongoDB to prevent RAM buildup
    for store_key in list(intermediate_results.keys()):
        STORED_DFS.pop(store_key, None)
        STORED_META.pop(store_key, None)
    
    result = {
        "execution_id": exec_id,
        "final_chunks": final_chunks_saved,
        "intermediate_results_saved": len(intermediate_chunks_saved),
        "created_at": created_at,
        "memory_cleared": True
    }

    # ── Emit workflow.completed event ──────────────────────────────────
    try:
        from triggers.event_bus import EventBus
        from triggers.trigger_models import EventType, PlatformEvent
        EventBus.get().emit_sync(PlatformEvent(
            event_type=EventType.WORKFLOW_COMPLETED,
            source="save_workflow_execution_to_mongo",
            correlation_id=process_id,
            payload={
                "workflow_id": workflow_id,
                "workflow_name": workflow_name,
                "status": status,
                "initial_rows": initial_rows,
                "final_rows": final_rows,
                "execution_id": exec_id,
            },
        ))
    except Exception:
        pass  # trigger failures must never block workflow execution

    return result


def load_workflow_execution_from_mongo(execution_id: str) -> Dict[str, Any]:
    """
    Load workflow execution metadata and data from MongoDB.
    """
    if MONGO_DB is None:
        raise RuntimeError("MongoDB not connected")
    
    meta = MONGO_WORKFLOW_EXECUTE.find_one({"_id": execution_id, "type": "workflow_execution"})
    if not meta:
        raise ValueError(f"Workflow execution '{execution_id}' not found")
    
    # Load final result
    final_chunks = list(MONGO_WORKFLOW_EXECUTE.find(
        {"type": "workflow_result_chunk", "execution_id": execution_id, "result_type": "final"}
    ).sort([("chunk_index", ASCENDING)]))
    
    final_rows = []
    for chunk in final_chunks:
        final_rows.extend(chunk.get("data", []))
    
    final_df = None
    if final_rows:
        final_df = pd.DataFrame(final_rows)
        final_df = ensure_outcome_labels(final_df)
    
    # Load intermediate results
    intermediate_dfs = {}
    for result_name, result_meta in meta.get("intermediate_results", {}).items():
        inter_chunks = list(MONGO_WORKFLOW_EXECUTE.find(
            {"type": "workflow_result_chunk", "execution_id": execution_id, "result_type": "intermediate", "result_name": result_name}
        ).sort([("chunk_index", ASCENDING)]))
        
        inter_rows = []
        for chunk in inter_chunks:
            inter_rows.extend(chunk.get("data", []))
        
        if inter_rows:
            inter_df = pd.DataFrame(inter_rows)
            inter_df = ensure_outcome_labels(inter_df)
            intermediate_dfs[result_name] = inter_df
    
    return {
        "metadata": meta,
        "final_df": final_df,
        "intermediate_dfs": intermediate_dfs
    }


def list_workflow_executions(
    workflow_id: Optional[str] = None,
    workflow_name: Optional[str] = None,
    status: Optional[str] = None,
    limit: int = 50,
    skip: int = 0
) -> List[Dict[str, Any]]:
    """
    List workflow executions with optional filters.
    """
    if MONGO_DB is None:
        return []
    
    query = {"type": "workflow_execution"}
    if workflow_id:
        query["workflow_id"] = workflow_id
    if workflow_name:
        query["workflow_name"] = workflow_name
    if status:
        query["status"] = status
    
    cursor = MONGO_WORKFLOW_EXECUTE.find(query).sort([("created_at", DESCENDING)]).skip(skip).limit(limit)
    return list(cursor)


# =========================
# Dev utilities, NL pipeline (memoryless and filtered-only persistence)
# =========================

def get_df_from_input(
    data: Optional[Any],
    source_name: Optional[str],
    *,
    normalize_cols: bool = False,
    ensure_labels: bool = True,
    coerce_nums: bool = True
) -> Tuple[pd.DataFrame, str]:
    """
    Memoryless helper for dev: only accepts an uploaded file.
    Returns a DataFrame and a source descriptor ("upload:<name>").
    """
    if source_name:
        # Disable in-memory source usage
        raise ValueError("In-memory sources are disabled. Use get_df_from_source with data_source=dataset|filtered.")
    if data is None:
        raise ValueError("Provide a 'data' file upload.")
    df = load_uploaded_to_df(data)
    src = f"upload:{getattr(data, 'filename', 'upload')}"
    if normalize_cols:
        df = normalize_dataframe(df)
    if coerce_nums:
        df = coerce_numeric(df)
    if ensure_labels:
        df = ensure_outcome_labels(df)
    return df, src


def resolve_tables_spec(
    spec: Optional[Dict[str, Any]],
    current_df: pd.DataFrame
) -> Dict[str, pd.DataFrame]:
    """
    Build a dict of named DataFrames for NL ops like concat/join without using in-memory tables.
    Pass values as:
      - a DataFrame (already loaded), or
      - "dataset:<nameOrId>", or
      - "filtered:<nameOrId>"
    """
    tbls: Dict[str, pd.DataFrame] = {"current": current_df}
    if not spec:
        return tbls

    for name, val in spec.items():
        if isinstance(val, pd.DataFrame):
            tbls[name] = ensure_outcome_labels(val.copy())
            continue

        if isinstance(val, str):
            v = val.strip()
            if v.startswith("dataset:"):
                key = v.split(":", 1)[1].strip()
                df, _ = load_dataset_df_from_mongo(
                    id=key if key.startswith("ds_") else None,
                    name=None if key.startswith("ds_") else key
                )
                tbls[name] = ensure_outcome_labels(df)
                continue
            if v.startswith("filtered:"):
                key = v.split(":", 1)[1].strip()
                df, _ = load_filtered_df_from_mongo(
                    id=key if key.startswith("fd_") else None,
                    name=None if key.startswith("fd_") else key
                )
                tbls[name] = ensure_outcome_labels(df)
                continue

        raise ValueError(f"Unsupported table spec for '{name}': {val!r}. Use 'dataset:<nameOrId>' or 'filtered:<nameOrId>' or pass a DataFrame.")

    return tbls


def run_nl_pipeline(
    df: pd.DataFrame,
    nl_query: str,
    *,
    outcome: Optional[str] = None,
    filter_matched: bool = True,   # store matched subset by default for filter
    save_name: Optional[str] = None,  # replaces legacy store_as
    store_as: Optional[str] = None,   # kept for backward-compat name fallback
    overwrite: bool = False,
    chunk_size: int = DEFAULT_CHUNK_SIZE,
    tables: Optional[Dict[str, Any]] = None,  # may include "dataset:<nameOrId>" or "filtered:<nameOrId>"
    process_id: Optional[str] = None,
    drop_outcome_labels: bool = False  # if True, remove Outcome_Labels column from saved data and preview
) -> Dict[str, Any]:
    """
    Memoryless NL pipeline that ALWAYS persists DataFrame outputs to the 'filtered' collection.
    - Non-filter ops: outcome = provided outcome OR 'Derived:<action>', matched_rows = len(result_df)
    - Filter ops: matched_rows computed from the mask; outcome uses provided value or 'Derived:filter'
    - Aggregations (non-DataFrame) are returned but not persisted automatically.
    """
    logs: List[str] = []
    result_name = save_name or store_as  # prefer save_name if provided

    # Ensure df shape
    df = coerce_numeric(ensure_outcome_labels(df))

    # Resolve any external tables by loading from Mongo
    tbls = resolve_tables_spec(tables, df)

    def save_filtered_result(out_df: pd.DataFrame, action: str) -> Dict[str, Any]:
        # Use caller-provided outcome if any; otherwise tag with the action
        out_label = outcome or f"Derived:{action}"
        matched = len(out_df)
        
        # Drop Outcome_Labels column if requested (for /nl/universal API)
        df_to_save = out_df.copy()
        if drop_outcome_labels and "Outcome_Labels" in df_to_save.columns:
            df_to_save = df_to_save.drop(columns=["Outcome_Labels"])
        
        return persist_df_to_mongo(
            df_to_save,
            name=result_name or f"{action}_{uuid.uuid4().hex[:6]}",
            persist_to="filtered",
            overwrite=overwrite,
            outcome=out_label,
            matched_rows=matched,
            chunk_size=chunk_size
        )

    # 1) Lambda application
    res, cap = capture_output(parse_lambda_apply, nl_query, df)
    if res is not None:
        out_df = ensure_outcome_labels(res)
        logs.append("lambda apply:\n" + cap.strip())
        action = "lambda_apply"
        saved = save_filtered_result(out_df, action)
        
        # Drop Outcome_Labels from preview if requested
        preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
        
        return {
            "action": action,
            "preview": df_preview(preview_df),
            "rows": len(out_df),
            "logs": "\n".join([l for l in logs if l]),
            "saved": saved
        }

    # 2) Create new column
    res, cap = capture_output(parse_create_column, nl_query, df)
    if res is not None:
        out_df = ensure_outcome_labels(res)
        logs.append("create column:\n" + cap.strip())
        action = "create_column"
        saved = save_filtered_result(out_df, action)
        
        # Drop Outcome_Labels from preview if requested
        preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
        
        return {
            "action": action,
            "preview": df_preview(preview_df),
            "rows": len(out_df),
            "logs": "\n".join([l for l in logs if l]),
            "saved": saved
        }

    # 3) Concat/stack any two tables
    res, cap = capture_output(parse_concat_tables, nl_query, tbls)
    if res is not None:
        out_df = ensure_outcome_labels(res)
        logs.append("concat/stack:\n" + cap.strip())
        action = "concat_tables"
        saved = save_filtered_result(out_df, action)
        
        # Drop Outcome_Labels from preview if requested
        preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
        
        return {
            "action": action,
            "preview": df_preview(preview_df),
            "rows": len(out_df),
            "logs": "\n".join([l for l in logs if l]),
            "saved": saved
        }

    # 4) Rename/delete column
    res, cap = capture_output(parse_column_ops, nl_query, df)
    if res is not None:
        out_df = ensure_outcome_labels(res)
        logs.append("column ops:\n" + cap.strip())
        action = "column_ops"
        saved = save_filtered_result(out_df, action)
        
        # Drop Outcome_Labels from preview if requested
        preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
        
        return {
            "action": action,
            "preview": df_preview(preview_df),
            "rows": len(out_df),
            "logs": "\n".join([l for l in logs if l]),
            "saved": saved
        }

    # 5) Transform column values
    res, cap = capture_output(parse_transform_column, nl_query, df)
    if res is not None:
        out_df = ensure_outcome_labels(res)
        logs.append("transform column:\n" + cap.strip())
        action = "transform_column"
        saved = save_filtered_result(out_df, action)
        
        # Drop Outcome_Labels from preview if requested
        preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
        
        return {
            "action": action,
            "preview": df_preview(preview_df),
            "rows": len(out_df),
            "logs": "\n".join([l for l in logs if l]),
            "saved": saved
        }

    # 6) LLM-Based Query Classification and Routing
    classification, cap = capture_output(classify_query_intent, nl_query, df)
    logs.append(f"Query Classification:\n{cap.strip()}")
    logs.append(f"  Intent: {classification.get('intent')}")
    logs.append(f"  Confidence: {classification.get('confidence', 0):.2f}")
    logs.append(f"  Reasoning: {classification.get('reasoning', 'N/A')}")
    
    intent = classification.get("intent", "filter")
    
    # Route based on LLM classification
    if intent == "aggregate":
        # Generate and execute aggregation code
        agg_code, cap = capture_output(generate_aggregate_code, nl_query, df)
        logs.append(f"Aggregation code generated:\n{agg_code}\n{cap.strip()}")
        
        try:
            result = eval(agg_code)
            logs.append(f"Aggregation result:\n{result}")
            
            return {
                "action": "aggregate",
                "code": agg_code,
                "result": to_json_safe(result),
                "logs": "\n".join([l for l in logs if l]),
                "classification": classification
            }
        except Exception as e:
            logs.append(f"❌ Aggregation execution failed: {e}")
            return {
                "action": "aggregate",
                "code": agg_code,
                "error": str(e),
                "logs": "\n".join([l for l in logs if l]),
                "classification": classification
            }
    
    elif intent == "transform":
        # Generate and execute transformation code
        transform_code, cap = capture_output(generate_transform_code, nl_query, df)
        logs.append(f"Transform code generated:\n{transform_code}\n{cap.strip()}")
        
        try:
            out_df = eval(transform_code)
            out_df = ensure_outcome_labels(out_df)
            
            action = "transform"
            saved = save_filtered_result(out_df, action)
            
            # Drop Outcome_Labels from preview if requested
            preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
            
            return {
                "action": action,
                "code": transform_code,
                "preview": df_preview(preview_df),
                "rows": len(out_df),
                "logs": "\n".join([l for l in logs if l]),
                "saved": saved,
                "classification": classification
            }
        except Exception as e:
            logs.append(f"❌ Transform execution failed: {e}")
            return {
                "action": "transform",
                "code": transform_code,
                "error": str(e),
                "logs": "\n".join([l for l in logs if l]),
                "classification": classification
            }
    
    elif intent == "join":
        # Use existing join logic
        op_result, cap = capture_output(parse_nl_operation, nl_query, df, tbls)
        logs.append(cap.strip())
        
        if isinstance(op_result, pd.DataFrame):
            out_df = ensure_outcome_labels(op_result)
            action = "join"
            saved = save_filtered_result(out_df, action)
            
            # Drop Outcome_Labels from preview if requested
            preview_df = out_df.drop(columns=["Outcome_Labels"]) if drop_outcome_labels and "Outcome_Labels" in out_df.columns else out_df
            
            return {
                "action": action,
                "preview": df_preview(preview_df),
                "rows": len(out_df),
                "logs": "\n".join([l for l in logs if l]),
                "saved": saved,
                "classification": classification
            }

    # 7) Default: treat as filter (intent == "filter" or "create_column" or fallback)
    if not outcome:
        # If the user didn't provide an outcome, tag it generically as "Derived:filter"
        derived_outcome = "Derived:filter"
    else:
        derived_outcome = outcome

    condition, cap = capture_output(nl_to_condition, nl_query, df)
    logs.append("nl_to_condition:\n" + cap.strip())

    out_df, cap = capture_output(safe_apply, df, condition, derived_outcome)
    logs.append(cap.strip())

    # Determine matched subset
    mask = out_df["Outcome_Labels"].apply(lambda x: derived_outcome in x)
    matched_rows = int(mask.sum())
    df_to_save = out_df[mask].reset_index(drop=True) if filter_matched else out_df

    # Drop Outcome_Labels column if requested (for /nl/universal API)
    if drop_outcome_labels and "Outcome_Labels" in df_to_save.columns:
        df_to_save = df_to_save.drop(columns=["Outcome_Labels"])

    action = "filter"
    saved = persist_df_to_mongo(
        df_to_save,
        name=result_name or f"filtered_{slugify(derived_outcome)}_{uuid.uuid4().hex[:6]}",
        persist_to="filtered",
        overwrite=overwrite,
        outcome=derived_outcome,
        matched_rows=matched_rows,
        chunk_size=chunk_size
    )

    return {
        "action": action,
        "condition": condition,
        "preview": df_preview(df_to_save),
        "rows": len(df_to_save),
        "matched_rows": matched_rows,
        "logs": "\n".join([l for l in logs if l]),
        "saved": saved,
        "classification": classification  # Include LLM classification info
    }


# --------- General helpers used by APIs ---------

def _parse_rules_text(text: str) -> List[Dict[str, Any]]:
    try:
        payload = json.loads(text)
        if isinstance(payload, dict) and isinstance(payload.get("rules"), list):
            return payload["rules"]
        if isinstance(payload, list):
            return payload
    except Exception:
        pass
    try:
        rows = list(csv.DictReader(io.StringIO(text)))
        if rows:
            return rows
    except Exception:
        pass
    return []