query_router.py

"""
query_router.py — Hybrid Intent Detection Pipeline
=====================================================
Priority order:
  1. Rule-based parser (fast, deterministic)
  2. Cached results (TTL cache)
  3. LLM parser via Groq (fallback only)

Returns the same dict structure as groq_service.parse_query_with_llm()
so it's a drop-in replacement upstream.
"""

import re
import logging
import hashlib
from datetime import datetime, timedelta
from typing import Dict, Any, Optional
from cachetools import TTLCache

import groq_service

logger = logging.getLogger("TradeSense.Router")

# Cache for router results (1 hour TTL)
_router_cache = TTLCache(maxsize=500, ttl=3600)


# ---------------------------------------------------------------------------
# KEYWORD / REGEX PATTERNS
# ---------------------------------------------------------------------------

# Guide intent patterns
GUIDE_PATTERNS = [
    r"\b(?:how\s+to\s+open|open\s+(?:a\s+)?(?:demat|trading|zerodha|groww))\b",
    r"\b(?:zerodha|groww)\s+(?:account|sign\s*up|register)\b",
    r"\b(?:account\s+open|create\s+account|new\s+account)\b",
    r"\b(?:how\s+to\s+start\s+(?:investing|trading|stock))\b",
    r"\b(?:beginner|getting\s+started|first\s+(?:investment|trade))\b",
    r"\b(?:demat\s+account|what\s+is\s+demat)\b",
    r"\b(?:kaise\s+(?:khole|shuru|kare))\b",  # Hindi support
    r"\b(?:steps\s+to\s+(?:open|start|invest))\b",
]

# Education intent patterns
EDUCATION_PATTERNS = [
    r"\b(?:what\s+is|what\s+are|what's|explain|define|meaning\s+of|tell\s+me\s+about)\b",
    r"\b(?:how\s+does|how\s+do)\b.*\b(?:work|function)\b",
    r"\b(?:kya\s+(?:hai|hota|hoti))\b",  # Hindi: "kya hai"
    r"\b(?:difference\s+between|compare)\b",
    r"\b(?:basics?\s+of|introduction\s+to|learn\s+about)\b",
    r"\b(?:sip|mutual\s+fund|index\s+fund|etf|ipo|bond|dividend)\b",
    r"\b(?:risk\s+management|diversification|portfolio\s+(?:allocation|management))\b",
    r"\b(?:bull\s+market|bear\s+market|market\s+cap|p\/e\s+ratio)\b",
    r"\b(?:capital\s+gains?\s+tax|stcg|ltcg|80c|tax\s+(?:saving|saver))\b",
]

# Strategy intent patterns
STRATEGY_PATTERNS = [
    r"\b(?:if\s+i\s+(?:bought|invested|put)|what\s+if\s+i)\b",
    r"\b(?:bought|purchased|invested|invest)\b.*\b(?:ago|year|month|on|in\s+\d{4})\b",
    r"\b(?:profit|loss|return|p&l|pnl)\b.*\b(?:on|from|since)\b",
    r"\b(?:maine|agar|kharide)\b",  # Hindi
    r"\b(?:investment\s+(?:of|worth)|how\s+much\s+(?:profit|return|would))\b",
    r"\b\d+[kKmMlLcCrR]*\s+(?:in|on|of)\s+[a-zA-Z]+.*\b(?:ago|year|month|on|in\s+\d{4})\b",
    r"\b(?:ago|past|last)\s+(?:year|month|week|day)s?\b.*\b(?:profit|return|did)\b",
]

# Irrelevant intent patterns
IRRELEVANT_PATTERNS = [
    r"\b(?:joke|funny|weather|movie|song|recipe|cricket|football)\b",
    r"\b(?:hello|hi|hey|good\s+morning|howdy)\s*$",
    r"\b(?:who\s+(?:are\s+you|made\s+you)|your\s+name)\b",
    r"\b(?:thank|thanks|bye|goodbye)\b",
]

