backend/core/mode_handler.py

"""
Production-Grade Mode Handler - $5M AI System
==============================================

5 MODES with proper query understanding:
1. RAG - Direct document-based answers
2. GraphRAG - Relationship and entity analysis
3. Hybrid - Comprehensive deep analysis
4. Vision - Image/chart analysis
5. Prediction - Forecasting with confidence

ZERO HARDCODING - Everything from user's real data.
"""

from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from enum import Enum


class AnalysisMode(Enum):
    RAG = "rag"
    GRAPHRAG = "graphrag"
    HYBRID = "hybrid"
    VISION = "vision"
    PREDICTION = "prediction"


@dataclass
class ModeResponse:
    """Response from any mode"""
    answer: str
    mode: AnalysisMode
    sources: List[str]
    confidence: float
    visualization: Optional[Dict] = None
    insights: Optional[List[str]] = None


# ============================================================================
# MODE-SPECIFIC PROMPTS - Each mode behaves differently
# ============================================================================

MODE_PROMPTS = {
    AnalysisMode.RAG: """You are in RAG MODE - Direct Document Search.

## YOUR ROLE:
Answer questions ONLY from the user's uploaded data.

## ANTI-HALLUCINATION RULES (CRITICAL):
1. Give DIRECT answers ONLY from data shown
2. If data doesn't contain the answer, say "This information is not in your uploaded data"
3. NEVER make up numbers, names, or facts
4. If uncertain, say "Based on available data, I can tell you..."

## RESPONSE STYLE:
- For numbers: Bold the key metric (e.g., **₹548,765**)
- For lists: Use tables for top N items
- For facts: One-line answers when possible
- Always cite which data source you used""",

    AnalysisMode.GRAPHRAG: """You are in GRAPHRAG MODE - Relationship Analysis.

## YOUR ROLE:
Analyze CONNECTIONS and RELATIONSHIPS between entities in the data.

## ANTI-HALLUCINATION RULES (CRITICAL):
1. Only describe relationships that EXIST in the data
2. NEVER invent connections between entities
3. Use phrases like "The data shows that X connects to Y"
4. If no relationship exists, state that clearly

## RESPONSE STYLE:
- Explain connections: "Customer_33 buys Product X, which is also purchased by..."
- Show network effects with evidence
- Highlight key relationship patterns with data backing""",

    AnalysisMode.HYBRID: """You are in HYBRID MODE - Deep Comprehensive Analysis.

## YOUR ROLE:
Provide THOROUGH analysis combining document search with relationship insights.

## ANTI-HALLUCINATION RULES (CRITICAL):
1. Every claim must be backed by data
2. Separate "data says" from "analysis suggests"
3. Mark speculative insights as "Potential insight: ..."
4. Never mix real data with assumptions

## RESPONSE STYLE:
- Start with executive summary
- Include key metrics with exact values
- Add relationship insights (labeled as such)
- End with actionable recommendations (marked as recommendations)""",

    AnalysisMode.VISION: """You are in VISION MODE - Visual Analysis.

## YOUR ROLE:
Analyze images, charts, and visual data accurately.

## ANTI-HALLUCINATION RULES (CRITICAL):
1. Only describe what is ACTUALLY visible
2. Don't invent data points not shown in images
3. If image quality prevents reading, say so
4. Estimates should be marked as "approximately"

## RESPONSE STYLE:
- Describe the visual first (chart type, axes, labels)
- Extract VISIBLE data points
- Provide analysis of observable patterns only""",

    AnalysisMode.PREDICTION: """You are in PREDICTION MODE - Forecasting with Confidence.

## YOUR ROLE:
Generate predictions and forecasts with MANDATORY confidence scoring.

## PREDICTION REQUIREMENTS (CRITICAL):
1. ALWAYS include confidence level (Low/Medium/High or percentage)
2. Show the trend/historical data driving the prediction
3. Provide prediction RANGE (e.g., "Expected: 100-120, most likely ~110")
4. List key factors affecting the forecast

## ANTI-HALLUCINATION RULES:
1. Base predictions ONLY on available historical data
2. If insufficient data for prediction, say so
3. Never guarantee future outcomes

## RESPONSE FORMAT:
📊 **Prediction**: [Value or Range]
📈 **Confidence**: [High/Medium/Low - X%]
📉 **Based on**: [Historical trend description]
⚠️ **Risk factors**: [What could change this]
💡 **Recommendation**: [Action to take]"""
}



