RULE/query_engine/query_interface.py

from __future__ import annotations
"""
query_engine/query_interface.py - Safe query interface for business users.

Provides a structured query interface that validates inputs against
a whitelist of allowed patterns and executes them via the rule engine.
"""


import logging
import re
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional

import pandas as pd

logger = logging.getLogger(__name__)

__all__ = ["QueryInterface", "QueryResult"]

# ─── Allowed pattern whitelist ─────────────────────────────────────────────────

ALLOWED_PATTERNS = [
    # Aggregation queries
    re.compile(r"(count|sum|average|avg|mean|max|min)\s+(of\s+)?[\w\s]+", re.I),
    re.compile(r"how many\s+[\w\s]+", re.I),
    re.compile(r"total\s+[\w\s]+", re.I),
    # Filtering
    re.compile(r"show\s+(me\s+)?[\w\s]+(where|with|having)\s+[\w\s]+", re.I),
    re.compile(r"filter\s+[\w\s]+", re.I),
    re.compile(r"list\s+(all\s+)?[\w\s]+", re.I),
    # Comparison / ranking
    re.compile(r"top\s+\d+\s+[\w\s]+", re.I),
    re.compile(r"[\w\s]+(greater|less|above|below|equal)\s+[\w\s]+", re.I),
    # Graph patterns
    re.compile(r"which\s+[\w\s]+(depend|trigger|link|connect)\w*", re.I),
    re.compile(r"[\w\s]+(depend|trigger|link|connect)\w+\s+on\s+[\w\s]+", re.I),
]


@dataclass
class QueryResult:
    """Result of a QueryInterface.execute() call."""
    success: bool
    query: str
    answer: str = ""
    data: List[Dict[str, Any]] = field(default_factory=list)
    row_count: int = 0
    error: Optional[str] = None
    engine: str = "dataframe"

    def to_dict(self) -> Dict[str, Any]:
        return {
            "success": self.success,
            "query": self.query,
            "answer": self.answer,
            "data": self.data,
            "row_count": self.row_count,
            "error": self.error,
            "engine": self.engine,
        }


class QueryInterface:
    """
    Safe business query interface with whitelist validation.

    Validates queries against allowed patterns before execution.
    Does NOT allow arbitrary AI-generated answers — only structured results.

    Usage:
        interface = QueryInterface(df=df)
        result = interface.execute("How many transactions have match_score > 90?")
    """

    def __init__(
        self,
        df: Optional[pd.DataFrame] = None,
        strict_whitelist: bool = False,
    ):
        """
        Args:
            df: DataFrame to query against.
            strict_whitelist: If True, reject queries not matching whitelist.
        """
        self._df = df
        self._strict = strict_whitelist

    def execute(self, query: str) -> QueryResult:
        """
        Execute a safe business query.

        Args:
            query: Natural language query string.

        Returns:
            QueryResult with structured data.
        """
        query = query.strip()
        if not query:
            return QueryResult(success=False, query=query, error="Empty query.")

        # Validate against whitelist
        if self._strict and not self._matches_whitelist(query):
            return QueryResult(
                success=False,
                query=query,
                error=(
                    "Query does not match allowed patterns. "
                    "Use structured queries like 'count of X', 'show Y where Z', etc."
                ),
            )

        if self._df is None or self._df.empty:
            return QueryResult(
                success=False,
                query=query,
                error="No dataset loaded. Upload a dataset first.",
            )

        # Parse and execute
        result = self._dispatch_query(query)
        return result

    def _matches_whitelist(self, query: str) -> bool:
        """Check if query matches any allowed pattern."""
        for pattern in ALLOWED_PATTERNS:
            if pattern.search(query):
                return True
        return False

    def _dispatch_query(self, query: str) -> QueryResult:
        """Simple pattern-based query dispatch for common business queries."""
        q_lower = query.lower()
        df = self._df

        try:
            # Count queries
            if any(kw in q_lower for kw in ["how many", "count"]):
                return self._execute_count(query, df, q_lower)

            # Aggregation queries
            if any(kw in q_lower for kw in ["sum", "total", "average", "avg", "mean", "max", "min"]):
                return self._execute_aggregation(query, df, q_lower)

            # Top N queries
            top_match = re.search(r"top\s+(\d+)", q_lower)
            if top_match:
                return self._execute_top_n(query, df, int(top_match.group(1)))

            # List / show all
            if any(kw in q_lower for kw in ["list", "show", "filter"]):
                return self._execute_list(query, df, q_lower)

            # Default: return preview
            return QueryResult(
                success=True,
                query=query,
                answer=f"Dataset has {len(df)} rows and {len(df.columns)} columns.",
                data=df.head(20).to_dict(orient="records"),
                row_count=len(df),
                engine="dataframe",
            )

        except Exception as exc:
            logger.error("[QueryInterface] Query dispatch failed: %s", exc)
            return QueryResult(
                success=False,
                query=query,
                error=str(exc),
            )

    def _execute_count(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult:
        """Handle count queries."""
        total = len(df)
        return QueryResult(
            success=True,
            query=query,
            answer=f"Total count: {total} rows.",
            data=[{"count": total}],
            row_count=1,
            engine="dataframe",
        )

    def _execute_aggregation(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult:
        """Handle aggregation queries on numeric columns."""
        numeric_cols = df.select_dtypes(include="number").columns.tolist()
        if not numeric_cols:
            return QueryResult(
                success=False, query=query, error="No numeric columns found."
            )

        stats = {}
        for col in numeric_cols[:5]:  # limit to first 5 numeric cols
            stats[col] = {
                "sum": float(df[col].sum()),
                "mean": float(df[col].mean()),
                "min": float(df[col].min()),
                "max": float(df[col].max()),
            }

        return QueryResult(
            success=True,
            query=query,
            answer=f"Aggregation statistics for {len(numeric_cols)} numeric columns.",
            data=[stats],
            row_count=len(numeric_cols),
            engine="dataframe",
        )

    def _execute_top_n(self, query: str, df: pd.DataFrame, n: int) -> QueryResult:
        """Handle top-N queries."""
        numeric_cols = df.select_dtypes(include="number").columns.tolist()
        if numeric_cols:
            sort_col = numeric_cols[0]
            result_df = df.nlargest(n, sort_col)
        else:
            result_df = df.head(n)

        return QueryResult(
            success=True,
            query=query,
            answer=f"Top {n} rows returned.",
            data=result_df.to_dict(orient="records"),
            row_count=len(result_df),
            engine="dataframe",
        )

    def _execute_list(self, query: str, df: pd.DataFrame, q_lower: str) -> QueryResult:
        """Handle list/show/filter queries."""
        preview = df.head(20).to_dict(orient="records")
        return QueryResult(
            success=True,
            query=query,
            answer=f"Showing first {len(preview)} of {len(df)} rows.",
            data=preview,
            row_count=len(df),
            engine="dataframe",
        )