core/intent_parser.py

"""
Intent parser — regex + keyword + fuzzy column resolution.

No LLM anywhere. Covers Student + Business tier operations:
  increase, decrease, filter, sort_asc, sort_desc,
  sum, average, count, min, max,
  find_replace, delete_column, rename_column, add_column,
  remove_duplicates, cast_type
"""
from __future__ import annotations

import re
from typing import Optional

from rapidfuzz import process as rf_process, fuzz as rf_fuzz
from core.column_registry import column_registry

# ── Operation keyword table ─────────────────────────────────────────
# Longer, more specific keywords score higher so "badhao" beats "bada"
# when both appear in a command.
OPERATION_KEYWORDS: dict[str, list[str]] = {
    "increase": [
        "badhao", "badha do", "increase", "barha do", "barhao",
        "zyada karo", "bada karo", "grow", "raise", "badha dijiye",
    ],
    "decrease": [
        "ghatao", "ghata do", "decrease", "kam karo", "kam kar do",
        "chhota karo", "reduce", "cut", "minus karo", "ghata dijiye",
    ],
    "filter": [
        "sirf", "only", "filter", "dikhao", "show only", "show me",
        "bas", "wale dikhao", "where", "jitne", "laao",
    ],
    "sort_asc": [
        "chota se bada", "ascending", "a to z", "low to high",
        "smallest first", "ascending order", "a-z", "asc",
    ],
    "sort_desc": [
        "bada se chota", "descending", "z to a", "high to low",
        "largest first", "descending order", "z-a", "desc",
        "bade se chhote",
    ],
    "sum": [
        "sum", "total", "jod", "yog", "add up", "total batao",
        "kul", "jama",
    ],
    "average": [
        "average", "avg", "mean", "samanya", "average nikalo",
    ],
    "count": [
        "count", "ginti", "kitne", "kitni rows", "count karo",
        "kitni", "rows kitne", "total rows",
    ],
    "min": [
        "minimum", "min", "sabse chhota", "lowest", "kam se kam",
    ],
    "max": [
        "maximum", "max", "sabse bada", "highest", "zyada se zyada",
    ],
    "find_replace": [
        "replace", "badlo", "change", "find", "dhundho",
        "substitute", "replace karo", "change karo", "naye se badlo",
    ],
    "delete_column": [
        "delete column", "column hatao", "column delete karo",
        "remove column", "column remove karo", "column hatado",
        "column drop karo",
    ],
    "rename_column": [
        "rename column", "column ka naam badlo", "column rename karo",
        "name change karo", "naam badlo", "rename karo",
    ],
    "add_column": [
        "add column", "naya column banao", "column add karo",
        "new column", "column create karo",
    ],
    "remove_duplicates": [
        "duplicate hatao", "duplicates remove karo", "unique rakho",
        "duplicate remove", "repeat hatao",
    ],
    "cast_type": [
        "type badlo", "data type change", "convert type",
        "type convert karo", "numeric banao", "string banao",
    ],
}

# ── Compiled regexes ───────────────────────────────────────────────
PERCENT_RE = re.compile(r"(\d+\.?\d*)\s*%", re.IGNORECASE)
NUMBER_RE = re.compile(r"(\d+\.?\d*)")

# Comparison operators (symbols)
_CMP_SYMBOLS = re.compile(r"([><=!]+)\s*([\d.]+|[^\s]+)")

HINDI_CMP_MAP: dict[str, str] = {
    "se zyada": ">",
    "se kam": "<",
    "ke barabar": "==",
    "se zyada ya barabar": ">=",
    "se kam ya barabar": "<=",
    "se bada": ">",
    "se chhota": "<",
    "ke equal": "==",
    "ke hi": "==",
}

# Words to strip when extracting filter values
_STOP_WORDS = [
    "sirf", "only", "filter", "dikhao", "show", "show only",
    "bas", "wale", "laao", "bhai", "ko", "ka", "ke", "ki",
    "mein", "hai", "hain", "karo", "karein", "sort",
    "bada", "chhota", "se", "nikalo", "batao",
]


# ══════════════════════════════════════════════════════════════════════
#  Public API
# ══════════════════════════════════════════════════════════════════════