# ============================================================================
# QUERY UNDERSTANDING - What does user want?
# ============================================================================

def understand_query(query: str, mode: AnalysisMode) -> Dict[str, Any]:
    """
    Understand what the user wants based on query and mode.
    
    Returns understanding dict with:
    - intent: what they want
    - entities: customers/products mentioned
    - metric: what metric they care about
    - format: how they want the answer
    """
    q_lower = query.lower()
    
    # Detect intent
    if any(w in q_lower for w in ['total', 'sum', 'count', 'how much', 'how many']):
        intent = "aggregation"
    elif any(w in q_lower for w in ['top', 'best', 'highest', 'most']):
        intent = "ranking"
    elif any(w in q_lower for w in ['compare', 'versus', 'vs', 'difference']):
        intent = "comparison"
    elif any(w in q_lower for w in ['trend', 'over time', 'monthly', 'growth']):
        intent = "trend"
    elif any(w in q_lower for w in ['predict', 'forecast', 'will', 'next']):
        intent = "prediction"
    elif any(w in q_lower for w in ['why', 'explain', 'reason']):
        intent = "explanation"
    elif any(w in q_lower for w in ['who', 'which customer', 'which product']):
        intent = "identification"
    elif any(w in q_lower for w in ['relationship', 'connected', 'buy together']):
        intent = "relationship"
    else:
        intent = "general"
    
    # Detect metric
    if any(w in q_lower for w in ['revenue', 'sales', 'amount', 'money']):
        metric = "revenue"
    elif any(w in q_lower for w in ['customer', 'client', 'buyer']):
        metric = "customers"
    elif any(w in q_lower for w in ['product', 'item', 'sku']):
        metric = "products"
    elif any(w in q_lower for w in ['order', 'transaction', 'purchase']):
        metric = "orders"
    else:
        metric = "revenue"  # Default
    
    # Detect response format
    if any(w in q_lower for w in ['brief', 'short', 'one line', 'quick']):
        format_type = "brief"
    elif any(w in q_lower for w in ['detail', 'explain', 'comprehensive']):
        format_type = "detailed"
    elif any(w in q_lower for w in ['table', 'list', 'show all']):
        format_type = "table"
    elif any(w in q_lower for w in ['chart', 'graph', 'visual']):
        format_type = "visual"
    else:
        format_type = "standard"
    
    # Extract number for limits
    import re
    limit_match = re.search(r'\b(\d+)\b', query)
    limit = int(limit_match.group(1)) if limit_match else 10
    
    return {
        "intent": intent,
        "metric": metric,
        "format": format_type,
        "limit": limit,
        "mode": mode.value
    }


# ============================================================================
# MODE HANDLERS - Each mode has its own logic
# ============================================================================

