backend/app/agents/query_understanding.py

"""
Query Understanding Agent - classifies user intent and extracts entities.
First agent in the pipeline. Determines routing for the rest of the graph.
"""

import structlog

from app.agents.intent_classifier import build_chat_response, classify_intent
from app.agents.state import AgentState

logger = structlog.get_logger()

# ── Complexity → retrieval_top_k mapping ────────────────
# Simple queries need fewer schema docs (less noise, faster).
# Complex queries need more context to join multiple tables correctly.
_COMPLEXITY_TOP_K: dict[str, int] = {
    "simple": 3,
    "moderate": 5,
    "complex": 8,
}


def _is_chat_input(query_lower: str) -> bool:
    """Backward-compatible wrapper for the old chat fast-path helper."""
    return classify_intent(query_lower).intent == "chat"


def _build_greeting_response(user_query: str) -> dict:
    """Backward-compatible wrapper for the old chat response helper."""
    return {
        "intent": "chat",
        "route_intent": "chat",
        "entities": [],
        "complexity": "simple",
        "friendly_message": build_chat_response(user_query),
    }


def query_understanding_node(state: AgentState, llm_router) -> dict:
    """
    Classify the user's intent and extract relevant entities.

    The chat/sql decision is rule-based and happens before any LLM call so
    casual messages cannot be forced into SQL generation.
    """
    user_query = state["user_query"]
    trace_id = state.get("trace_id", "unknown")

    logger.info("agent_started", agent="query_understanding", trace_id=trace_id, query=user_query)

    classification = classify_intent(user_query)

    if classification.intent == "chat":
        logger.info("intent_classified", intent="chat", method=classification.reason)
        return {
            "intent": "chat",
            "route_intent": "chat",
            "entities": [],
            "complexity": classification.complexity,
            "retrieval_top_k": _COMPLEXITY_TOP_K.get(classification.complexity, 5),
            "friendly_message": build_chat_response(user_query),
        }

    if classification.intent == "ambiguous":
        logger.info("intent_classified", intent="ambiguous", method=classification.reason)
        return {
            "intent": "ambiguous",
            "route_intent": "chat",
            "entities": _extract_entities_basic(user_query),
            "complexity": "simple",
            "retrieval_top_k": 5,  # default — ambiguous queries don't reach schema retrieval
            "friendly_message": _build_ambiguous_response(user_query),
        }

    if classification.route_intent == "meta_query":
        logger.info("intent_classified", intent="sql", route_intent="meta_query", method=classification.reason)
        return {
            "intent": "sql",
            "route_intent": "meta_query",
            "entities": _extract_entities_basic(user_query),
            "complexity": classification.complexity,
            "retrieval_top_k": _COMPLEXITY_TOP_K.get(classification.complexity, 5),
        }

    # ── Heuristic-only classification (no LLM call) ──────────
    # The rule-based classifier already determines the correct SQL sub-intent.
    # An LLM refinement call was previously made here but added ~800ms of latency
    # without improving downstream SQL generation quality — all SQL intents
    # follow the same pipeline path (retrieve → generate → validate → execute).
    route = classification.route_intent
    entities = _extract_entities_basic(user_query)
    complexity = classification.complexity
    top_k = _COMPLEXITY_TOP_K.get(complexity, 5)

    logger.info(
        "intent_classified",
        intent="sql",
        route_intent=route,
        entities=entities,
        complexity=complexity,
        retrieval_top_k=top_k,
        method="heuristic_fast",
    )

    return {
        "intent": "sql",
        "route_intent": route,
        "entities": entities,
        "complexity": complexity,
        "retrieval_top_k": top_k,
    }