# Stock/asset name extraction patterns
STOCK_KEYWORDS = {
    "reliance": ("RELIANCE.NS", "Reliance Industries"),
    "tcs": ("TCS.NS", "Tata Consultancy Services"),
    "infosys": ("INFY.NS", "Infosys"),
    "infy": ("INFY.NS", "Infosys"),
    "hdfc": ("HDFCBANK.NS", "HDFC Bank"),
    "hdfc bank": ("HDFCBANK.NS", "HDFC Bank"),
    "sbi": ("SBIN.NS", "State Bank of India"),
    "icici": ("ICICIBANK.NS", "ICICI Bank"),
    "wipro": ("WIPRO.NS", "Wipro"),
    "hul": ("HINDUNILVR.NS", "Hindustan Unilever"),
    "bajaj": ("BAJFINANCE.NS", "Bajaj Finance"),
    "titan": ("TITAN.NS", "Titan Company"),
    "bharti": ("BHARTIARTL.NS", "Bharti Airtel"),
    "airtel": ("BHARTIARTL.NS", "Bharti Airtel"),
    "adani": ("ADANIENT.NS", "Adani Enterprises"),
    "maruti": ("MARUTI.NS", "Maruti Suzuki"),
    "tata motors": ("TATAMOTORS.NS", "Tata Motors"),
    "tata steel": ("TATASTEEL.NS", "Tata Steel"),
    "tata power": ("TATAPOWER.NS", "Tata Power"),
    "coal india": ("COALINDIA.NS", "Coal India"),
    "itc": ("ITC.NS", "ITC Limited"),
    "asian paints": ("ASIANPAINT.NS", "Asian Paints"),
    "sun pharma": ("SUNPHARMA.NS", "Sun Pharmaceutical"),
    "mrf": ("MRF.NS", "MRF Limited"),
    "kotak": ("KOTAKBANK.NS", "Kotak Mahindra Bank"),
    "axis bank": ("AXISBANK.NS", "Axis Bank"),
    "indusind": ("INDUSINDBK.NS", "IndusInd Bank"),
    "tech mahindra": ("TECHM.NS", "Tech Mahindra"),
    "hcl": ("HCLTECH.NS", "HCL Technologies"),
    "l&t": ("LT.NS", "Larsen & Toubro"),
    "larsen": ("LT.NS", "Larsen & Toubro"),
    "power grid": ("POWERGRID.NS", "Power Grid Corporation"),
    "ntpc": ("NTPC.NS", "NTPC Limited"),
    "ongc": ("ONGC.NS", "Oil & Natural Gas Corporation"),
    "gold": ("GC=F", "Gold"),
    "silver": ("SI=F", "Silver"),
    "crude": ("CL=F", "Crude Oil"),
    "crude oil": ("CL=F", "Crude Oil"),
    "bitcoin": ("BTC-USD", "Bitcoin"),
    "btc": ("BTC-USD", "Bitcoin"),
    "ethereum": ("ETH-USD", "Ethereum"),
    "eth": ("ETH-USD", "Ethereum"),
    "nifty": ("^NSEI", "Nifty 50"),
    "sensex": ("^BSESN", "BSE Sensex"),
    "apple": ("AAPL", "Apple Inc."),
    "aapl": ("AAPL", "Apple Inc."),
    "tesla": ("TSLA", "Tesla Inc."),
    "tsla": ("TSLA", "Tesla Inc."),
    "google": ("GOOGL", "Alphabet Inc."),
    "microsoft": ("MSFT", "Microsoft Corp."),
    "amazon": ("AMZN", "Amazon.com Inc."),
    "nvidia": ("NVDA", "NVIDIA Corp."),
    "nvda": ("NVDA", "NVIDIA Corp."),
    "meta": ("META", "Meta Platforms"),
    "facebook": ("META", "Meta Platforms"),
}

# Timeframe extraction patterns
TIMEFRAME_MAP = {
    r"1\s*(?:month|mo)": "1mo",
    r"3\s*(?:month|mo)": "3mo",
    r"6\s*(?:month|mo)": "6mo",
    r"1\s*(?:year|yr)": "1y",
    r"2\s*(?:year|yr)": "1y",
    r"5\s*(?:year|yr)": "5y",
    r"10\s*(?:year|yr)": "5y",
    r"(?:last|past)\s+(?:week|7\s*day)": "1mo",
    r"(?:last|past)\s+month": "1mo",
    r"(?:last|past)\s+3\s*month": "3mo",
    r"(?:last|past)\s+6\s*month": "6mo",
    r"(?:last|past)\s+year": "1y",
}


