src/context_manager.py

# context_manager.py
import sqlite3
import json
import logging
import uuid
import hashlib
import threading
from contextlib import contextmanager
from datetime import datetime, timedelta

logger = logging.getLogger(__name__)


class TransactionManager:
    """Manage database transactions with proper locking"""
    
    def __init__(self, db_path):
        self.db_path = db_path
        self._lock = threading.RLock()
        self._connections = {}
    
    @contextmanager
    def transaction(self, session_id=None):
        """Context manager for database transactions with automatic rollback"""
        conn = None
        cursor = None
        
        try:
            with self._lock:
                conn = sqlite3.connect(self.db_path, isolation_level='IMMEDIATE')
                conn.execute('PRAGMA journal_mode=WAL')  # Write-Ahead Logging for better concurrency
                conn.execute('PRAGMA busy_timeout=5000')  # 5 second timeout for locks
                cursor = conn.cursor()
                
                yield cursor
                
                conn.commit()
                logger.debug(f"Transaction committed for session {session_id}")
                
        except Exception as e:
            if conn:
                conn.rollback()
                logger.error(f"Transaction rolled back for session {session_id}: {e}")
            raise
        finally:
            if conn:
                conn.close()

class EfficientContextManager:
    def __init__(self, llm_router=None):
        self.session_cache = {}  # In-memory for active sessions
        self.cache_config = {
            "max_session_size": 10,  # MB per session
            "ttl": 3600,  # 1 hour
            "compression": "gzip",
            "eviction_policy": "LRU"
        }
        self.db_path = "sessions.db"
        self.llm_router = llm_router  # For generating context summaries
        logger.info(f"Initializing ContextManager with DB path: {self.db_path}")
        self.transaction_manager = TransactionManager(self.db_path)
        self._init_database()
        self.optimize_database_indexes()
    
    def _init_database(self):
        """Initialize database and create tables"""
        try:
            logger.info("Initializing database...")
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Create sessions table if not exists
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    user_id TEXT DEFAULT 'Test_Any',
                    created_at TIMESTAMP,
                    last_activity TIMESTAMP,
                    context_data TEXT,
                    user_metadata TEXT
                )
            """)
            
            # Add user_id column to existing sessions table if it doesn't exist
            try:
                cursor.execute("ALTER TABLE sessions ADD COLUMN user_id TEXT DEFAULT 'Test_Any'")
                logger.info("✓ Added user_id column to sessions table")
            except sqlite3.OperationalError:
                # Column already exists
                pass
            
            logger.info("✓ Sessions table ready")
            
            # Create interactions table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS interactions (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT REFERENCES sessions(session_id),
                    user_input TEXT,
                    context_snapshot TEXT,
                    created_at TIMESTAMP,
                    FOREIGN KEY(session_id) REFERENCES sessions(session_id)
                )
            """)
            logger.info("✓ Interactions table ready")
            
            # Create user_contexts table (persistent user persona summaries)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS user_contexts (
                    user_id TEXT PRIMARY KEY,
                    persona_summary TEXT,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
            """)
            logger.info("✓ User contexts table ready")
            
            # Create session_contexts table (session summaries)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS session_contexts (
                    session_id TEXT PRIMARY KEY,
                    user_id TEXT,
                    session_summary TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY(session_id) REFERENCES sessions(session_id),
                    FOREIGN KEY(user_id) REFERENCES user_contexts(user_id)
                )
            """)
            logger.info("✓ Session contexts table ready")
            
            # Create interaction_contexts table (individual interaction summaries)
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS interaction_contexts (
                    interaction_id TEXT PRIMARY KEY,
                    session_id TEXT,
                    user_input TEXT,
                    system_response TEXT,
                    interaction_summary TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    FOREIGN KEY(session_id) REFERENCES sessions(session_id)
                )
            """)
            logger.info("✓ Interaction contexts table ready")
            
            conn.commit()
            conn.close()
            
            # Update schema with new columns and tables for user change tracking
            self._update_database_schema()
            
            logger.info("Database initialization complete")
            
        except Exception as e:
            logger.error(f"Database initialization error: {e}", exc_info=True)
    
    def _update_database_schema(self):
        """Add missing columns and tables for user change tracking"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Add needs_refresh column to interaction_contexts
            try:
                cursor.execute("""
                    ALTER TABLE interaction_contexts 
                    ADD COLUMN needs_refresh INTEGER DEFAULT 0
                """)
                logger.info("✓ Added needs_refresh column to interaction_contexts")
            except sqlite3.OperationalError:
                pass  # Column already exists
            
            # Create user change log table
            cursor.execute("""
                CREATE TABLE IF NOT EXISTS user_change_log (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    old_user_id TEXT,
                    new_user_id TEXT,
                    timestamp TIMESTAMP,
                    FOREIGN KEY(session_id) REFERENCES sessions(session_id)
                )
            """)
            
            conn.commit()
            conn.close()
            logger.info("✓ Database schema updated successfully for user change tracking")
            
            # Update interactions table for deduplication
            self._update_interactions_table()
            
        except Exception as e:
            logger.error(f"Schema update error: {e}", exc_info=True)
    
    def _update_interactions_table(self):
        """Add interaction_hash column for deduplication"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Check if column already exists
            cursor.execute("PRAGMA table_info(interactions)")
            columns = [row[1] for row in cursor.fetchall()]
            
            # Add interaction_hash column if it doesn't exist
            if 'interaction_hash' not in columns:
                try:
                    cursor.execute("""
                        ALTER TABLE interactions 
                        ADD COLUMN interaction_hash TEXT
                    """)
                    logger.info("✓ Added interaction_hash column to interactions table")
                except sqlite3.OperationalError:
                    pass  # Column already exists
            
            # Create unique index for deduplication (this enforces uniqueness)
            try:
                cursor.execute("""
                    CREATE UNIQUE INDEX IF NOT EXISTS idx_interaction_hash_unique 
                    ON interactions(interaction_hash)
                """)
                logger.info("✓ Created unique index on interaction_hash")
            except sqlite3.OperationalError:
                # Index might already exist, try non-unique index as fallback
                cursor.execute("""
                    CREATE INDEX IF NOT EXISTS idx_interaction_hash 
                    ON interactions(interaction_hash)
                """)
            
            conn.commit()
            conn.close()
            logger.info("✓ Interactions table updated for deduplication")
            
        except Exception as e:
            logger.error(f"Error updating interactions table: {e}", exc_info=True)
    
    async def manage_context(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict:
        """
        Efficient context management with separated session/user caching
        STEP 1: Fetch User Context (if available) 
        STEP 2: Get Previous Interaction Contexts
        STEP 3: Combine for workflow use
        """
        # Use session-only cache key to prevent user_id conflicts
        session_cache_key = f"session_{session_id}"
        user_cache_key = f"user_{user_id}"
        
        # Get session context from cache
        session_context = self._get_from_memory_cache(session_cache_key)
        
        # Check if cached session context matches current user_id
        if session_context and session_context.get("user_id") != user_id:
            # User changed, invalidate session cache
            logger.info(f"User mismatch in cache for session {session_id}, invalidating cache")
            session_context = None
            if session_cache_key in self.session_cache:
                del self.session_cache[session_cache_key]
        
        # Get user context separately
        user_context = self._get_from_memory_cache(user_cache_key)
        
        if not session_context:
            # Retrieve from database with user context
            session_context = await self._retrieve_from_db(session_id, user_input, user_id)
            
            # Cache session context (cache invalidation for user changes is handled in _retrieve_from_db)
            self._warm_memory_cache(session_cache_key, session_context)
        
        # Handle user context separately to prevent loops
        if not user_context or not user_context.get("user_context_loaded"):
            user_context_data = await self.get_user_context(user_id)
            user_context = {
                "user_context": user_context_data,
                "user_context_loaded": True,
                "user_id": user_id
            }
            # Cache user context separately
            self._warm_memory_cache(user_cache_key, user_context)
        
        # Merge contexts without duplication
        merged_context = {
            **session_context,
            "user_context": user_context.get("user_context", ""),
            "user_context_loaded": True,
            "user_id": user_id  # Ensure current user_id is used
        }
        
        # Update context with new interaction
        updated_context = self._update_context(merged_context, user_input, user_id=user_id)
        
        return self._optimize_context(updated_context)
    
    async def get_user_context(self, user_id: str) -> str:
        """
        STEP 1: Fetch or generate User Context (500-token persona summary)
        Available for all interactions except first time per user
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Check if user context exists
            cursor.execute("""
                SELECT persona_summary FROM user_contexts WHERE user_id = ?
            """, (user_id,))
            
            row = cursor.fetchone()
            if row and row[0]:
                # Existing user context found
                conn.close()
                logger.info(f"✓ User context loaded for {user_id}")
                return row[0]
            
            # Generate new user context from all historical data
            logger.info(f"Generating new user context for {user_id}")
            
            # Fetch all historical Session and Interaction contexts for this user
            all_session_summaries = []
            all_interaction_summaries = []
            
            # Get all session contexts
            cursor.execute("""
                SELECT session_summary FROM session_contexts WHERE user_id = ?
                ORDER BY created_at DESC LIMIT 50
            """, (user_id,))
            for row in cursor.fetchall():
                if row[0]:
                    all_session_summaries.append(row[0])
            
            # Get all interaction contexts
            cursor.execute("""
                SELECT ic.interaction_summary 
                FROM interaction_contexts ic
                JOIN sessions s ON ic.session_id = s.session_id
                WHERE s.user_id = ?
                ORDER BY ic.created_at DESC LIMIT 100
            """, (user_id,))
            for row in cursor.fetchall():
                if row[0]:
                    all_interaction_summaries.append(row[0])
            
            conn.close()
            
            if not all_session_summaries and not all_interaction_summaries:
                # First time user - no context to generate
                logger.info(f"No historical data for {user_id} - first time user")
                return ""
            
            # Generate persona summary using LLM (500 tokens)
            historical_data = "\n\n".join(all_session_summaries + all_interaction_summaries[:20])
            
            if self.llm_router:
                prompt = f"""Generate a concise 500-token persona summary for user {user_id} based on their interaction history:

Historical Context:
{historical_data}

Create a persona summary that captures:
- Communication style and preferences
- Common topics and interests
- Interaction patterns
- Key information shared across sessions

Keep the summary concise and focused (approximately 500 tokens)."""
                
                try:
                    persona_summary = await self.llm_router.route_inference(
                        task_type="general_reasoning",
                        prompt=prompt,
                        max_tokens=500,
                        temperature=0.7
                    )
                    
                    if persona_summary and isinstance(persona_summary, str) and persona_summary.strip():
                        # Store in database
                        conn = sqlite3.connect(self.db_path)
                        cursor = conn.cursor()
                        cursor.execute("""
                            INSERT OR REPLACE INTO user_contexts (user_id, persona_summary, updated_at)
                            VALUES (?, ?, ?)
                        """, (user_id, persona_summary.strip(), datetime.now().isoformat()))
                        conn.commit()
                        conn.close()
                        
                        logger.info(f"✓ Generated and stored user context for {user_id}")
                        return persona_summary.strip()
                except Exception as e:
                    logger.error(f"Error generating user context: {e}", exc_info=True)
            
            # Fallback: Return empty if LLM fails
            logger.warning(f"Could not generate user context for {user_id} - using empty")
            return ""
            
        except Exception as e:
            logger.error(f"Error getting user context: {e}", exc_info=True)
            return ""
    
    async def generate_interaction_context(self, interaction_id: str, session_id: str, 
                                         user_input: str, system_response: str, 
                                         user_id: str = "Test_Any") -> str:
        """
        STEP 2: Generate Interaction Context (50-token summary)
        Called after each response
        """
        try:
            if not self.llm_router:
                return ""
            
            prompt = f"""Summarize this interaction in approximately 50 tokens:

User Input: {user_input[:200]}
System Response: {system_response[:300]}

Provide a brief summary capturing the key exchange."""
            
            try:
                summary = await self.llm_router.route_inference(
                    task_type="general_reasoning",
                    prompt=prompt,
                    max_tokens=50,
                    temperature=0.7
                )
                
                if summary and isinstance(summary, str) and summary.strip():
                    # Store in database
                    conn = sqlite3.connect(self.db_path)
                    cursor = conn.cursor()
                    cursor.execute("""
                        INSERT OR REPLACE INTO interaction_contexts 
                        (interaction_id, session_id, user_input, system_response, interaction_summary, created_at)
                        VALUES (?, ?, ?, ?, ?, ?)
                    """, (
                        interaction_id, 
                        session_id, 
                        user_input[:500], 
                        system_response[:1000],
                        summary.strip(),
                        datetime.now().isoformat()
                    ))
                    conn.commit()
                    conn.close()
                    
                    logger.info(f"✓ Generated interaction context for {interaction_id}")
                    return summary.strip()
            except Exception as e:
                logger.error(f"Error generating interaction context: {e}", exc_info=True)
            
            # Fallback on LLM failure
            return ""
            
        except Exception as e:
            logger.error(f"Error in generate_interaction_context: {e}", exc_info=True)
            return ""
    
    async def generate_session_context(self, session_id: str, user_id: str = "Test_Any") -> str:
        """
        FINAL STEP: Generate Session Context (100-token summary)
        Called at session end
        """
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Get all interaction contexts for this session
            cursor.execute("""
                SELECT interaction_summary FROM interaction_contexts 
                WHERE session_id = ?
                ORDER BY created_at ASC
            """, (session_id,))
            
            interaction_summaries = [row[0] for row in cursor.fetchall() if row[0]]
            conn.close()
            
            if not interaction_summaries:
                logger.info(f"No interactions to summarize for session {session_id}")
                return ""
            
            # Generate session summary using LLM (100 tokens)
            if self.llm_router:
                combined_context = "\n".join(interaction_summaries)
                
                prompt = f"""Summarize this session's interactions in approximately 100 tokens:

Interaction Summaries:
{combined_context}

Create a concise session summary capturing:
- Main topics discussed
- Key outcomes or information shared
- User's focus areas

Keep the summary concise (approximately 100 tokens)."""
                
                try:
                    session_summary = await self.llm_router.route_inference(
                        task_type="general_reasoning",
                        prompt=prompt,
                        max_tokens=100,
                        temperature=0.7
                    )
                    
                    if session_summary and isinstance(session_summary, str) and session_summary.strip():
                        # Store in database
                        conn = sqlite3.connect(self.db_path)
                        cursor = conn.cursor()
                        cursor.execute("""
                            INSERT OR REPLACE INTO session_contexts 
                            (session_id, user_id, session_summary, created_at)
                            VALUES (?, ?, ?, ?)
                        """, (session_id, user_id, session_summary.strip(), datetime.now().isoformat()))
                        conn.commit()
                        conn.close()
                        
                        logger.info(f"✓ Generated session context for {session_id}")
                        return session_summary.strip()
                except Exception as e:
                    logger.error(f"Error generating session context: {e}", exc_info=True)
            
            # Fallback on LLM failure
            return ""
            
        except Exception as e:
            logger.error(f"Error in generate_session_context: {e}", exc_info=True)
            return ""
    
    async def end_session(self, session_id: str, user_id: str = "Test_Any"):
        """
        FINAL STEP: Generate Session Context and clear cache
        """
        try:
            # Generate session context
            await self.generate_session_context(session_id, user_id)
            
            # Clear in-memory cache for this session (session-only key)
            session_cache_key = f"session_{session_id}"
            if session_cache_key in self.session_cache:
                del self.session_cache[session_cache_key]
                logger.info(f"✓ Cleared cache for session {session_id}")
            
        except Exception as e:
            logger.error(f"Error ending session: {e}", exc_info=True)
    
    def _clear_user_cache_on_change(self, session_id: str, new_user_id: str, old_user_id: str):
        """Clear cache entries when user changes"""
        if new_user_id != old_user_id:
            # Clear old composite cache keys
            old_cache_key = f"{session_id}_{old_user_id}"
            if old_cache_key in self.session_cache:
                del self.session_cache[old_cache_key]
                logger.info(f"Cleared old cache for user {old_user_id} on session {session_id}")
    
    def _optimize_context(self, context: dict) -> dict:
        """
        Optimize context for LLM consumption
        Format: [Interaction Context #N, #N-1, ...] + User Context
        """
        user_context = context.get("user_context", "")
        interaction_contexts = context.get("interaction_contexts", [])
        
        # Format interaction contexts as requested
        formatted_interactions = []
        for idx, ic in enumerate(interaction_contexts[:10]):  # Last 10 interactions
            formatted_interactions.append(f"[Interaction Context #{len(interaction_contexts) - idx}]\n{ic.get('summary', '')}")
        
        # Combine User Context + Interaction Contexts
        combined_context = ""
        if user_context:
            combined_context += f"[User Context]\n{user_context}\n\n"
        if formatted_interactions:
            combined_context += "\n\n".join(formatted_interactions)
        
        return {
            "session_id": context.get("session_id"),
            "user_id": context.get("user_id", "Test_Any"),
            "user_context": user_context,
            "interaction_contexts": interaction_contexts,
            "combined_context": combined_context,  # For direct use in prompts
            "preferences": context.get("preferences", {}),
            "active_tasks": context.get("active_tasks", []),
            "last_activity": context.get("last_activity")
        }
    
    def _get_from_memory_cache(self, cache_key: str) -> dict:
        """
        Retrieve context from in-memory session cache
        """
        return self.session_cache.get(cache_key)
    
    async def _retrieve_from_db(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict:
        """
        Retrieve session context with proper user_id synchronization
        Uses transactions to ensure atomic updates of database and cache
        """
        conn = None
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Use transaction to ensure atomic updates
            cursor.execute("BEGIN TRANSACTION")
            
            # Get session data (SQLite doesn't support FOR UPDATE, but transaction ensures consistency)
            cursor.execute("""
                SELECT context_data, user_metadata, last_activity, user_id
                FROM sessions 
                WHERE session_id = ?
            """, (session_id,))
            
            row = cursor.fetchone()
            
            if row:
                context_data = json.loads(row[0]) if row[0] else {}
                user_metadata = json.loads(row[1]) if row[1] else {}
                last_activity = row[2]
                session_user_id = row[3] if len(row) > 3 else user_id
                
                # Check for user_id change and update atomically
                user_changed = False
                if session_user_id != user_id:
                    logger.info(f"User change detected: {session_user_id} -> {user_id} for session {session_id}")
                    user_changed = True
                    
                    # Update session with new user_id
                    cursor.execute("""
                        UPDATE sessions 
                        SET user_id = ?, last_activity = ?
                        WHERE session_id = ?
                    """, (user_id, datetime.now().isoformat(), session_id))
                    
                    # Clear any cached interaction contexts for old user by marking for refresh
                    try:
                        cursor.execute("""
                            UPDATE interaction_contexts
                            SET needs_refresh = 1
                            WHERE session_id = ?
                        """, (session_id,))
                    except sqlite3.OperationalError:
                        # Column might not exist yet, will be created by schema update
                        pass
                    
                    # Log user change event
                    try:
                        cursor.execute("""
                            INSERT INTO user_change_log (session_id, old_user_id, new_user_id, timestamp)
                            VALUES (?, ?, ?, ?)
                        """, (session_id, session_user_id, user_id, datetime.now().isoformat()))
                    except sqlite3.OperationalError:
                        # Table might not exist yet, will be created by schema update
                        pass
                    
                    # Clear old cache entries when user changes
                    self._clear_user_cache_on_change(session_id, user_id, session_user_id)
                
                cursor.execute("COMMIT")
                
                # Get interaction contexts with refresh flag check
                try:
                    cursor.execute("""
                        SELECT interaction_summary, created_at, needs_refresh
                        FROM interaction_contexts
                        WHERE session_id = ? AND (needs_refresh IS NULL OR needs_refresh = 0)
                        ORDER BY created_at DESC
                        LIMIT 20
                    """, (session_id,))
                except sqlite3.OperationalError:
                    # Column might not exist yet, fall back to query without needs_refresh
                    cursor.execute("""
                        SELECT interaction_summary, created_at
                        FROM interaction_contexts
                        WHERE session_id = ?
                        ORDER BY created_at DESC
                        LIMIT 20
                    """, (session_id,))
                
                interaction_contexts = []
                for ic_row in cursor.fetchall():
                    # Handle both query formats (with and without needs_refresh)
                    if len(ic_row) >= 2:
                        summary = ic_row[0]
                        timestamp = ic_row[1]
                        needs_refresh = ic_row[2] if len(ic_row) > 2 else 0
                        
                        if summary and not needs_refresh:
                            interaction_contexts.append({
                                "summary": summary,
                                "timestamp": timestamp
                            })
                
                context = {
                    "session_id": session_id,
                    "user_id": user_id,
                    "interaction_contexts": interaction_contexts,
                    "preferences": user_metadata.get("preferences", {}),
                    "active_tasks": user_metadata.get("active_tasks", []),
                    "last_activity": last_activity,
                    "user_context_loaded": False,
                    "user_changed": user_changed
                }
                
                conn.close()
                return context
            else:
                # Create new session with transaction
                cursor.execute("""
                    INSERT INTO sessions (session_id, user_id, created_at, last_activity, context_data, user_metadata)
                    VALUES (?, ?, ?, ?, ?, ?)
                """, (session_id, user_id, datetime.now().isoformat(), datetime.now().isoformat(), "{}", "{}"))
                
                cursor.execute("COMMIT")
                conn.close()
                
                return {
                    "session_id": session_id,
                    "user_id": user_id,
                    "interaction_contexts": [],
                    "preferences": {},
                    "active_tasks": [],
                    "user_context_loaded": False,
                    "user_changed": False
                }
                
        except sqlite3.Error as e:
            logger.error(f"Database transaction error: {e}", exc_info=True)
            if conn:
                try:
                    conn.rollback()
                except:
                    pass
                conn.close()
            # Return safe fallback
            return {
                "session_id": session_id,
                "user_id": user_id,
                "interaction_contexts": [],
                "preferences": {},
                "active_tasks": [],
                "user_context_loaded": False,
                "error": str(e),
                "user_changed": False
            }
        except Exception as e:
            logger.error(f"Database retrieval error: {e}", exc_info=True)
            if conn:
                try:
                    conn.rollback()
                except:
                    pass
                conn.close()
            # Return safe fallback
            return {
                "session_id": session_id,
                "user_id": user_id,
                "interaction_contexts": [],
                "preferences": {},
                "active_tasks": [],
                "user_context_loaded": False,
                "error": str(e),
                "user_changed": False
            }
    
    def _warm_memory_cache(self, cache_key: str, context: dict):
        """
        Warm the in-memory cache with retrieved context
        """
        self.session_cache[cache_key] = context
    
    def _update_context(self, context: dict, user_input: str, response: str = None, user_id: str = "Test_Any") -> dict:
        """
        Update context with deduplication and idempotency checks
        Prevents duplicate context updates using interaction hashes
        """
        try:
            # Generate unique interaction hash to prevent duplicates
            interaction_hash = self._generate_interaction_hash(user_input, context["session_id"], user_id)
            
            # Check if this interaction was already processed
            if self._is_duplicate_interaction(interaction_hash):
                logger.info(f"Duplicate interaction detected, skipping update: {interaction_hash[:8]}")
                return context
            
            # Use transaction for atomic updates
            current_time = datetime.now().isoformat()
            with self.transaction_manager.transaction(context["session_id"]) as cursor:
                # Update session activity (only if last_activity is older to prevent unnecessary updates)
                cursor.execute("""
                    UPDATE sessions 
                    SET last_activity = ?, user_id = ?
                    WHERE session_id = ? AND (last_activity IS NULL OR last_activity < ?)
                """, (current_time, user_id, context["session_id"], current_time))
                
                # Store interaction with duplicate prevention using INSERT OR IGNORE
                session_context = {
                    "preferences": context.get("preferences", {}),
                    "active_tasks": context.get("active_tasks", [])
                }
                
                cursor.execute("""
                    INSERT OR IGNORE INTO interactions (
                        interaction_hash,
                        session_id, 
                        user_input, 
                        context_snapshot, 
                        created_at
                    ) VALUES (?, ?, ?, ?, ?)
                """, (
                    interaction_hash,
                    context["session_id"],
                    user_input,
                    json.dumps(session_context),
                    current_time
                ))
            
            # Mark interaction as processed (outside transaction)
            self._mark_interaction_processed(interaction_hash)
            
            # Update in-memory context
            context["last_interaction"] = user_input
            context["last_update"] = current_time
            
            logger.info(f"Context updated for session {context['session_id']} with hash {interaction_hash[:8]}")
            
            return context
            
        except Exception as e:
            logger.error(f"Error updating context: {e}", exc_info=True)
            return context
    
    def _generate_interaction_hash(self, user_input: str, session_id: str, user_id: str) -> str:
        """Generate unique hash for interaction to prevent duplicates"""
        # Use session_id, user_id, and user_input for exact duplicate detection
        # Normalize user input by stripping whitespace
        normalized_input = user_input.strip()
        content = f"{session_id}:{user_id}:{normalized_input}"
        return hashlib.sha256(content.encode()).hexdigest()
    
    def _is_duplicate_interaction(self, interaction_hash: str) -> bool:
        """Check if interaction was already processed"""
        # Keep a rolling window of recent interaction hashes in memory
        if not hasattr(self, '_processed_interactions'):
            self._processed_interactions = set()
        
        # Check in-memory cache first
        if interaction_hash in self._processed_interactions:
            return True
        
        # Also check database for persistent duplicates
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            # Check if interaction_hash column exists and query for duplicates
            cursor.execute("PRAGMA table_info(interactions)")
            columns = [row[1] for row in cursor.fetchall()]
            if 'interaction_hash' in columns:
                cursor.execute("""
                    SELECT COUNT(*) FROM interactions 
                    WHERE interaction_hash IS NOT NULL AND interaction_hash = ?
                """, (interaction_hash,))
                count = cursor.fetchone()[0]
                conn.close()
                return count > 0
            else:
                conn.close()
                return False
        except sqlite3.OperationalError:
            # Column might not exist yet, only check in-memory
            return interaction_hash in self._processed_interactions
    
    def _mark_interaction_processed(self, interaction_hash: str):
        """Mark interaction as processed"""
        if not hasattr(self, '_processed_interactions'):
            self._processed_interactions = set()
        self._processed_interactions.add(interaction_hash)
        
        # Limit memory usage by keeping only last 1000 hashes
        if len(self._processed_interactions) > 1000:
            # Keep most recent 500 entries (simple truncation)
            self._processed_interactions = set(list(self._processed_interactions)[-500:])
    
    async def manage_context_optimized(self, session_id: str, user_input: str, user_id: str = "Test_Any") -> dict:
        """
        Efficient context management with transaction optimization
        """
        # Use session-only cache key
        session_cache_key = f"session_{session_id}"
        
        # Try to get from cache first (no DB access)
        cached_context = self._get_from_memory_cache(session_cache_key)
        if cached_context and self._is_cache_valid(cached_context):
            logger.debug(f"Using cached context for session {session_id}")
            return cached_context
        
        # Use transaction for all DB operations
        with self.transaction_manager.transaction(session_id) as cursor:
            # Atomic session retrieval and update
            cursor.execute("""
                SELECT s.context_data, s.user_metadata, s.last_activity, s.user_id,
                       COUNT(ic.interaction_id) as interaction_count
                FROM sessions s
                LEFT JOIN interaction_contexts ic ON s.session_id = ic.session_id
                WHERE s.session_id = ?
                GROUP BY s.session_id
            """, (session_id,))
            
            row = cursor.fetchone()
            
            if row:
                # Parse existing session data
                context_data = json.loads(row[0] or '{}')
                user_metadata = json.loads(row[1] or '{}')
                last_activity = row[2]
                stored_user_id = row[3] or user_id
                interaction_count = row[4] or 0
                
                # Handle user change atomically
                if stored_user_id != user_id:
                    self._handle_user_change_atomic(cursor, session_id, stored_user_id, user_id)
                
                # Get interaction contexts efficiently
                interaction_contexts = self._get_interaction_contexts_atomic(cursor, session_id)
                
            else:
                # Create new session atomically
                cursor.execute("""
                    INSERT INTO sessions (session_id, user_id, created_at, last_activity, context_data, user_metadata)
                    VALUES (?, ?, datetime('now'), datetime('now'), '{}', '{}')
                """, (session_id, user_id))
                
                context_data = {}
                user_metadata = {}
                interaction_contexts = []
                interaction_count = 0
        
        # Load user context asynchronously (outside transaction)
        user_context = await self._load_user_context_async(user_id)
        
        # Build final context
        final_context = {
            "session_id": session_id,
            "user_id": user_id,
            "interaction_contexts": interaction_contexts,
            "user_context": user_context,
            "preferences": user_metadata.get("preferences", {}),
            "active_tasks": user_metadata.get("active_tasks", []),
            "interaction_count": interaction_count,
            "cache_timestamp": datetime.now().isoformat()
        }
        
        # Update cache
        self._warm_memory_cache(session_cache_key, final_context)
        
        return self._optimize_context(final_context)
    
    def _handle_user_change_atomic(self, cursor, session_id: str, old_user_id: str, new_user_id: str):
        """Handle user change within transaction"""
        logger.info(f"Handling user change in transaction: {old_user_id} -> {new_user_id}")
        
        # Update session
        cursor.execute("""
            UPDATE sessions 
            SET user_id = ?, last_activity = datetime('now')
            WHERE session_id = ?
        """, (new_user_id, session_id))
        
        # Log the change
        try:
            cursor.execute("""
                INSERT INTO user_change_log (session_id, old_user_id, new_user_id, timestamp)
                VALUES (?, ?, ?, datetime('now'))
            """, (session_id, old_user_id, new_user_id))
        except sqlite3.OperationalError:
            # Table might not exist yet
            pass
        
        # Invalidate related caches
        try:
            cursor.execute("""
                UPDATE interaction_contexts
                SET needs_refresh = 1
                WHERE session_id = ?
            """, (session_id,))
        except sqlite3.OperationalError:
            # Column might not exist yet
            pass
    
    def _get_interaction_contexts_atomic(self, cursor, session_id: str, limit: int = 20):
        """Get interaction contexts within transaction"""
        try:
            cursor.execute("""
                SELECT interaction_summary, created_at, interaction_id
                FROM interaction_contexts
                WHERE session_id = ? AND (needs_refresh IS NULL OR needs_refresh = 0)
                ORDER BY created_at DESC
                LIMIT ?
            """, (session_id, limit))
        except sqlite3.OperationalError:
            # Fallback if needs_refresh column doesn't exist
            cursor.execute("""
                SELECT interaction_summary, created_at, interaction_id
                FROM interaction_contexts
                WHERE session_id = ?
                ORDER BY created_at DESC
                LIMIT ?
            """, (session_id, limit))
        
        contexts = []
        for row in cursor.fetchall():
            if row[0]:
                contexts.append({
                    "summary": row[0],
                    "timestamp": row[1],
                    "id": row[2] if len(row) > 2 else None
                })
        
        return contexts
    
    async def _load_user_context_async(self, user_id: str):
        """Load user context asynchronously to avoid blocking"""
        try:
            # Check memory cache first
            user_cache_key = f"user_{user_id}"
            cached = self._get_from_memory_cache(user_cache_key)
            if cached:
                return cached.get("user_context", "")
            
            # Load from database
            return await self.get_user_context(user_id)
        except Exception as e:
            logger.error(f"Error loading user context: {e}")
            return ""
    
    def _is_cache_valid(self, cached_context: dict, max_age_seconds: int = 60) -> bool:
        """Check if cached context is still valid"""
        if not cached_context:
            return False
        
        cache_timestamp = cached_context.get("cache_timestamp")
        if not cache_timestamp:
            return False
        
        try:
            cache_time = datetime.fromisoformat(cache_timestamp)
            age = (datetime.now() - cache_time).total_seconds()
            return age < max_age_seconds
        except:
            return False
    
    def optimize_database_indexes(self):
        """Create database indexes for better query performance"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Create indexes for frequently queried columns
            indexes = [
                "CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)",
                "CREATE INDEX IF NOT EXISTS idx_sessions_last_activity ON sessions(last_activity)",
                "CREATE INDEX IF NOT EXISTS idx_interactions_session_id ON interactions(session_id)",
                "CREATE INDEX IF NOT EXISTS idx_interaction_contexts_session_id ON interaction_contexts(session_id)",
                "CREATE INDEX IF NOT EXISTS idx_interaction_contexts_created_at ON interaction_contexts(created_at)",
                "CREATE INDEX IF NOT EXISTS idx_user_change_log_session_id ON user_change_log(session_id)",
                "CREATE INDEX IF NOT EXISTS idx_user_contexts_updated_at ON user_contexts(updated_at)"
            ]
            
            for index in indexes:
                try:
                    cursor.execute(index)
                except sqlite3.OperationalError as e:
                    # Table might not exist yet, skip this index
                    logger.debug(f"Skipping index creation (table may not exist): {e}")
            
            # Analyze database for query optimization
            try:
                cursor.execute("ANALYZE")
            except sqlite3.OperationalError:
                # ANALYZE might not be available in all SQLite versions
                pass
            
            conn.commit()
            conn.close()
            
            logger.info("✓ Database indexes optimized successfully")
            
        except Exception as e:
            logger.error(f"Error optimizing database indexes: {e}", exc_info=True)
    
    def _extract_entities(self, context: dict) -> list:
        """
        Extract essential entities from context
        """
        # TODO: Implement entity extraction
        return []
    
    def _generate_summary(self, context: dict) -> str:
        """
        Generate conversation summary
        """
        # TODO: Implement summary generation
        return ""