def _extract_entities_basic(query: str) -> list[str]:
    """
    Extract table names and relevant entities based on keyword matching
    covering all 22 database tables and their common business synonyms.
    """
    query_lower = query.lower()
    
    # Map keywords and synonyms to the actual database table names
    keyword_map = {
        "subscriptions": ["subscriptions"],
        "subscription": ["subscriptions"],
        "arr": ["subscriptions"],
        "mrr": ["subscriptions"],
        "nrr": ["subscriptions", "accounts"],
        "retention": ["subscriptions", "accounts"],
        "contracted_arr": ["subscriptions"],
        "contract": ["subscriptions"],
        
        "invoices": ["invoices"],
        "invoice": ["invoices"],
        "billing": ["invoices"],
        "bill": ["invoices"],
        "revenue": ["invoices", "subscriptions"],
        
        "payments": ["payments"],
        "payment": ["payments"],
        "paid": ["payments"],
        "transaction": ["payments"],
        
        "opportunities": ["opportunities"],
        "opportunity": ["opportunities"],
        "deal": ["opportunities"],
        "deals": ["opportunities"],
        "pipeline": ["opportunities"],
        "forecast": ["opportunities"],
        
        "support_tickets": ["support_tickets"],
        "ticket": ["support_tickets"],
        "tickets": ["support_tickets"],
        "support": ["support_tickets"],
        "csat": ["support_tickets"],
        
        "ticket_events": ["ticket_events"],
        "sla": ["ticket_events", "support_tickets"],
        "breach": ["ticket_events"],
        
        "incidents": ["incidents"],
        "incident": ["incidents"],
        "outage": ["incidents"],
        "downtime": ["incidents"],
        "severity": ["incidents"],
        
        "workspaces": ["workspaces"],
        "workspace": ["workspaces"],
        "environment": ["workspaces"],
        "env": ["workspaces"],
        
        "workspace_users": ["workspace_users"],
        "workspace_user": ["workspace_users"],
        
        "accounts": ["accounts"],
        "account": ["accounts"],
        "customer": ["accounts"],
        "customers": ["accounts"],
        "segment": ["accounts"],
        "industry": ["accounts"],
        "region": ["accounts"],
        "health": ["accounts"],
        "churn": ["accounts"],
        
        "contacts": ["contacts"],
        "contact": ["contacts"],
        "sponsor": ["contacts"],
        
        "employees": ["employees"],
        "employee": ["employees"],
        "manager": ["employees"],
        "rep": ["employees"],
        "staff": ["employees"],
        "hire": ["employees"],
        "salary": ["employees"],
        "salaries": ["employees"],
        
        "departments": ["departments"],
        "department": ["departments"],
        "cost center": ["departments"],
        
        "plans": ["plans"],
        "plan": ["plans"],
        "pricing": ["plans"],
        
        "products": ["products"],
        "product": ["products"],
        
        "feature_catalog": ["feature_catalog"],
        "feature": ["feature_catalog", "products"],
        "features": ["feature_catalog", "products"],
        
        "product_usage_daily": ["product_usage_daily"],
        "usage": ["product_usage_daily"],
        "queries": ["product_usage_daily"],
        "latency": ["product_usage_daily"],
        
        "query_audit_log": ["query_audit_log"],
        "audit": ["query_audit_log"],
        
        "users": ["workspace_users", "plainsql_users"],
        "user": ["workspace_users", "plainsql_users"],
    }
    
    matched_tables = set()
    
    # 1. First, check for exact/partial word matching
    import re
    words = re.findall(r"[a-z_0-9]+", query_lower)
    
    for word in words:
        if word in keyword_map:
            matched_tables.update(keyword_map[word])
            
    # 2. Check for multi-word phrases (like "cost center", "ticket event")
    for key, tables in keyword_map.items():
        if " " in key and key in query_lower:
            matched_tables.update(tables)
            
    # 3. Fallback: check if the actual table name is in the query text directly
    all_table_names = [
        "departments", "employees", "plans", "products", "feature_catalog",
        "accounts", "contacts", "workspaces", "workspace_users", "subscriptions",
        "invoices", "payments", "opportunities", "product_usage_daily", "support_tickets",
        "ticket_events", "incidents", "query_audit_log", "query_feedback", "conversations",
        "messages", "plainsql_users"
    ]
    for table in all_table_names:
        if table in query_lower:
            matched_tables.add(table)
            
    return list(matched_tables)


def _build_ambiguous_response(user_query: str) -> str:
    """Return a helpful response for queries that are too vague to generate SQL."""
    entities = _extract_entities_basic(user_query)

    if entities:
        table_name = entities[0]
        return (
            f"I found a reference to **{table_name}**, but your question is a bit vague. "
            f"Could you be more specific? For example:\n\n"
            f"• \"Show all {table_name}\"\n"
            f"• \"How many {table_name} are there?\"\n"
            f"• \"Show top 5 {table_name} by name\""
        )

    return (
        "I'm not sure what data you're looking for. I can query these tables: "
        "**employees**, **departments**, **products**, **customers**, **sales**.\n\n"
        "Try asking something specific like:\n"
        "• \"Show top 5 employees by salary\"\n"
        "• \"Total sales revenue by region\"\n"
        "• \"List products with low stock\"\n"
        "• \"Which department has the highest average salary?\""
    )