"""
Conversation Memory Module
Author: RSK World (https://rskworld.in)
Founder: Molla Samser
Designer & Tester: Rima Khatun
Contact: help@rskworld.in, +91 93305 39277
Year: 2026
Description: Advanced conversation context memory with persistence and learning
"""

import json
import sqlite3
import logging
import pickle
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Any, Tuple
import hashlib
import os
from collections import defaultdict, deque
import threading

logger = logging.getLogger(__name__)

class ConversationMemory:
    def __init__(self, db_path: str = "conversations.db"):
        self.db_path = db_path
        self.memory_cache = {}
        self.user_contexts = defaultdict(dict)
        self.global_context = defaultdict(list)
        self.lock = threading.Lock()
        
        # Memory settings
        self.max_short_term_memory = 50  # messages
        self.max_long_term_memory = 1000  # interactions
        self.context_retention_days = 30
        self.similarity_threshold = 0.7
        
        # Initialize database
        self._init_database()
        
        # Load existing contexts
        self._load_contexts()
    
    def _init_database(self):
        """Initialize SQLite database for conversation storage"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Create tables
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS conversations (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT NOT NULL,
                    session_id TEXT,
                    message_id TEXT UNIQUE,
                    role TEXT NOT NULL,
                    content TEXT NOT NULL,
                    language TEXT,
                    sentiment TEXT,
                    emotions TEXT,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    context_data TEXT,
                    importance_score REAL DEFAULT 1.0
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS user_profiles (
                    user_id TEXT PRIMARY KEY,
                    name TEXT,
                    preferences TEXT,
                    language_preference TEXT,
                    personality_traits TEXT,
                    interaction_count INTEGER DEFAULT 0,
                    first_interaction DATETIME DEFAULT CURRENT_TIMESTAMP,
                    last_interaction DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS conversation_summaries (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT NOT NULL,
                    session_id TEXT,
                    summary TEXT NOT NULL,
                    key_topics TEXT,
                    sentiment_summary TEXT,
                    timestamp DATETIME DEFAULT CURRENT_TIMESTAMP,
                    message_count INTEGER
                )
            ''')
            
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS knowledge_base (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    user_id TEXT,
                    question TEXT NOT NULL,
                    answer TEXT NOT NULL,
                    category TEXT,
                    confidence REAL,
                    usage_count INTEGER DEFAULT 0,
                    last_used DATETIME DEFAULT CURRENT_TIMESTAMP,
                    created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                )
            ''')
            
            # Create indexes for better performance
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_conversations_user_id ON conversations(user_id)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_conversations_timestamp ON conversations(timestamp)')
            cursor.execute('CREATE INDEX IF NOT EXISTS idx_knowledge_base_question ON knowledge_base(question)')
            
            conn.commit()
            conn.close()
            
            logger.info("Database initialized successfully")
            
        except Exception as e:
            logger.error(f"Database initialization error: {str(e)}")
    
    def add_message(self, user_id: str, role: str, content: str, 
                   language: str = None, sentiment: Dict = None, 
                   emotions: Dict = None, context_data: Dict = None,
                   session_id: str = None, importance_score: float = 1.0):
        """Add a message to conversation memory"""
        try:
            with self.lock:
                message_id = self._generate_message_id(content, user_id)
                
                # Store in database
                conn = sqlite3.connect(self.db_path)
                cursor = conn.cursor()
                
                cursor.execute('''
                    INSERT OR REPLACE INTO conversations 
                    (user_id, session_id, message_id, role, content, language, 
                     sentiment, emotions, context_data, importance_score)
                    VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
                ''', (
                    user_id, session_id, message_id, role, content, language,
                    json.dumps(sentiment) if sentiment else None,
                    json.dumps(emotions) if emotions else None,
                    json.dumps(context_data) if context_data else None,
                    importance_score
                ))
                
                conn.commit()
                conn.close()
                
                # Update user context
                self._update_user_context(user_id, {
                    'role': role,
                    'content': content,
                    'language': language,
                    'sentiment': sentiment,
                    'emotions': emotions,
                    'timestamp': datetime.now().isoformat()
                })
                
                # Update user profile
                self._update_user_profile(user_id, language)
                
                logger.debug(f"Message added to memory for user {user_id}")
                
        except Exception as e:
            logger.error(f"Error adding message to memory: {str(e)}")
    
    def get_conversation_history(self, user_id: str, limit: int = 20, 
                               session_id: str = None) -> List[Dict]:
        """Get conversation history for a user"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            query = '''
                SELECT role, content, language, sentiment, emotions, 
                       timestamp, context_data, importance_score
                FROM conversations
                WHERE user_id = ?
            '''
            params = [user_id]
            
            if session_id:
                query += ' AND session_id = ?'
                params.append(session_id)
            
            query += ' ORDER BY timestamp DESC LIMIT ?'
            params.append(limit)
            
            cursor.execute(query, params)
            rows = cursor.fetchall()
            conn.close()
            
            conversations = []
            for row in rows:
                conversations.append({
                    'role': row[0],
                    'content': row[1],
                    'language': row[2],
                    'sentiment': json.loads(row[3]) if row[3] else None,
                    'emotions': json.loads(row[4]) if row[4] else None,
                    'timestamp': row[5],
                    'context_data': json.loads(row[6]) if row[6] else None,
                    'importance_score': row[7]
                })
            
            return conversations[::-1]  # Return in chronological order
            
        except Exception as e:
            logger.error(f"Error getting conversation history: {str(e)}")
            return []
    
    def get_context_summary(self, user_id: str, session_id: str = None) -> Dict:
        """Get a summary of conversation context"""
        try:
            # Get recent conversations
            conversations = self.get_conversation_history(
                user_id, limit=10, session_id=session_id
            )
            
            if not conversations:
                return {'summary': 'No previous conversation', 'topics': [], 'sentiment': 'neutral'}
            
            # Extract key information
            topics = self._extract_topics(conversations)
            sentiment_summary = self._analyze_sentiment_trends(conversations)
            language_preference = self._detect_language_preference(conversations)
            
            # Generate summary
            summary = self._generate_context_summary(conversations)
            
            return {
                'summary': summary,
                'topics': topics,
                'sentiment': sentiment_summary,
                'language_preference': language_preference,
                'message_count': len(conversations),
                'time_span': self._get_time_span(conversations)
            }
            
        except Exception as e:
            logger.error(f"Error getting context summary: {str(e)}")
            return {'summary': 'Error generating context', 'topics': [], 'sentiment': 'neutral'}
    
    def find_similar_conversations(self, user_id: str, query: str, 
                                 limit: int = 5) -> List[Dict]:
        """Find similar past conversations"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Simple text-based similarity search
            cursor.execute('''
                SELECT content, role, timestamp, sentiment, emotions
                FROM conversations
                WHERE user_id = ? AND content LIKE ?
                ORDER BY timestamp DESC
                LIMIT ?
            ''', (user_id, f'%{query}%', limit))
            
            rows = cursor.fetchall()
            conn.close()
            
            similar_conversations = []
            for row in rows:
                similar_conversations.append({
                    'content': row[0],
                    'role': row[1],
                    'timestamp': row[2],
                    'sentiment': json.loads(row[3]) if row[3] else None,
                    'emotions': json.loads(row[4]) if row[4] else None,
                    'similarity_score': self._calculate_similarity(query, row[0])
                })
            
            # Sort by similarity score
            similar_conversations.sort(key=lambda x: x['similarity_score'], reverse=True)
            
            return similar_conversations
            
        except Exception as e:
            logger.error(f"Error finding similar conversations: {str(e)}")
            return []
    
    def learn_from_conversation(self, user_id: str, question: str, 
                             answer: str, category: str = None, 
                             confidence: float = 0.8):
        """Learn from conversation and add to knowledge base"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Check if already exists
            cursor.execute('''
                SELECT id FROM knowledge_base 
                WHERE question = ? AND user_id = ?
            ''', (question, user_id))
            
            if cursor.fetchone():
                # Update existing entry
                cursor.execute('''
                    UPDATE knowledge_base 
                    SET answer = ?, confidence = ?, usage_count = usage_count + 1,
                        last_used = CURRENT_TIMESTAMP
                    WHERE question = ? AND user_id = ?
                ''', (answer, confidence, question, user_id))
            else:
                # Insert new entry
                cursor.execute('''
                    INSERT INTO knowledge_base 
                    (user_id, question, answer, category, confidence)
                    VALUES (?, ?, ?, ?, ?)
                ''', (user_id, question, answer, category, confidence))
            
            conn.commit()
            conn.close()
            
            logger.info(f"Learned from conversation: {question[:50]}...")
            
        except Exception as e:
            logger.error(f"Error learning from conversation: {str(e)}")
    
    def get_knowledge_base(self, user_id: str = None, 
                         category: str = None) -> List[Dict]:
        """Get knowledge base entries"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            query = 'SELECT question, answer, category, confidence, usage_count FROM knowledge_base WHERE 1=1'
            params = []
            
            if user_id:
                query += ' AND user_id = ?'
                params.append(user_id)
            
            if category:
                query += ' AND category = ?'
                params.append(category)
            
            query += ' ORDER BY usage_count DESC, confidence DESC'
            
            cursor.execute(query, params)
            rows = cursor.fetchall()
            conn.close()
            
            knowledge_base = []
            for row in rows:
                knowledge_base.append({
                    'question': row[0],
                    'answer': row[1],
                    'category': row[2],
                    'confidence': row[3],
                    'usage_count': row[4]
                })
            
            return knowledge_base
            
        except Exception as e:
            logger.error(f"Error getting knowledge base: {str(e)}")
            return []
    
    def get_user_profile(self, user_id: str) -> Dict:
        """Get user profile with preferences and traits"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT name, preferences, language_preference, personality_traits,
                       interaction_count, first_interaction, last_interaction
                FROM user_profiles
                WHERE user_id = ?
            ''', (user_id,))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                return {
                    'user_id': user_id,
                    'name': row[0],
                    'preferences': json.loads(row[1]) if row[1] else {},
                    'language_preference': row[2],
                    'personality_traits': json.loads(row[3]) if row[3] else {},
                    'interaction_count': row[4],
                    'first_interaction': row[5],
                    'last_interaction': row[6]
                }
            else:
                return self._create_default_profile(user_id)
                
        except Exception as e:
            logger.error(f"Error getting user profile: {str(e)}")
            return self._create_default_profile(user_id)
    
    def update_user_preferences(self, user_id: str, preferences: Dict):
        """Update user preferences"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                UPDATE user_profiles
                SET preferences = ?, last_interaction = CURRENT_TIMESTAMP
                WHERE user_id = ?
            ''', (json.dumps(preferences), user_id))
            
            conn.commit()
            conn.close()
            
            # Update cache
            if user_id in self.user_contexts:
                self.user_contexts[user_id]['preferences'] = preferences
                
        except Exception as e:
            logger.error(f"Error updating user preferences: {str(e)}")
    
    def cleanup_old_memories(self, days: int = None):
        """Clean up old conversation memories"""
        try:
            if days is None:
                days = self.context_retention_days
            
            cutoff_date = datetime.now() - timedelta(days=days)
            
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                DELETE FROM conversations
                WHERE timestamp < ?
            ''', (cutoff_date,))
            
            deleted_count = cursor.rowcount
            conn.commit()
            conn.close()
            
            logger.info(f"Cleaned up {deleted_count} old conversation records")
            
        except Exception as e:
            logger.error(f"Error cleaning up old memories: {str(e)}")
    
    def get_memory_statistics(self) -> Dict:
        """Get memory usage statistics"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Get various statistics
            stats = {}
            
            cursor.execute('SELECT COUNT(*) FROM conversations')
            stats['total_conversations'] = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(DISTINCT user_id) FROM conversations')
            stats['unique_users'] = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM knowledge_base')
            stats['knowledge_base_entries'] = cursor.fetchone()[0]
            
            cursor.execute('SELECT COUNT(*) FROM user_profiles')
            stats['user_profiles'] = cursor.fetchone()[0]
            
            # Get language distribution
            cursor.execute('''
                SELECT language, COUNT(*) 
                FROM conversations 
                WHERE language IS NOT NULL
                GROUP BY language
                ORDER BY COUNT(*) DESC
            ''')
            stats['language_distribution'] = dict(cursor.fetchall())
            
            # Get recent activity
            cursor.execute('''
                SELECT COUNT(*) FROM conversations 
                WHERE timestamp > datetime('now', '-1 day')
            ''')
            stats['last_24h_messages'] = cursor.fetchone()[0]
            
            conn.close()
            
            return stats
            
        except Exception as e:
            logger.error(f"Error getting memory statistics: {str(e)}")
            return {}
    
    def _generate_message_id(self, content: str, user_id: str) -> str:
        """Generate unique message ID"""
        timestamp = datetime.now().isoformat()
        content_hash = hashlib.md5(content.encode()).hexdigest()[:8]
        return f"{user_id}_{timestamp}_{content_hash}"
    
    def _update_user_context(self, user_id: str, message_data: Dict):
        """Update user context in memory"""
        if user_id not in self.user_contexts:
            self.user_contexts[user_id] = {'messages': deque(maxlen=self.max_short_term_memory)}
        
        self.user_contexts[user_id]['messages'].append(message_data)
        
        # Update other context data
        if 'language' in message_data and message_data['language']:
            self.user_contexts[user_id]['last_language'] = message_data['language']
        
        if 'sentiment' in message_data and message_data['sentiment']:
            self.user_contexts[user_id]['last_sentiment'] = message_data['sentiment']
    
    def _update_user_profile(self, user_id: str, language: str):
        """Update user profile with interaction data"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Check if profile exists
            cursor.execute('SELECT user_id FROM user_profiles WHERE user_id = ?', (user_id,))
            exists = cursor.fetchone()
            
            if exists:
                cursor.execute('''
                    UPDATE user_profiles
                    SET interaction_count = interaction_count + 1,
                        last_interaction = CURRENT_TIMESTAMP,
                        language_preference = COALESCE(?, language_preference)
                    WHERE user_id = ?
                ''', (language, user_id))
            else:
                cursor.execute('''
                    INSERT INTO user_profiles
                    (user_id, language_preference, interaction_count)
                    VALUES (?, ?, 1)
                ''', (user_id, language))
            
            conn.commit()
            conn.close()
            
        except Exception as e:
            logger.error(f"Error updating user profile: {str(e)}")
    
    def _extract_topics(self, conversations: List[Dict]) -> List[str]:
        """Extract key topics from conversations"""
        # Simple keyword extraction - could be enhanced with NLP
        topics = set()
        common_words = {'the', 'a', 'an', 'and', 'or', 'but', 'in', 'on', 'at', 'to', 'for', 'of', 'with', 'by'}
        
        for conv in conversations:
            words = conv['content'].lower().split()
            for word in words:
                word = word.strip('.,!?()[]{}"\'')
                if len(word) > 3 and word not in common_words:
                    topics.add(word)
        
        return list(topics)[:10]  # Return top 10 topics
    
    def _analyze_sentiment_trends(self, conversations: List[Dict]) -> str:
        """Analyze sentiment trends in conversations"""
        sentiments = []
        for conv in conversations:
            if conv.get('sentiment') and 'sentiment' in conv['sentiment']:
                sentiments.append(conv['sentiment']['sentiment'])
        
        if not sentiments:
            return 'neutral'
        
        positive_count = sentiments.count('positive')
        negative_count = sentiments.count('negative')
        
        if positive_count > negative_count:
            return 'positive'
        elif negative_count > positive_count:
            return 'negative'
        else:
            return 'neutral'
    
    def _detect_language_preference(self, conversations: List[Dict]) -> str:
        """Detect user's language preference"""
        languages = [conv['language'] for conv in conversations if conv.get('language')]
        if languages:
            # Return most common language
            from collections import Counter
            return Counter(languages).most_common(1)[0][0]
        return 'en'
    
    def _generate_context_summary(self, conversations: List[Dict]) -> str:
        """Generate a summary of conversation context"""
        if not conversations:
            return "No previous conversation"
        
        # Get first and last messages
        first_msg = conversations[0]['content'][:100]
        last_msg = conversations[-1]['content'][:100]
        
        return f"Conversation started with: '{first_msg}...' and most recent: '{last_msg}...'"
    
    def _get_time_span(self, conversations: List[Dict]) -> str:
        """Get time span of conversations"""
        if not conversations:
            return "No time span"
        
        start_time = datetime.fromisoformat(conversations[0]['timestamp'])
        end_time = datetime.fromisoformat(conversations[-1]['timestamp'])
        
        duration = end_time - start_time
        
        if duration.days > 0:
            return f"{duration.days} days"
        elif duration.seconds > 3600:
            hours = duration.seconds // 3600
            return f"{hours} hours"
        else:
            minutes = duration.seconds // 60
            return f"{minutes} minutes"
    
    def _calculate_similarity(self, query: str, text: str) -> float:
        """Calculate similarity between query and text"""
        # Simple word-based similarity
        query_words = set(query.lower().split())
        text_words = set(text.lower().split())
        
        if not query_words or not text_words:
            return 0.0
        
        intersection = query_words.intersection(text_words)
        union = query_words.union(text_words)
        
        return len(intersection) / len(union) if union else 0.0
    
    def _create_default_profile(self, user_id: str) -> Dict:
        """Create default user profile"""
        return {
            'user_id': user_id,
            'name': None,
            'preferences': {},
            'language_preference': 'en',
            'personality_traits': {},
            'interaction_count': 0,
            'first_interaction': None,
            'last_interaction': None
        }
    
    def _load_contexts(self):
        """Load existing contexts into memory"""
        try:
            # This could be enhanced to load recent contexts into memory
            logger.info("Contexts loaded successfully")
        except Exception as e:
            logger.error(f"Error loading contexts: {str(e)}")