# ---------------------------------------------------------------------------
# DATE PARSING
# ---------------------------------------------------------------------------

def _parse_relative_date(query: str, today: str = "") -> Optional[str]:
    """Extract a buy_date from relative date references in the query."""
    if not today:
        today = datetime.now().strftime("%Y-%m-%d")
    today_dt = datetime.strptime(today, "%Y-%m-%d")
    
    # "X years ago"
    m = re.search(r"(\d+)\s*(?:year|yr)s?\s*ago", query, re.IGNORECASE)
    if m:
        years = int(m.group(1))
        d = today_dt - timedelta(days=years * 365)
        return d.strftime("%Y-%m-%d")
    
    # "X months ago"
    m = re.search(r"(\d+)\s*(?:month|mo)s?\s*ago", query, re.IGNORECASE)
    if m:
        months = int(m.group(1))
        d = today_dt - timedelta(days=months * 30)
        return d.strftime("%Y-%m-%d")
    
    # "last year" / "since last year"
    if re.search(r"(?:last|past|previous)\s+year", query, re.IGNORECASE):
        d = today_dt - timedelta(days=365)
        return d.strftime("%Y-%m-%d")
    
    # "in 2023", "in January 2023", "on 5th March 2023"
    m = re.search(
        r"(?:on|in|since)?\s*(\d{1,2})?\s*(?:st|nd|rd|th)?\s*"
        r"(january|february|march|april|may|june|july|august|september|october|november|december|"
        r"jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec)?\s*"
        r"\b(\d{4})\b",
        query, re.IGNORECASE
    )
    if m:
        day = int(m.group(1)) if m.group(1) else 1
        month_str = m.group(2)
        year = int(m.group(3))
        
        month_map = {
            "january": 1, "jan": 1, "february": 2, "feb": 2,
            "march": 3, "mar": 3, "april": 4, "apr": 4,
            "may": 5, "june": 6, "jun": 6,
            "july": 7, "jul": 7, "august": 8, "aug": 8,
            "september": 9, "sep": 9, "october": 10, "oct": 10,
            "november": 11, "nov": 11, "december": 12, "dec": 12,
        }
        month = month_map.get(month_str.lower(), 1) if month_str else 1
        
        try:
            d = datetime(year, month, min(day, 28))
            return d.strftime("%Y-%m-%d")
        except ValueError:
            return f"{year}-{month:02d}-01"
    
    return None


# ---------------------------------------------------------------------------
# QUANTITY / AMOUNT PARSING
# ---------------------------------------------------------------------------

def _parse_quantity(query: str) -> Dict[str, Any]:
    """Extract quantity, unit, and investment amount from query."""
    result = {"quantity": 1, "unit": "shares", "investment_amount": None}
    
    # Investment amount: "invested 1 lakh", "₹50000", "10k", "10000 rupees"
    m = re.search(r"(?:₹|rs\.?|inr)?\s*(\d+(?:\.\d+)?)\s*(?:lakh|lac)", query, re.IGNORECASE)
    if m:
        result["investment_amount"] = float(m.group(1)) * 100000
        return result
    
    m = re.search(r"(?:₹|rs\.?|inr)?\s*(\d+(?:\.\d+)?)\s*(?:crore|cr)", query, re.IGNORECASE)
    if m:
        result["investment_amount"] = float(m.group(1)) * 10000000
        return result

    m = re.search(r"(?:₹|rs\.?|inr)?\s*(\d+(?:\.\d+)?)\s*(?:k|thousand)\b", query, re.IGNORECASE)
    if m:
        result["investment_amount"] = float(m.group(1)) * 1000
        return result
    
    m = re.search(r"(?:₹|rs\.?|inr)\s*(\d+(?:,?\d+)*(?:\.\d+)?)", query, re.IGNORECASE)
    if m:
        amt = m.group(1).replace(",", "")
        result["investment_amount"] = float(amt)
        return result
    
    m = re.search(r"(\d+(?:,?\d+)*(?:\.\d+)?)\s*(?:rupees?|rs|inr|ruppp?e[s]*|bucks)", query, re.IGNORECASE)
    if m:
        amt = m.group(1).replace(",", "")
        result["investment_amount"] = float(amt)
        return result
    
    # Quantity: "10 shares", "5 grams", "2 units"
    m = re.search(r"(\d+(?:\.\d+)?)\s*(shares?|units?|grams?|g\b|kg|oz)", query, re.IGNORECASE)
    if m:
        result["quantity"] = int(float(m.group(1)))
        unit = m.group(2).lower()
        if unit.startswith("gram") or unit == "g":
            result["unit"] = "gram"
        elif unit == "kg":
            result["unit"] = "kg"
        elif unit == "oz":
            result["unit"] = "oz"
        else:
            result["unit"] = "shares"
        return result
    
    # Just a number before a stock name: "10 reliance"
    m = re.search(r"(\d+)\s+(?:shares?\s+(?:of|in)\s+)?([a-zA-Z])", query, re.IGNORECASE)
    if m:
        result["quantity"] = int(m.group(1))
    
    return result


