session_manager.py

"""
Session Management System for Law RAG Chatbot
"""

import uuid
import time
import json
from typing import Dict, List, Any, Optional
from datetime import datetime, timedelta
from pathlib import Path
import sqlite3
import logging

logger = logging.getLogger(__name__)

class SessionManager:
    """Manages user sessions and chat history"""
    
    def __init__(self, db_path: str = "chat_sessions.db"):
        self.db_path = db_path
        self.sessions: Dict[str, Dict[str, Any]] = {}
        self._init_database()
    
    def _init_database(self):
        """Initialize SQLite database for session storage"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Create sessions table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS sessions (
                    session_id TEXT PRIMARY KEY,
                    created_at TIMESTAMP,
                    last_activity TIMESTAMP,
                    user_info TEXT,
                    metadata TEXT
                )
            ''')
            
            # Create chat_history table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS chat_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    question TEXT,
                    answer TEXT,
                    sources TEXT,
                    confidence REAL,
                    processing_time REAL,
                    timestamp TIMESTAMP,
                    FOREIGN KEY (session_id) REFERENCES sessions (session_id)
                )
            ''')
            
            # Create search_history table
            cursor.execute('''
                CREATE TABLE IF NOT EXISTS search_history (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id TEXT,
                    query TEXT,
                    results_count INTEGER,
                    timestamp TIMESTAMP,
                    FOREIGN KEY (session_id) REFERENCES sessions (session_id)
                )
            ''')
            
            conn.commit()
            conn.close()
            logger.info("Session database initialized successfully")
            
        except Exception as e:
            logger.error(f"Failed to initialize session database: {e}")
            raise
    
    def create_session(self, user_info: Optional[str] = None, metadata: Optional[Dict] = None) -> str:
        """Create a new session"""
        try:
            session_id = str(uuid.uuid4())
            current_time = datetime.now()
            
            # Store in memory
            self.sessions[session_id] = {
                "session_id": session_id,
                "created_at": current_time,
                "last_activity": current_time,
                "user_info": user_info or "anonymous",
                "metadata": metadata or {},
                "chat_count": 0,
                "search_count": 0
            }
            
            # Store in database
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO sessions (session_id, created_at, last_activity, user_info, metadata)
                VALUES (?, ?, ?, ?, ?)
            ''', (
                session_id,
                current_time.isoformat(),
                current_time.isoformat(),
                user_info or "anonymous",
                json.dumps(metadata or {})
            ))
            
            conn.commit()
            conn.close()
            
            logger.info(f"Created new session: {session_id}")
            return session_id
            
        except Exception as e:
            logger.error(f"Failed to create session: {e}")
            raise
    
    def get_session(self, session_id: str) -> Optional[Dict[str, Any]]:
        """Get session information"""
        # Check memory first
        if session_id in self.sessions:
            return self.sessions[session_id]
        
        # Check database
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT session_id, created_at, last_activity, user_info, metadata
                FROM sessions WHERE session_id = ?
            ''', (session_id,))
            
            row = cursor.fetchone()
            conn.close()
            
            if row:
                session_data = {
                    "session_id": row[0],
                    "created_at": datetime.fromisoformat(row[1]),
                    "last_activity": datetime.fromisoformat(row[2]),
                    "user_info": row[3],
                    "metadata": json.loads(row[4]) if row[4] else {},
                    "chat_count": 0,
                    "search_count": 0
                }
                
                # Load counts
                session_data["chat_count"] = self._get_chat_count(session_id)
                session_data["search_count"] = self._get_search_count(session_id)
                
                # Store in memory
                self.sessions[session_id] = session_data
                return session_data
            
            return None
            
        except Exception as e:
            logger.error(f"Failed to get session {session_id}: {e}")
            return None
    
    def update_session_activity(self, session_id: str):
        """Update session last activity"""
        current_time = datetime.now()
        
        # Update memory
        if session_id in self.sessions:
            self.sessions[session_id]["last_activity"] = current_time
        
        # Update database
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                UPDATE sessions SET last_activity = ? WHERE session_id = ?
            ''', (current_time.isoformat(), session_id))
            
            conn.commit()
            conn.close()
            
        except Exception as e:
            logger.error(f"Failed to update session activity: {e}")
    
    def store_chat_response(self, session_id: str, question: str, answer: str, 
                           sources: List[Dict], confidence: float, processing_time: float):
        """Store a chat response in the session"""
        try:
            current_time = datetime.now()
            
            # Update session activity
            self.update_session_activity(session_id)
            
            # Store in database
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO chat_history 
                (session_id, question, answer, sources, confidence, processing_time, timestamp)
                VALUES (?, ?, ?, ?, ?, ?, ?)
            ''', (
                session_id,
                question,
                answer,
                json.dumps(sources),
                confidence,
                processing_time,
                current_time.isoformat()
            ))
            
            conn.commit()
            conn.close()
            
            # Update memory
            if session_id in self.sessions:
                self.sessions[session_id]["chat_count"] += 1
            
            logger.info(f"Stored chat response for session {session_id}")
            
        except Exception as e:
            logger.error(f"Failed to store chat response: {e}")
    
    def store_search_query(self, session_id: str, query: str, results_count: int):
        """Store a search query in the session"""
        try:
            current_time = datetime.now()
            
            # Update session activity
            self.update_session_activity(session_id)
            
            # Store in database
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                INSERT INTO search_history 
                (session_id, query, results_count, timestamp)
                VALUES (?, ?, ?, ?)
            ''', (
                session_id,
                query,
                results_count,
                current_time.isoformat()
            ))
            
            conn.commit()
            conn.close()
            
            # Update memory
            if session_id in self.sessions:
                self.sessions[session_id]["search_count"] += 1
            
            logger.info(f"Stored search query for session {session_id}")
            
        except Exception as e:
            logger.error(f"Failed to store search query: {e}")
    
    def get_chat_history(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Get chat history for a session"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT question, answer, sources, confidence, processing_time, timestamp
                FROM chat_history 
                WHERE session_id = ? 
                ORDER BY timestamp DESC 
                LIMIT ?
            ''', (session_id, limit))
            
            rows = cursor.fetchall()
            conn.close()
            
            history = []
            for row in rows:
                history.append({
                    "question": row[0],
                    "answer": row[1],
                    "sources": json.loads(row[2]) if row[2] else [],
                    "confidence": row[3],
                    "processing_time": row[4],
                    "timestamp": row[5]
                })
            
            return history
            
        except Exception as e:
            logger.error(f"Failed to get chat history: {e}")
            return []
    
    def get_search_history(self, session_id: str, limit: int = 10) -> List[Dict[str, Any]]:
        """Get search history for a session"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('''
                SELECT query, results_count, timestamp
                FROM search_history 
                WHERE session_id = ? 
                ORDER BY timestamp DESC 
                LIMIT ?
            ''', (session_id, limit))
            
            rows = cursor.fetchall()
            conn.close()
            
            history = []
            for row in rows:
                history.append({
                    "query": row[0],
                    "results_count": row[1],
                    "timestamp": row[2]
                })
            
            return history
            
        except Exception as e:
            logger.error(f"Failed to get search history: {e}")
            return []
    
    def _get_chat_count(self, session_id: str) -> int:
        """Get chat count for a session"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('SELECT COUNT(*) FROM chat_history WHERE session_id = ?', (session_id,))
            count = cursor.fetchone()[0]
            
            conn.close()
            return count
            
        except Exception as e:
            logger.error(f"Failed to get chat count: {e}")
            return 0
    
    def _get_search_count(self, session_id: str) -> int:
        """Get search count for a session"""
        try:
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('SELECT COUNT(*) FROM search_history WHERE session_id = ?', (session_id,))
            count = cursor.fetchone()[0]
            
            conn.close()
            return count
            
        except Exception as e:
            logger.error(f"Failed to get search count: {e}")
            return 0
    
    def get_session_stats(self, session_id: str) -> Dict[str, Any]:
        """Get comprehensive session statistics"""
        session = self.get_session(session_id)
        if not session:
            return {}
        
        chat_history = self.get_chat_history(session_id, limit=100)
        search_history = self.get_search_history(session_id, limit=100)
        
        # Calculate average confidence
        confidences = [chat["confidence"] for chat in chat_history if chat["confidence"] > 0]
        avg_confidence = sum(confidences) / len(confidences) if confidences else 0
        
        # Calculate average processing time
        processing_times = [chat["processing_time"] for chat in chat_history if chat["processing_time"] > 0]
        avg_processing_time = sum(processing_times) / len(processing_times) if processing_times else 0
        
        return {
            "session_id": session_id,
            "created_at": session["created_at"].isoformat(),
            "last_activity": session["last_activity"].isoformat(),
            "user_info": session["user_info"],
            "total_chats": len(chat_history),
            "total_searches": len(search_history),
            "average_confidence": round(avg_confidence, 3),
            "average_processing_time": round(avg_processing_time, 3),
            "recent_questions": [chat["question"] for chat in chat_history[:5]],
            "recent_searches": [search["query"] for search in search_history[:5]]
        }
    
    def cleanup_old_sessions(self, days: int = 30):
        """Clean up sessions older than specified days"""
        try:
            cutoff_date = datetime.now() - timedelta(days=days)
            
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            # Delete old chat history
            cursor.execute('''
                DELETE FROM chat_history 
                WHERE session_id IN (
                    SELECT session_id FROM sessions 
                    WHERE last_activity < ?
                )
            ''', (cutoff_date.isoformat(),))
            
            # Delete old search history
            cursor.execute('''
                DELETE FROM search_history 
                WHERE session_id IN (
                    SELECT session_id FROM sessions 
                    WHERE last_activity < ?
                )
            ''', (cutoff_date.isoformat(),))
            
            # Delete old sessions
            cursor.execute('DELETE FROM sessions WHERE last_activity < ?', (cutoff_date.isoformat(),))
            
            conn.commit()
            conn.close()
            
            logger.info(f"Cleaned up sessions older than {days} days")
            
        except Exception as e:
            logger.error(f"Failed to cleanup old sessions: {e}")
    
    def delete_session(self, session_id: str):
        """Delete a session and all its data"""
        try:
            # Remove from memory
            if session_id in self.sessions:
                del self.sessions[session_id]
            
            # Remove from database
            conn = sqlite3.connect(self.db_path)
            cursor = conn.cursor()
            
            cursor.execute('DELETE FROM chat_history WHERE session_id = ?', (session_id,))
            cursor.execute('DELETE FROM search_history WHERE session_id = ?', (session_id,))
            cursor.execute('DELETE FROM sessions WHERE session_id = ?', (session_id,))
            
            conn.commit()
            conn.close()
            
            logger.info(f"Deleted session: {session_id}")
            
        except Exception as e:
            logger.error(f"Failed to delete session {session_id}: {e}")