def parse_intent(session_id: str, command: str) -> Optional[dict]:
    """Parse a natural-language command into a structured intent dict.

    Returns ``None`` when nothing can be resolved (the caller should
    return an *unresolved* response with column suggestions).
    """
    columns = column_registry.get_columns(session_id)
    if not columns:
        return None

    operation = _match_operation(command)
    if operation is None:
        return None

    # ── Operations that need special parsing ──────────────────────
    if operation == "remove_duplicates":
        col = _best_column(session_id, command, columns)
        return {"operation": "remove_duplicates", "column": col}

    if operation == "delete_column":
        col = _best_column(session_id, command, columns)
        if col is None:
            return None
        return {"operation": "delete_column", "column": col}

    if operation == "rename_column":
        return _parse_rename(session_id, command, columns)

    if operation == "find_replace":
        return _parse_find_replace(session_id, command, columns)

    if operation == "filter":
        return _parse_filter(session_id, command, columns)

    if operation == "cast_type":
        return _parse_cast(session_id, command, columns)

    # ── Standard: operation + column + optional value ────────────
    col = _best_column(session_id, command, columns)
    if col is None:
        return None

    value = _parse_value(command)

    return {"operation": operation, "column": col, "value": value}


# ══════════════════════════════════════════════════════════════════════
#  Internal helpers
# ══════════════════════════════════════════════════════════════════════

def _match_operation(command: str) -> Optional[str]:
    """Pick the operation with the highest keyword-match score."""
    cmd = command.lower()
    scores: dict[str, int] = {}
    for op, keywords in OPERATION_KEYWORDS.items():
        for kw in keywords:
            if kw in cmd:
                # Weight by keyword length so specific phrases beat short ones
                scores[op] = scores.get(op, 0) + len(kw)
    if not scores:
        return None
    return max(scores, key=scores.get)  # type: ignore[arg-type]


# Flat set of all operation keywords — used to skip them during column resolution
_ALL_OP_KEYWORDS: set[str] = set()
for _kws in OPERATION_KEYWORDS.values():
    _ALL_OP_KEYWORDS.update(_kws)
_ALL_OP_KEYWORDS.update(_STOP_WORDS)
_ALL_OP_KEYWORDS.update(["ko", "ka", "ke", "ki", "karo", "nikalo", "batao",
                          "hai", "hain", "mein", "se", "do", "dijiye"])


def _best_column(session_id: str, command: str, columns: list[str]) -> Optional[str]:
    """Fuzzy-resolve the best column from the command text.

    1. Try column_registry (O(1) alias + cached fuzzy).
    2. Fall back to direct rapidfuzz scan.

    Skips tokens that are known operation keywords (e.g. "average" won't
    false-match to column "Age").
    """
    tokens = command.split()

    # Filter out keyword tokens and pure-number tokens
    clean_tokens = [
        t for t in tokens
        if not re.fullmatch(r"[\d.]+%?", t)
        and t.lower() not in _ALL_OP_KEYWORDS
    ]
    candidates = clean_tokens + [
        " ".join(clean_tokens[i : i + 2]) for i in range(len(clean_tokens) - 1)
    ]

    # 1. Registry first (O(1) alias + cached fuzzy)
    for token in candidates:
        resolved = column_registry.resolve(session_id, token)
        if resolved:
            return resolved

    # 2. Direct rapidfuzz scan as fallback
    best_match, best_score = None, 0
    for cand in candidates:
        hit = rf_process.extractOne(cand, columns, scorer=rf_fuzz.WRatio)
        if hit and hit[1] > best_score:
            best_match, best_score = hit[0], hit[1]
    return best_match if best_score >= 78 else None


def _parse_value(command: str) -> Optional[float]:
    """Extract a numeric value. Percentage wins over absolute."""
    m = PERCENT_RE.search(command)
    if m:
        return float(m.group(1))
    m = NUMBER_RE.search(command)
    if m:
        return float(m.group(1))
    return None