# ---------------------------------------------------------------------------
# SYMBOL EXTRACTION
# ---------------------------------------------------------------------------

def _extract_symbol(query: str) -> Dict[str, Optional[str]]:
    """Extract stock/asset symbol and name from query using keyword map."""
    query_lower = query.lower().strip()
    
    # Check exact and partial matches (longer keys first for specificity)
    sorted_keywords = sorted(STOCK_KEYWORDS.keys(), key=len, reverse=True)
    
    for keyword in sorted_keywords:
        if keyword in query_lower:
            symbol, name = STOCK_KEYWORDS[keyword]
            return {"symbol": symbol, "asset_name": name}
    
    # Check if query contains a direct ticker-like pattern (e.g., "AAPL", "TCS.NS")
    m = re.search(r"\b([A-Z]{2,10}(?:\.NS|\.BO)?)\b", query.upper())
    if m:
        potential = m.group(1)
        # Don't match common English words
        skip = {"THE", "AND", "FOR", "HOW", "WHAT", "WHO", "WHY", "HAS", "HAD", "WAS",
                "ARE", "NOT", "BUT", "CAN", "DID", "HIS", "HER", "ITS", "OUT", "OUR",
                "ALL", "ANY", "NOW", "OLD", "NEW", "USE", "WAY", "DAY", "MAY",
                "SIP", "IPO", "TAX", "AGO"}
        if potential not in skip and len(potential) >= 2:
            return {"symbol": potential, "asset_name": potential}
    
    return {"symbol": None, "asset_name": None}


# ---------------------------------------------------------------------------
# TIMEFRAME EXTRACTION
# ---------------------------------------------------------------------------

def _extract_timeframe(query: str) -> str:
    """Extract analysis timeframe from query. Defaults to 6mo."""
    query_lower = query.lower()
    for pattern, tf in TIMEFRAME_MAP.items():
        if re.search(pattern, query_lower):
            return tf
    return "6mo"


# ---------------------------------------------------------------------------
# MAIN ROUTER — THE PIPELINE
# ---------------------------------------------------------------------------

def route_query(query: str, today: str = "") -> Dict[str, Any]:
    """
    Hybrid intent detection pipeline.
    
    Priority:
      1. Cache hit → return immediately
      2. Rule-based regex matching → deterministic, fast
      3. LLM fallback via Groq → only when rules can't determine intent
    
    Returns dict compatible with groq_service.parse_query_with_llm() output.
    """
    if not today:
        today = datetime.now().strftime("%Y-%m-%d")
    
    query_clean = query.strip()
    if not query_clean:
        return _default_result(query_clean, today)
    
    # --- 1. CACHE CHECK ---
    cache_key = hashlib.md5(query_clean.lower().encode()).hexdigest()
    if cache_key in _router_cache:
        logger.info(f"Router cache hit: {query_clean[:40]}...")
        return _router_cache[cache_key]
    
    # --- 2. RULE-BASED PARSING ---
    result = _rule_based_parse(query_clean, today)
    
    if result is not None:
        logger.info(f"Rule-based routing: intent={result.get('intent')}, symbol={result.get('symbol')}")
        _router_cache[cache_key] = result
        return result
    
    # --- 3. LLM FALLBACK ---
    logger.info(f"Rule-based routing failed, falling back to LLM for: {query_clean[:40]}...")
    llm_result = groq_service.parse_query_with_llm(query_clean, today)
    
    if llm_result:
        # Mark this as LLM-parsed for meta tracking
        llm_result["_parsed_by"] = "llm"
        _router_cache[cache_key] = llm_result
        return llm_result
    
    # --- 4. ABSOLUTE FALLBACK ---
    logger.warning(f"All parsing failed for: {query_clean[:40]}. Using safe default.")
    fallback = _default_result(query_clean, today)
    _router_cache[cache_key] = fallback
    return fallback


