#!/usr/bin/env python3
"""
Context-Aware Memory Retrieval - Adaptive memory retrieval based on question types
==================================================================================

Implements intelligent memory retrieval that adapts its strategy based on the type
of question being asked, prioritizing different memory types and relevance factors.
"""

import os
import json
import re
from datetime import datetime, timedelta
from typing import Dict, List, Any, Tuple, Optional
from enum import Enum

from intelligence.predictive_memory_surfacing import get_predictive_memory_surfacer, MemoryRelevanceScorer
from core.memory.memory_manager import MemoryManager
from intelligence.work_context_intelligence import get_work_context_intelligence
from utils.logging import setup_logger

logger = setup_logger(__name__)


class QuestionType(Enum):
    """Different types of questions that require different retrieval strategies"""
    TECHNICAL = "technical"           # Code, debugging, implementation
    CONCEPTUAL = "conceptual"         # Understanding, explanations, theory
    HISTORICAL = "historical"         # Past events, what happened before
    DECISION = "decision"             # Making choices, planning
    CREATIVE = "creative"             # Ideas, brainstorming, innovation
    TROUBLESHOOTING = "troubleshooting"  # Problem solving, debugging
    LEARNING = "learning"             # How-to, tutorials, knowledge acquisition
    EMOTIONAL = "emotional"           # Feelings, relationships, personal
    STRATEGIC = "strategic"           # Long-term planning, goals
    OPERATIONAL = "operational"       # Day-to-day tasks, procedures


class ContextualRetrievalStrategy:
    """Base class for context-specific retrieval strategies"""
    
    def __init__(self, question_type: QuestionType):
        self.question_type = question_type
        self.memory_manager = MemoryManager()
        self.surfacer = get_predictive_memory_surfacer()
        
    def retrieve_memories(self, query: str, context: Dict[str, Any], max_memories: int = 10) -> List[Dict[str, Any]]:
        """Retrieve memories using strategy-specific logic"""
        raise NotImplementedError("Subclasses must implement retrieve_memories")
    
    def get_memory_filters(self) -> Dict[str, Any]:
        """Get filters specific to this strategy"""
        return {}
    
    def adjust_relevance_weights(self, base_weights: Dict[str, float]) -> Dict[str, float]:
        """Adjust relevance scoring weights for this question type"""
        return base_weights


class TechnicalRetrievalStrategy(ContextualRetrievalStrategy):
    """Strategy for technical questions - prioritizes code, bugs, solutions"""
    
    def retrieve_memories(self, query: str, context: Dict[str, Any], max_memories: int = 10) -> List[Dict[str, Any]]:
        # Enhance context for technical queries
        tech_context = context.copy()
        tech_context.update({
            'type': 'technical',
            'priority_types': ['solution', 'bug_fix', 'implementation', 'code_insight'],
            'keywords': self._extract_technical_keywords(query) + context.get('keywords', [])
        })
        
        # Get memories with technical focus
        memories = self.surfacer.surface_relevant_memories(tech_context, max_memories * 2)
        
        # Filter and re-score for technical relevance
        technical_memories = []
        for memory in memories:
            if self._is_technically_relevant(memory, query):
                # Boost score for technical memories
                memory['predicted_relevance'] = min(1.0, memory.get('predicted_relevance', 0) * 1.2)
                technical_memories.append(memory)
        
        return technical_memories[:max_memories]
    
    def _extract_technical_keywords(self, query: str) -> List[str]:
        """Extract technical keywords from query"""
        technical_patterns = [
            r'\b(function|class|method|variable|bug|error|exception)\b',
            r'\b(code|implementation|algorithm|debug|fix)\b',
            r'\b(API|database|server|client|frontend|backend)\b',
            r'\b(Python|JavaScript|TypeScript|SQL|HTML|CSS)\b'
        ]
        
        keywords = []
        for pattern in technical_patterns:
            matches = re.findall(pattern, query, re.IGNORECASE)
            keywords.extend(matches)
        
        return list(set(keywords))
    
    def _is_technically_relevant(self, memory: Dict[str, Any], query: str) -> bool:
        """Check if memory is technically relevant"""
        memory_content = memory.get('content', '').lower()
        memory_type = memory.get('type', '')
        
        # Check for technical content indicators
        technical_indicators = ['code', 'function', 'bug', 'error', 'implementation', 'debug', 'solution']
        
        for indicator in technical_indicators:
            if indicator in memory_content or indicator in memory_type:
                return True
        
        return False


