core/parser/local_parser/parser.py

"""
Orchestration parser module.
Integrates spelling correction, synonym mapping, ONNX embeddings matching, and parameter extraction.
"""
from __future__ import annotations

import re
from typing import Any
from core.column_registry import column_registry
from core.parser.local_parser.spelling import SymSpell
from core.parser.local_parser.synonyms import SynonymMapper
from core.parser.local_parser.embeddings import EmbeddingModel
from core.parser.local_parser.ast_extractor import SafeMathEvaluator

# Standard operation keywords (from intent_parser.py)
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"],
}

class LocalIntentParser:
    """Orchestrates local heuristic parsing using spelling correction, synonyms, embeddings, and AST."""
    def __init__(self):
        self.synonym_mapper = SynonymMapper()
        self.embedding_model = EmbeddingModel()
        self.math_evaluator = SafeMathEvaluator()
        
        # Load spelling corrector with keywords
        self.sym_spell = SymSpell(max_edit_distance=2)
        for op_list in OPERATION_KEYWORDS.values():
            for kw in op_list:
                # Add word tokens to spelling index
                for token in re.findall(r'[a-zA-Z]+', kw):
                    self.sym_spell.add_word(token)

    def _get_columns(self, session_id: str) -> list[str]:
        return column_registry.get_columns(session_id) or []

    def parse_intent(self, session_id: str, command: str) -> dict[str, Any] | None:
        """Parse natural language command into structured dataframe operation JSON."""
        columns = self._get_columns(session_id)
        if not columns:
            return None

        # 1. Spelling correction
        # Build a spelling corrector with current column names dynamically included
        local_sym_spell = SymSpell(max_edit_distance=2)
        # copy keywords
        for w in self.sym_spell.words:
            local_sym_spell.add_word(w)
        # Add column names
        for col in columns:
            local_sym_spell.add_word(col)
            for token in re.findall(r'[a-zA-Z]+', col):
                local_sym_spell.add_word(token)

        corrected_cmd = local_sym_spell.correct_query(command)

        # 2. Synonym mapping and Hinglish normalization
        norm_cmd = self.synonym_mapper.normalize_text(corrected_cmd)

        # 3. Match operations
        operation = self._match_operation(norm_cmd)
        if not operation:
            return None

        # 4. Column matching (Exact or Semantic)
        column_match, confidence = self._resolve_column(norm_cmd, columns)

        # 5. Parameter extraction based on matched operation
        result: dict[str, Any] = {"operation": operation}

        if operation in ["remove_duplicates", "count"]:
            # Column is optional for remove_duplicates and count
            if column_match:
                result["column"] = column_match
            else:
                result["column"] = None
            result["confidence"] = "high"
            return result

        if operation == "delete_column":
            if not column_match:
                return None
            result["column"] = column_match
            result["confidence"] = "high"
            return result

        if operation == "rename_column":
            # Extract new name from patterns like "rename X to Y"
            rename_match = re.search(r"rename\s+(?:column\s+)?(\w+)\s+(?:to|as)\s+(\w+)", norm_cmd, re.IGNORECASE)
            if rename_match:
                old_name_cand = rename_match.group(1)
                new_name = rename_match.group(2)
                # Resolve old name using columns list
                resolved_old, _ = self._resolve_column(old_name_cand, columns)
                result["column"] = resolved_old or column_match
                result["new_name"] = new_name
                result["confidence"] = "high" if result["column"] else "low"
                return result
            # Try fallback: split by "to" or "as"
            parts = re.split(r"\b(?:to|as)\b", norm_cmd)
            if len(parts) >= 2:
                new_name = parts[-1].strip().split()[-1]
                result["column"] = column_match
                result["new_name"] = new_name
                result["confidence"] = "high" if column_match else "low"
                return result
            return None

        if operation == "find_replace":
            # Look for patterns: "replace A with B"
            replace_match = re.search(r"replace\s+(.+?)\s+with\s+(.+)", norm_cmd, re.IGNORECASE)
            if replace_match:
                old_val = replace_match.group(1).strip()
                new_val = replace_match.group(2).strip()
                # Remove column name references from old_val if present
                if column_match and old_val.startswith(column_match.lower()):
                    old_val = old_val[len(column_match):].strip()
                
                result["column"] = column_match
                result["old_value"] = self._try_parse_numeric(old_val)
                result["new_value"] = self._try_parse_numeric(new_val)
                result["confidence"] = "high" if column_match else "low"
                return result
            return None

        if operation == "add_column":
            # Patterns: "add column X with value Y", "new column X = Y"
            add_match = re.search(r"(?:add|new)\s+column\s+(\w+)(?:\s+(?:with|value|=)\s+(.+))?", norm_cmd, re.IGNORECASE)
            if add_match:
                col_name = add_match.group(1)
                default_val_str = add_match.group(2)
                default_val = self._try_parse_numeric(default_val_str) if default_val_str else None
                result["column"] = col_name
                result["value"] = default_val
                result["confidence"] = "high"
                return result
            return None

        if operation == "cast_type":
            # Check target datatype
            target_dtype = self._resolve_dtype(norm_cmd)
            if not target_dtype:
                return None
            result["column"] = column_match
            result["target_dtype"] = target_dtype
            result["confidence"] = "high" if column_match else "low"
            return result

        if operation in ["increase", "decrease"]:
            if not column_match:
                return None
            # Check if percentage
            is_percent = "%" in command or "percent" in norm_cmd
            # Find numbers
            num_match = re.search(r"(\d+(?:\.\d+)?)", norm_cmd)
            value = float(num_match.group(1)) if num_match else 0.0
            
            result["column"] = column_match
            result["value"] = value
            result["is_percent"] = is_percent
            result["confidence"] = "high"
            return result

        if operation == "filter":
            if not column_match:
                return None
            
            # Resolve comparison operator
            condition = "=="
            for op_sym in [">=", "<=", ">", "<", "!=", "=="]:
                if op_sym in norm_cmd:
                    condition = op_sym
                    break
            else:
                if "contains" in norm_cmd or "like" in norm_cmd:
                    condition = "contains"
                elif "equal" in norm_cmd:
                    condition = "=="
                elif "greater" in norm_cmd:
                    condition = ">"
                elif "less" in norm_cmd:
                    condition = "<"

            # Try to extract filter value
            # Split query by operator or column name to isolate the value
            filter_val_str = ""
            if condition in norm_cmd:
                parts = norm_cmd.split(condition, 1)
                if len(parts) == 2:
                    filter_val_str = parts[1].strip()
            else:
                # Fallback to finding numeric or text token at the end
                words = norm_cmd.split()
                if words:
                    filter_val_str = words[-1]

            # Clean filter val string
            filter_val_str = re.sub(r"\b(?:karo|dikhao|bas|only|show|hai|hain)\b", "", filter_val_str).strip()
            # Remove any trailing periods
            filter_val_str = filter_val_str.rstrip(".")

            filter_value = self._try_parse_numeric(filter_val_str)
            
            result["column"] = column_match
            result["condition"] = condition
            result["filter_value"] = filter_value
            result["confidence"] = "high"
            return result

        # Standard aggregation / simple operations (sum, average, min, max, sort_asc, sort_desc)
        if not column_match:
            return None
            
        result["column"] = column_match
        result["confidence"] = "high" if confidence >= 0.5 else "low"
        return result

    def _match_operation(self, command: str) -> str | None:
        """Choose 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 to favor more specific matches
                    scores[op] = scores.get(op, 0) + len(kw)
        if not scores:
            return None
        return max(scores, key=scores.get)

    def _resolve_column(self, command: str, columns: list[str]) -> tuple[str | None, float]:
        """Match the query to a column name, supporting exact and semantic resolution."""
        cmd_lower = command.lower()
        
        # 1. Exact / Substring Match (Case-insensitive)
        for col in columns:
            if col.lower() in cmd_lower:
                return col, 1.0

        # 2. Semantic Similarity Fallback
        # Extract keywords to reduce noise in the query string
        clean_text = cmd_lower
        for op_list in OPERATION_KEYWORDS.values():
            for kw in op_list:
                clean_text = re.sub(rf'\b{re.escape(kw)}\b', "", clean_text)
                
        # Clean extra spaces
        clean_text = " ".join(clean_text.split())
        if not clean_text:
            clean_text = cmd_lower

        try:
            return self.embedding_model.match_column(clean_text, columns, threshold=0.4)
        except Exception as e:
            print(f"[Local Parser] Semantic matching failed: {e}")
            # Fallback to first column or None
            return None, 0.0

    def _try_parse_numeric(self, val_str: str) -> Any:
        """Helper to cast string to int or float if applicable, strip quotes if string."""
        val_str = val_str.strip().strip("'\"")
        try:
            if "." in val_str:
                return float(val_str)
            return int(val_str)
        except ValueError:
            # Check boolean values
            if val_str.lower() == "true":
                return True
            if val_str.lower() == "false":
                return False
            return val_str

    def _resolve_dtype(self, command: str) -> str | None:
        """Match string to target datatype name."""
        cmd = command.lower()
        if "int" in cmd or "integer" in cmd or "numeric" in cmd or "number" in cmd:
            return "Int64"
        if "float" in cmd or "double" in cmd or "decimal" in cmd:
            return "Float64"
        if "string" in cmd or "text" in cmd or "character" in cmd:
            return "String"
        if "bool" in cmd or "boolean" in cmd:
            return "Boolean"
        if "date" in cmd or "time" in cmd:
            return "Date"
        return None