def _rule_based_parse(query: str, today: str) -> Optional[Dict[str, Any]]:
    """
    Attempt to parse query using regex/keyword rules only.
    Returns None if no confident match (triggers LLM fallback).
    """
    query_lower = query.lower().strip()
    
    # --- CHECK IRRELEVANT FIRST ---
    for pattern in IRRELEVANT_PATTERNS:
        if re.search(pattern, query_lower):
            return {
                "intent": "irrelevant",
                "asset_name": None,
                "symbol": None,
                "cleaned_query": query,
                "_parsed_by": "rules",
            }
    
    # --- CHECK GUIDE INTENT ---
    for pattern in GUIDE_PATTERNS:
        if re.search(pattern, query_lower):
            return {
                "intent": "guide",
                "asset_name": None,
                "symbol": None,
                "cleaned_query": query,
                "_parsed_by": "rules",
            }
    
    # --- CHECK STRATEGY INTENT ---
    for pattern in STRATEGY_PATTERNS:
        if re.search(pattern, query_lower):
            sym_info = _extract_symbol(query)
            qty_info = _parse_quantity(query)
            buy_date = _parse_relative_date(query, today)
            
            return {
                "intent": "strategy",
                "asset_name": sym_info["asset_name"],
                "symbol": sym_info["symbol"],
                "quantity": qty_info["quantity"],
                "unit": qty_info["unit"],
                "investment_amount": qty_info["investment_amount"],
                "buy_date": buy_date,
                "timeframe": _extract_timeframe(query),
                "cleaned_query": query,
                "_parsed_by": "rules",
            }
    
    # --- CHECK EDUCATION INTENT ---
    for pattern in EDUCATION_PATTERNS:
        if re.search(pattern, query_lower):
            # Check if it's actually asking about a specific stock price
            sym_info = _extract_symbol(query)
            if sym_info["symbol"] and not _is_educational_topic(query_lower):
                # "What is the price of Reliance" → analysis, not education
                break
            
            return {
                "intent": "education",
                "asset_name": sym_info.get("asset_name"),
                "symbol": None,
                "cleaned_query": query,
                "_parsed_by": "rules",
            }
    
    # --- CHECK IF IT'S A STOCK/ASSET QUERY (ANALYSIS) ---
    sym_info = _extract_symbol(query)
    if sym_info["symbol"]:
        return {
            "intent": "analysis",
            "asset_name": sym_info["asset_name"],
            "symbol": sym_info["symbol"],
            "timeframe": _extract_timeframe(query),
            "cleaned_query": query,
            "_parsed_by": "rules",
        }
    
    # --- NO CONFIDENT MATCH → return None to trigger LLM ---
    return None


def _is_educational_topic(query_lower: str) -> bool:
    """Check if the query is about educational financial topics vs specific asset prices."""
    edu_terms = [
        "sip", "mutual fund", "index fund", "risk management", "diversification",
        "ipo", "bond", "dividend", "portfolio", "asset allocation", "compounding",
        "capital gains", "tax", "demat", "trading account", "bull market", "bear market",
        "p/e ratio", "market cap", "fundamentals", "technical analysis",
    ]
    return any(term in query_lower for term in edu_terms)


def _default_result(query: str, today: str) -> Dict[str, Any]:
    """Safe fallback result when nothing can be parsed."""
    return {
        "intent": "analysis",
        "asset_name": "Nifty 50",
        "symbol": "^NSEI",
        "timeframe": "6mo",
        "cleaned_query": query or "market overview",
        "_parsed_by": "fallback",
    }