class ConceptualRetrievalStrategy(ContextualRetrievalStrategy):
    """Strategy for conceptual questions - prioritizes explanations, insights"""
    
    def retrieve_memories(self, query: str, context: Dict[str, Any], max_memories: int = 10) -> List[Dict[str, Any]]:
        conceptual_context = context.copy()
        conceptual_context.update({
            'type': 'conceptual',
            'priority_types': ['insight', 'learning', 'explanation', 'understanding'],
            'temporal_preference': 'significance_over_recency'  # Older insights may be more valuable
        })
        
        memories = self.surfacer.surface_relevant_memories(conceptual_context, max_memories)
        
        # Boost memories with explanatory content
        for memory in memories:
            if self._has_explanatory_content(memory):
                memory['predicted_relevance'] = min(1.0, memory.get('predicted_relevance', 0) * 1.15)
        
        return sorted(memories, key=lambda x: x.get('predicted_relevance', 0), reverse=True)
    
    def _has_explanatory_content(self, memory: Dict[str, Any]) -> bool:
        """Check if memory contains explanatory content"""
        content = memory.get('content', '').lower()
        explanatory_phrases = [
            'because', 'this means', 'in other words', 'the reason',
            'explained', 'understanding', 'concept', 'principle'
        ]
        
        return any(phrase in content for phrase in explanatory_phrases)


class HistoricalRetrievalStrategy(ContextualRetrievalStrategy):
    """Strategy for historical questions - prioritizes temporal accuracy and chronology"""
    
    def retrieve_memories(self, query: str, context: Dict[str, Any], max_memories: int = 10) -> List[Dict[str, Any]]:
        historical_context = context.copy()
        historical_context.update({
            'type': 'historical',
            'temporal_preference': 'chronological',
            'priority_types': ['event', 'milestone', 'decision', 'change']
        })
        
        # Extract time references from query
        time_refs = self._extract_time_references(query)
        if time_refs:
            historical_context['time_constraints'] = time_refs
        
        memories = self.surfacer.surface_relevant_memories(historical_context, max_memories * 2)
        
        # Sort by temporal relevance and importance
        sorted_memories = self._sort_by_historical_relevance(memories, time_refs)
        
        return sorted_memories[:max_memories]
    
    def _extract_time_references(self, query: str) -> List[str]:
        """Extract time references from query"""
        time_patterns = [
            r'\b(yesterday|today|last week|last month|ago)\b',
            r'\b(before|after|during|when|since)\b',
            r'\b(january|february|march|april|may|june|july|august|september|october|november|december)\b',
            r'\b(2023|2024|2025)\b'
        ]
        
        time_refs = []
        for pattern in time_patterns:
            matches = re.findall(pattern, query, re.IGNORECASE)
            time_refs.extend(matches)
        
        return time_refs
    
    def _sort_by_historical_relevance(self, memories: List[Dict[str, Any]], time_refs: List[str]) -> List[Dict[str, Any]]:
        """Sort memories by historical relevance"""
        # For now, just sort by timestamp and relevance
        return sorted(memories, 
                     key=lambda x: (x.get('predicted_relevance', 0), 
                                   x.get('timestamp', '')), 
                     reverse=True)


class DecisionRetrievalStrategy(ContextualRetrievalStrategy):
    """Strategy for decision-making questions - prioritizes outcomes, lessons learned"""
    
    def retrieve_memories(self, query: str, context: Dict[str, Any], max_memories: int = 10) -> List[Dict[str, Any]]:
        decision_context = context.copy()
        decision_context.update({
            'type': 'decision',
            'priority_types': ['decision', 'outcome', 'lesson_learned', 'reflection'],
            'emotional_weight': 1.2  # Decisions often have emotional components
        })
        
        memories = self.surfacer.surface_relevant_memories(decision_context, max_memories)
        
        # Boost memories with decision outcomes
        for memory in memories:
            if self._contains_decision_content(memory):
                memory['predicted_relevance'] = min(1.0, memory.get('predicted_relevance', 0) * 1.3)
        
        return sorted(memories, key=lambda x: x.get('predicted_relevance', 0), reverse=True)
    
    def _contains_decision_content(self, memory: Dict[str, Any]) -> bool:
        """Check if memory contains decision-related content"""
        content = memory.get('content', '').lower()
        decision_indicators = [
            'decided', 'chose', 'selected', 'picked', 'went with',
            'outcome', 'result', 'consequence', 'learned',
            'should have', 'next time', 'mistake', 'success'
        ]
        
        return any(indicator in content for indicator in decision_indicators)