def _parse_filter(session_id: str, command: str, columns: list[str]) -> Optional[dict]:
    """Resolve a filter command into {column, condition, filter_value}."""
    col = _best_column(session_id, command, columns)
    if col is None:
        return None

    cmd_lower = command.lower()

    # 1. Symbol comparison:  salary > 50000
    sym = _CMP_SYMBOLS.search(cmd_lower)
    if sym and col.lower() in cmd_lower:
        op_str, val_str = sym.group(1), sym.group(2)
        try:
            fval: str | float = float(val_str)
        except ValueError:
            fval = val_str.strip("'\"")
        return {
            "operation": "filter",
            "column": col,
            "condition": op_str,
            "filter_value": fval,
        }

    # 2. Hindi comparison:  salary 50000 se zyada
    for hindi_op, symbol in HINDI_CMP_MAP.items():
        if hindi_op in cmd_lower:
            num_match = re.search(
                r"(\d+\.?\d*)\s+" + re.escape(hindi_op), cmd_lower
            )
            if num_match:
                return {
                    "operation": "filter",
                    "column": col,
                    "condition": symbol,
                    "filter_value": float(num_match.group(1)),
                }

    # 3. Equality by presence:  "city Mumbai dikhao" → city == Mumbai
    stripped = cmd_lower
    for kw in _STOP_WORDS:
        stripped = stripped.replace(kw, "")
    stripped = stripped.replace(col.lower(), "", 1).strip()
    if stripped:
        stripped = re.sub(r"^[><=!]+\s*", "", stripped).strip()
        return {
            "operation": "filter",
            "column": col,
            "condition": "==",
            "filter_value": stripped,
        }

    return {"operation": "filter", "column": col, "condition": None, "filter_value": None}


def _parse_find_replace(session_id: str, command: str, columns: list[str]) -> Optional[dict]:
    """Extract old_value and new_value for find & replace."""
    col = _best_column(session_id, command, columns)
    if col is None:
        return None

    # Try quoted values first
    quoted = re.findall(r"""['"]([^'"]+)['"]""", command)
    if len(quoted) >= 2:
        return {
            "operation": "find_replace",
            "column": col,
            "old_value": quoted[0],
            "new_value": quoted[1],
        }

    # Try "X ko Y se badlo" / "replace X with Y" / "X ko Y replace karo"
    m = re.search(
        r"(\S+)\s+ko\s+(\S+)\s+(?:se\s+)?badlo"
        r"|replace\s+(\S+)\s+with\s+(\S+)"
        r"|(\S+)\s+ko\s+(\S+)\s+replace",
        command,
        re.IGNORECASE,
    )
    if m:
        groups = [g for g in m.groups() if g is not None]
        if len(groups) >= 2:
            return {
                "operation": "find_replace",
                "column": col,
                "old_value": groups[0],
                "new_value": groups[1],
            }

    return None


def _parse_rename(session_id: str, command: str, columns: list[str]) -> Optional[dict]:
    """Extract current column and desired new name."""
    col = _best_column(session_id, command, columns)
    if col is None:
        return None

    m = re.search(
        r"(?:naam|name)\s+(?:ko\s+)?(\S+)\s+(?:se\s+)?badlo"
        r"|rename\s+\S+\s+to\s+(\S+)",
        command,
        re.IGNORECASE,
    )
    if m:
        new_name = m.group(1) or m.group(2)
        if new_name:
            return {
                "operation": "rename_column",
                "column": col,
                "new_name": new_name.strip("'\" "),
            }
    return None


def _parse_cast(session_id: str, command: str, columns: list[str]) -> Optional[dict]:
    """Extract column and target type for type casting."""
    col = _best_column(session_id, command, columns)
    if col is None:
        return None

    cmd_lower = command.lower()
    target_dtype: str | None = None
    if "int" in cmd_lower or "numeric" in cmd_lower or "number" in cmd_lower:
        target_dtype = "Int64"
    elif "float" in cmd_lower or "decimal" in cmd_lower:
        target_dtype = "Float64"
    elif "str" in cmd_lower or "string" in cmd_lower or "text" in cmd_lower:
        target_dtype = "String"
    elif "bool" in cmd_lower:
        target_dtype = "Boolean"
    elif "date" in cmd_lower or "datetime" in cmd_lower:
        target_dtype = "Date"

    if target_dtype:
        return {"operation": "cast_type", "column": col, "target_dtype": target_dtype}
    return None