RULE/query_engine/query_router.py

from __future__ import annotations
"""
query_engine/query_router.py - QueryWeaver query routing orchestrator.

Routes natural-language queries to either:
  - The Pandas/DataFrame engine (existing RULE pipeline)
  - The FalkorDB graph engine (via QueryExecutor)

Uses QueryClassifier to determine intent, then dispatches accordingly.
Errors degrade gracefully — never breaks the existing RULE pipeline.
"""


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

import pandas as pd

from RULE.query_engine.dataset_registry import DatasetRegistry, get_registry
from RULE.query_engine.falkordb_client import FalkorDBClient, get_client
from RULE.query_engine.query_classifier import ClassificationResult, QueryClassifier
from RULE.query_engine.query_executor import ExecutionResult, QueryExecutor
from RULE.query_engine.schema_manager import SchemaManager, get_schema_manager

logger = logging.getLogger(__name__)

__all__ = ["QueryRouter", "RouterResult", "get_router"]


@dataclass
class RouterResult:
    """Result of routing a natural-language query."""
    engine: str                       # "dataframe" or "graph"
    handled: bool                     # True if query was handled by graph engine
    classification: Optional[ClassificationResult] = None
    execution: Optional[ExecutionResult] = None
    message: str = ""
    rows: List[Dict[str, Any]] = field(default_factory=list)
    cypher: str = ""

    def row_count(self) -> int:
        return len(self.rows)


class QueryRouter:
    """
    Routes NL queries to DataFrame or FalkorDB engine.

    Usage:
        router = QueryRouter()
        result = router.route("Which rules depend on AML_RULE?", df=df)
        if result.handled:
            print(result.rows)
        else:
            # Fall through to existing DataFrame pipeline
    """

    def __init__(
        self,
        client: Optional[FalkorDBClient] = None,
        schema_manager: Optional[SchemaManager] = None,
        registry: Optional[DatasetRegistry] = None,
    ):
        self._client = client or get_client()
        self._schema_manager = schema_manager or get_schema_manager()
        self._registry = registry or get_registry()
        self._classifier = QueryClassifier()
        self._executor = QueryExecutor(
            client=self._client,
            schema_manager=self._schema_manager,
        )

    def route(
        self,
        nl_query: str,
        df: Optional[pd.DataFrame] = None,
        graph_name_override: Optional[str] = None,
    ) -> RouterResult:
        """
        Route a natural-language query to the appropriate engine.

        Args:
            nl_query:            The natural-language query string.
            df:                  Currently loaded DataFrame (for context).
            graph_name_override: Force routing to a specific FalkorDB graph.

        Returns:
            RouterResult — always returned, check .handled and .engine.
        """
        if not nl_query or not nl_query.strip():
            return RouterResult(
                engine="none",
                handled=False,
                message="Empty query provided.",
            )

        # ── Step 1: Classify ──────────────────────────────────────────────
        try:
            df_columns = list(df.columns) if df is not None and not df.empty else []
            registry_summary = self._registry.get_registry_summary()

            classification = self._classifier.classify(
                nl_query=nl_query,
                df_columns=df_columns,
                registry_summary=registry_summary,
            )
        except Exception as exc:
            logger.warning("[QueryRouter] Classification failed: %s", exc)
            return RouterResult(
                engine="dataframe",
                handled=False,
                message=f"Classification failed — falling back to DataFrame engine: {exc}",
            )

        # ── Step 2: Route ─────────────────────────────────────────────────
        if not classification.is_graph() and not graph_name_override:
            return RouterResult(
                engine="dataframe",
                handled=False,
                classification=classification,
                message="Query classified as DataFrame engine.",
            )

        # ── Step 3: Resolve graph name ────────────────────────────────────
        if graph_name_override:
            graph_name = graph_name_override
        else:
            dataset = classification.dataset
            if dataset:
                graph_key = self._registry.get_graph_key(dataset)
                graph_name = graph_key or self._client.default_graph
            else:
                graph_name = self._client.default_graph

        # ── Step 4: Execute graph query ───────────────────────────────────
        try:
            execution = self._executor.execute(nl_query, graph_name=graph_name)
        except Exception as exc:
            logger.error("[QueryRouter] Execution failed: %s", exc)
            return RouterResult(
                engine="graph",
                handled=False,
                classification=classification,
                message=f"Graph query execution failed: {exc}. Falling back to DataFrame engine.",
            )

        if not execution.success:
            logger.warning(
                "[QueryRouter] Graph execution unsuccessful: %s", execution.error
            )
            return RouterResult(
                engine="graph",
                handled=False,
                classification=classification,
                execution=execution,
                message=execution.error or "Graph query failed.",
                cypher=execution.cypher,
            )

        return RouterResult(
            engine="graph",
            handled=True,
            classification=classification,
            execution=execution,
            message=f"Graph query executed on '{graph_name}': {execution.row_count()} rows.",
            rows=execution.rows,
            cypher=execution.cypher,
        )


# ─── Singleton ─────────────────────────────────────────────────────────────────

_router_instance: Optional[QueryRouter] = None


def get_router() -> QueryRouter:
    """Return the singleton QueryRouter instance."""
    global _router_instance
    if _router_instance is None:
        _router_instance = QueryRouter()
    return _router_instance