class ProductionModeHandler:
    """
    Production-grade mode handler.
    Routes queries to appropriate mode with proper understanding.
    """
    
    def __init__(self, user_id: str, currency_symbol: str = "₹"):
        self.user_id = user_id
        self.currency = currency_symbol
    
    def handle(
        self,
        query: str,
        mode: str = "rag",
        data_context: str = ""
    ) -> ModeResponse:
        """
        Handle query with appropriate mode.
        
        Args:
            query: User's question
            mode: Mode name (rag, graphrag, hybrid, vision, prediction)
            data_context: Available data context
            
        Returns:
            ModeResponse with answer and metadata
        """
        # Parse mode
        mode_enum = self._parse_mode(mode)
        
        # Understand query
        understanding = understand_query(query, mode_enum)
        print(f"[MODE] Query understanding: {understanding}")
        
        # Get mode-specific prompt
        mode_prompt = MODE_PROMPTS.get(mode_enum, MODE_PROMPTS[AnalysisMode.RAG])
        
        # Build prompt with context
        full_prompt = self._build_prompt(
            query, mode_prompt, data_context, understanding
        )
        
        # Generate response
        try:
            from core.llm import chat
            answer = chat(full_prompt, temperature=0.3, max_tokens=2000)
        except Exception as e:
            answer = f"I encountered an error: {str(e)[:100]}"
        
        # Get sources
        sources = self._get_sources(mode_enum)
        
        return ModeResponse(
            answer=answer,
            mode=mode_enum,
            sources=sources,
            confidence=0.85
        )
    
    def _parse_mode(self, mode: str) -> AnalysisMode:
        """Parse mode string to enum"""
        mode_map = {
            "rag": AnalysisMode.RAG,
            "graph": AnalysisMode.GRAPHRAG,
            "graphrag": AnalysisMode.GRAPHRAG,
            "hybrid": AnalysisMode.HYBRID,
            "vision": AnalysisMode.VISION,
            "prediction": AnalysisMode.PREDICTION,
        }
        return mode_map.get(mode.lower(), AnalysisMode.RAG)
    
    def _build_prompt(
        self,
        query: str,
        mode_prompt: str,
        data_context: str,
        understanding: Dict
    ) -> str:
        """Build complete prompt for LLM"""
        
        format_instructions = {
            "brief": "Answer in ONE sentence maximum.",
            "detailed": "Provide comprehensive analysis with bullet points.",
            "table": "Format answer as a clean markdown table.",
            "visual": "Describe the visualization that would best represent this.",
            "standard": "Give a clear, direct answer."
        }
        
        format_inst = format_instructions.get(
            understanding.get("format", "standard"),
            format_instructions["standard"]
        )
        
        return f"""{mode_prompt}

## USER'S DATA:
{data_context if data_context else "No specific data context provided."}

## QUERY UNDERSTANDING:
- Intent: {understanding.get('intent', 'general')}
- Metric focus: {understanding.get('metric', 'revenue')}
- Limit: {understanding.get('limit', 10)}

## FORMAT REQUIREMENT:
{format_inst}

## CRITICAL RULES:
1. Use ONLY the data provided above
2. NEVER make up numbers or facts
3. If data is missing, say what you CAN tell from available data
4. Currency: {self.currency}

## USER QUESTION:
{query}

## YOUR RESPONSE:"""
    
    def _get_sources(self, mode: AnalysisMode) -> List[str]:
        """Get source attribution for mode"""
        source_map = {
            AnalysisMode.RAG: ["Document Search", "Vector Index"],
            AnalysisMode.GRAPHRAG: ["Knowledge Graph", "Relationship Analysis"],
            AnalysisMode.HYBRID: ["Document Search", "Knowledge Graph"],
            AnalysisMode.VISION: ["Image Analysis", "Visual Processing"],
            AnalysisMode.PREDICTION: ["Prediction Engine", "Trend Analysis"],
        }
        return source_map.get(mode, ["AI Analysis"])


# ============================================================================
# CONVENIENCE FUNCTIONS
# ============================================================================

def handle_mode_query(
    user_id: str,
    query: str,
    mode: str = "rag",
    data_context: str = "",
    currency: str = "₹"
) -> Dict[str, Any]:
    """
    Handle a query with the appropriate mode.
    
    Usage:
        result = handle_mode_query(
            user_id="user123",
            query="What is total revenue?",
            mode="rag",
            data_context="Total revenue: ₹548,765"
        )
    """
    handler = ProductionModeHandler(user_id, currency)
    response = handler.handle(query, mode, data_context)
    
    return {
        "answer": response.answer,
        "mode": response.mode.value,
        "sources": response.sources,
        "confidence": response.confidence,
        "visualization": response.visualization,
        "insights": response.insights
    }


def get_mode_prompt(mode: str) -> str:
    """Get the prompt for a specific mode"""
    mode_map = {
        "rag": AnalysisMode.RAG,
        "graphrag": AnalysisMode.GRAPHRAG,
        "hybrid": AnalysisMode.HYBRID,
        "vision": AnalysisMode.VISION,
        "prediction": AnalysisMode.PREDICTION,
    }
    mode_enum = mode_map.get(mode.lower(), AnalysisMode.RAG)
    return MODE_PROMPTS.get(mode_enum, MODE_PROMPTS[AnalysisMode.RAG])