class ContextAwareMemoryRetriever:
    """Main class for context-aware memory retrieval"""
    
    def __init__(self):
        self.strategies = {
            QuestionType.TECHNICAL: TechnicalRetrievalStrategy(QuestionType.TECHNICAL),
            QuestionType.CONCEPTUAL: ConceptualRetrievalStrategy(QuestionType.CONCEPTUAL),
            QuestionType.HISTORICAL: HistoricalRetrievalStrategy(QuestionType.HISTORICAL),
            QuestionType.DECISION: DecisionRetrievalStrategy(QuestionType.DECISION),
            # Add more strategies as needed
        }
        
        # Default strategy for unclassified questions
        self.default_strategy = ConceptualRetrievalStrategy(QuestionType.CONCEPTUAL)
        
    def retrieve_contextual_memories(self, query: str, context: Dict[str, Any] = None, max_memories: int = 10) -> Dict[str, Any]:
        """Retrieve memories using context-aware strategy"""
        if context is None:
            context = {}
        
        # Classify the question type
        question_type = self._classify_question(query)
        
        # Get appropriate strategy
        strategy = self.strategies.get(question_type, self.default_strategy)
        
        # Retrieve memories using the strategy
        memories = strategy.retrieve_memories(query, context, max_memories)
        
        # Add metadata about the retrieval
        return {
            'memories': memories,
            'question_type': question_type.value,
            'strategy_used': strategy.__class__.__name__,
            'context_analysis': self._analyze_context(query, context),
            'retrieval_confidence': self._calculate_retrieval_confidence(memories, query)
        }
    
    def _classify_question(self, query: str) -> QuestionType:
        """Classify the type of question being asked"""
        query_lower = query.lower()
        
        # Technical indicators
        technical_patterns = [
            r'\b(how to|how do|implement|code|function|bug|error|debug)\b',
            r'\b(python|javascript|typescript|sql|api|database)\b',
            r'\b(fix|solve|troubleshoot|issue|problem)\b'
        ]
        
        # Historical indicators
        historical_patterns = [
            r'\b(what happened|when did|before|after|previously)\b',
            r'\b(last time|ago|history|past|remember when)\b'
        ]
        
        # Decision indicators
        decision_patterns = [
            r'\b(should I|which|choose|decide|recommend|best option)\b',
            r'\b(pros and cons|compare|versus|better)\b'
        ]
        
        # Conceptual indicators
        conceptual_patterns = [
            r'\b(what is|explain|understand|concept|why|meaning)\b',
            r'\b(theory|principle|idea|definition)\b'
        ]
        
        # Check patterns in order of specificity
        for pattern in technical_patterns:
            if re.search(pattern, query_lower):
                return QuestionType.TECHNICAL
        
        for pattern in historical_patterns:
            if re.search(pattern, query_lower):
                return QuestionType.HISTORICAL
        
        for pattern in decision_patterns:
            if re.search(pattern, query_lower):
                return QuestionType.DECISION
        
        for pattern in conceptual_patterns:
            if re.search(pattern, query_lower):
                return QuestionType.CONCEPTUAL
        
        # Default to conceptual for unclassified questions
        return QuestionType.CONCEPTUAL
    
    def _analyze_context(self, query: str, context: Dict[str, Any]) -> Dict[str, Any]:
        """Analyze the context for additional insights"""
        return {
            'query_length': len(query.split()),
            'context_keys': list(context.keys()),
            'question_complexity': 'simple' if len(query.split()) < 10 else 'complex',
            'has_temporal_reference': bool(re.search(r'\b(when|time|date|ago|before|after)\b', query.lower()))
        }
    
    def _calculate_retrieval_confidence(self, memories: List[Dict[str, Any]], query: str) -> float:
        """Calculate confidence in the retrieval results"""
        if not memories:
            return 0.0
        
        # Base confidence on average relevance score
        avg_relevance = sum(m.get('predicted_relevance', 0) for m in memories) / len(memories)
        
        # Boost confidence if we have high-relevance memories
        high_relevance_count = sum(1 for m in memories if m.get('predicted_relevance', 0) > 0.7)
        confidence_boost = min(0.2, high_relevance_count * 0.1)
        
        return min(1.0, avg_relevance + confidence_boost)


def get_context_aware_retriever() -> ContextAwareMemoryRetriever:
    """Get the global context-aware memory retriever instance"""
    return ContextAwareMemoryRetriever()


if __name__ == "__main__":
    # Test the context-aware retrieval
    retriever = get_context_aware_retriever()
    
    test_queries = [
        "How do I fix this Python error?",
        "What happened in our last meeting?",
        "Should I use React or Vue for this project?",
        "Explain the concept of neural networks"
    ]
    
    for query in test_queries:
        print(f"\nQuery: {query}")
        result = retriever.retrieve_contextual_memories(query)
        print(f"Question Type: {result['question_type']}")
        print(f"Strategy: {result['strategy_used']}")
        print(f"Memories Found: {len(result['memories'])}")
        print(f"Confidence: {result['retrieval_confidence']:.2f}")