"""
Retrieval-Augmented Generation (RAG) system for Arabic travel agency chatbot.
Handles document indexing, vector search, and context retrieval.
"""

import os
import json
import logging
import pickle
import time
import re
from typing import List, Dict, Optional
import numpy as np
from pathlib import Path

try:
    import faiss
    FAISS_AVAILABLE = True
except ImportError:
    FAISS_AVAILABLE = False

logger = logging.getLogger(__name__)

class RAGSystem:
    """RAG system for document indexing and retrieval."""
    
    def __init__(self, gemini_client, data_dir='data', storage_dir='storage', chunk_size=500, overlap=50):
        self.gemini_client = gemini_client
        self.data_dir = Path(data_dir)
        self.storage_dir = Path(storage_dir)
        self.chunk_size = chunk_size
        self.overlap = overlap
        
        # Ensure directories exist
        self.data_dir.mkdir(exist_ok=True)
        self.storage_dir.mkdir(exist_ok=True)
        
        # Initialize index and metadata
        self.index = None
        self.metadata = []
        self.embedding_dim = 768  # Default Gemini embedding dimension
        
        # Conversation memory storage (session_id -> conversation_history)
        self.conversations = {}
        
        # Pattern mappings for incomplete messages
        self.message_patterns = {
            # Price inquiries
            'price_patterns': [
                r'بكم\s*سعر|كم\s*سعر|بكم\s*ثمن|كم\s*ثمن|بكم\s*تكلفة|كم\s*تكلفة',
                r'بكم|كم\s*(.*?)\s*\?*',
                r'السعر|الثمن|التكلفة'
            ],
            # Requirements inquiries  
            'requirements_patterns': [
                r'ماهي\s*متطلبات|ما\s*هي\s*متطلبات|متطلبات',
                r'ماهي\s*الشروط|ما\s*هي\s*الشروط|الشروط',
                r'ماذا\s*احتاج|ما\s*احتاج|احتاج',
                r'الوثائق|المستندات|الأوراق'
            ],
            # Offer/details inquiries
            'details_patterns': [
                r'ماهي\s*عروض|ما\s*هي\s*عروض|عروض',
                r'ماهي\s*تفاصيل|ما\s*هي\s*تفاصيل|تفاصيل',
                r'معلومات\s*عن|معلومات',
                r'اخبرني\s*عن|قل\s*لي\s*عن'
            ],
            # Booking inquiries
            'booking_patterns': [
                r'كيف\s*احجز|كيف\s*الحجز|احجز|حجز',
                r'كيف\s*اسجل|التسجيل|سجل'
            ]
        }
        
    def chunk_text(self, text: str, chunk_size: int = None, overlap: int = None) -> List[Dict]:
        """Split text into overlapping chunks."""
        if chunk_size is None:
            chunk_size = self.chunk_size
        if overlap is None:
            overlap = self.overlap
            
        # Simple word-based chunking
        words = text.split()
        chunks = []
        
        for i in range(0, len(words), chunk_size - overlap):
            chunk_words = words[i:i + chunk_size]
            chunk_text = ' '.join(chunk_words)
            
            chunks.append({
                'text': chunk_text,
                'start_word': i,
                'end_word': min(i + chunk_size, len(words)),
                'word_count': len(chunk_words)
            })
            
            # Break if we've reached the end
            if i + chunk_size >= len(words):
                break
                
        return chunks
    
    def load_documents(self) -> List[Dict]:
        """Load and chunk all text documents from data directory."""
        documents = []
        
        logger.info(f"Loading documents from {self.data_dir}")
        
        for file_path in self.data_dir.glob('*.txt'):
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read().strip()
                
                if not content:
                    logger.warning(f"Empty file: {file_path}")
                    continue
                
                # Chunk the document
                chunks = self.chunk_text(content)
                
                for i, chunk in enumerate(chunks):
                    documents.append({
                        'source_file': file_path.name,
                        'chunk_id': i,
                        'text': chunk['text'],
                        'start_word': chunk['start_word'],
                        'end_word': chunk['end_word'],
                        'word_count': chunk['word_count']
                    })
                
                logger.info(f"Loaded {len(chunks)} chunks from {file_path.name}")
                
            except Exception as e:
                logger.error(f"Error loading {file_path}: {e}")
                continue
        
        logger.info(f"Total documents loaded: {len(documents)}")
        return documents
    
    def create_embeddings(self, texts: List[str]) -> np.ndarray:
        """Create embeddings for list of texts using Gemini."""
        try:
            embeddings = self.gemini_client.embed_texts(texts)
            return np.array(embeddings)
        except Exception as e:
            logger.error(f"Error creating embeddings: {e}")
            # Fallback to random embeddings for development
            logger.warning("Using random embeddings as fallback")
            return np.random.rand(len(texts), self.embedding_dim).astype(np.float32)
    
    def build_index(self, embeddings: np.ndarray):
        """Build FAISS index from embeddings."""
        if not FAISS_AVAILABLE:
            logger.warning("FAISS not available, using simple similarity search")
            self.index = embeddings
            return
        
        # Create FAISS index
        dimension = embeddings.shape[1]
        self.index = faiss.IndexFlatIP(dimension)  # Inner product for cosine similarity
        
        # Normalize embeddings for cosine similarity
        faiss.normalize_L2(embeddings)
        self.index.add(embeddings)
        
        logger.info(f"Built FAISS index with {self.index.ntotal} vectors")
    
    def save_index(self):
        """Save index and metadata to disk."""
        try:
            # Save metadata
            metadata_path = self.storage_dir / 'metadata.json'
            with open(metadata_path, 'w', encoding='utf-8') as f:
                json.dump(self.metadata, f, ensure_ascii=False, indent=2)
            
            # Save index
            if FAISS_AVAILABLE and hasattr(self.index, 'ntotal'):
                index_path = self.storage_dir / 'index.faiss'
                faiss.write_index(self.index, str(index_path))
            else:
                # Save numpy array as pickle
                index_path = self.storage_dir / 'index.pkl'
                with open(index_path, 'wb') as f:
                    pickle.dump(self.index, f)
            
            logger.info("Index and metadata saved successfully")
            
        except Exception as e:
            logger.error(f"Error saving index: {e}")
            raise
    
    def load_index(self) -> bool:
        """Load existing index and metadata from disk."""
        try:
            metadata_path = self.storage_dir / 'metadata.json'
            
            # Check if files exist
            if not metadata_path.exists():
                logger.info("No existing index found")
                return False
            
            # Load metadata
            with open(metadata_path, 'r', encoding='utf-8') as f:
                self.metadata = json.load(f)
            
            # Load index
            if FAISS_AVAILABLE:
                index_path = self.storage_dir / 'index.faiss'
                if index_path.exists():
                    self.index = faiss.read_index(str(index_path))
                else:
                    return False
            else:
                index_path = self.storage_dir / 'index.pkl'
                if index_path.exists():
                    with open(index_path, 'rb') as f:
                        self.index = pickle.load(f)
                else:
                    return False
            
            logger.info(f"Loaded index with {len(self.metadata)} documents")
            return True
            
        except Exception as e:
            logger.error(f"Error loading index: {e}")
            return False
    
    def create_index(self, force: bool = False):
        """Create new index from documents in data directory."""
        if not force and self.load_index():
            logger.info("Index already exists and force=False")
            return
        
        logger.info("Creating new index...")
        
        # Load documents
        documents = self.load_documents()
        
        if not documents:
            logger.warning("No documents found to index")
            return
        
        # Extract texts for embedding
        texts = [doc['text'] for doc in documents]
        
        # Create embeddings
        logger.info("Creating embeddings...")
        embeddings = self.create_embeddings(texts)
        
        # Build index
        logger.info("Building search index...")
        self.build_index(embeddings)
        
        # Store metadata
        self.metadata = documents
        
        # Save to disk
        self.save_index()
        
        logger.info("Index creation completed")
    
    def simple_similarity_search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:
        """Simple similarity search when FAISS is not available."""
        if self.index is None:
            return []
        
        # Compute cosine similarity
        similarities = np.dot(self.index, query_embedding.T).flatten()
        
        # Get top k indices
        top_indices = np.argsort(similarities)[::-1][:k]
        return top_indices.tolist()
    
    def search(self, query_embedding: np.ndarray, k: int = 4) -> List[int]:
        """Search for similar documents."""
        if self.index is None:
            logger.warning("No index available for search")
            return []
        
        try:
            if FAISS_AVAILABLE and hasattr(self.index, 'search'):
                # FAISS search
                query_embedding = query_embedding.reshape(1, -1).astype(np.float32)
                faiss.normalize_L2(query_embedding)
                
                scores, indices = self.index.search(query_embedding, k)
                return indices[0].tolist()
            else:
                # Simple similarity search
                return self.simple_similarity_search(query_embedding, k)
                
        except Exception as e:
            logger.error(f"Error during search: {e}")
            return []
    
    def retrieve(self, query: str, k: int = 4) -> List[Dict]:
        """Retrieve relevant documents for a query."""
        try:
            # Create query embedding
            query_embeddings = self.create_embeddings([query])
            query_embedding = query_embeddings[0]
            
            # Search for similar documents
            indices = self.search(query_embedding, k)
            
            # Return relevant documents with metadata
            results = []
            for idx in indices:
                if 0 <= idx < len(self.metadata):
                    results.append(self.metadata[idx])
            
            logger.info(f"Retrieved {len(results)} documents for query")
            return results
            
        except Exception as e:
            logger.error(f"Error during retrieval: {e}")
            return []
    
    def process_message_with_context(self, message: str, session_id: str) -> str:
        """Process user message with conversation context and pattern recognition."""
        # Get conversation history
        history = self.get_conversation_history(session_id)
        
        # Check if message is incomplete/abbreviated
        enhanced_message = self._enhance_incomplete_message(message, history)
        
        logger.info(f"Original: '{message}' -> Enhanced: '{enhanced_message}'")
        return enhanced_message
    
    def _enhance_incomplete_message(self, message: str, history: List[Dict]) -> str:
        """Enhance incomplete messages using pattern recognition and context."""
        message_lower = message.lower().strip()
        
        # If message is complete (more than 10 characters and has context), return as-is
        if len(message.strip()) > 10 and any(word in message_lower for word in ['عن', 'في', 'إلى', 'من', 'مع']):
            return message
        
        # Get the last topic from conversation history
        last_topic = self._extract_last_topic(history)
        
        # Pattern matching for incomplete messages
        enhanced_message = message
        
        # Price inquiry patterns
        for pattern in self.message_patterns['price_patterns']:
            if re.search(pattern, message, re.IGNORECASE):
                if last_topic:
                    enhanced_message = f"ما هو سعر {last_topic}؟"
                else:
                    enhanced_message = "ما هي أسعار العروض المتاحة؟"
                break
        
        # Requirements inquiry patterns  
        for pattern in self.message_patterns['requirements_patterns']:
            if re.search(pattern, message, re.IGNORECASE):
                if last_topic:
                    enhanced_message = f"ما هي متطلبات {last_topic}؟"
                else:
                    enhanced_message = "ما هي متطلبات السفر والحجز؟"
                break
        
        # Details inquiry patterns
        for pattern in self.message_patterns['details_patterns']:
            if re.search(pattern, message, re.IGNORECASE):
                if last_topic:
                    enhanced_message = f"ما هي تفاصيل {last_topic}؟"
                else:
                    enhanced_message = "ما هي العروض والخدمات المتاحة؟"
                break
        
        # Booking inquiry patterns
        for pattern in self.message_patterns['booking_patterns']:
            if re.search(pattern, message, re.IGNORECASE):
                if last_topic:
                    enhanced_message = f"كيف يمكنني حجز {last_topic}؟"
                else:
                    enhanced_message = "كيف يمكنني حجز رحلة؟"
                break
        
        # Handle very short messages like "الخ..." or "وماذا أيضاً"
        if len(message.strip()) < 5 or message.strip() in ['الخ', 'الخ...', 'وماذا', 'وماذا أيضاً', 'ماذا أيضاً']:
            if last_topic:
                enhanced_message = f"أخبرني المزيد عن {last_topic}"
            else:
                enhanced_message = "أخبرني المزيد عن خدماتكم"
        
        return enhanced_message
    
    def _extract_last_topic(self, history: List[Dict]) -> Optional[str]:
        """Extract the main topic from recent conversation history."""
        if not history:
            return None
        
        # Look at the last few messages to find topics
        topics = []
        for entry in history[-3:]:  # Check last 3 exchanges
            user_msg = entry.get('user_message', '').lower()
            
            # Common travel topics to look for
            travel_keywords = {
                'تركيا': 'عروض تركيا',
                'دبي': 'عروض دبي', 
                'باريس': 'عروض باريس',
                'اليابان': 'عروض اليابان',
                'المالديف': 'عروض المالديف',
                'القاهرة': 'عروض القاهرة',
                'طوكيو': 'عروض طوكيو',
                'إسطنبول': 'عروض إسطنبول',
                'العمرة': 'رحلات العمرة',
                'الحج': 'رحلات الحج',
                'شهر العسل': 'رحلات شهر العسل'
            }
            
            for keyword, topic in travel_keywords.items():
                if keyword in user_msg:
                    topics.append(topic)
        
        # Return the most recent topic
        return topics[-1] if topics else None
    
    def add_to_conversation_history(self, session_id: str, user_message: str, assistant_response: str):
        """Add a conversation turn to history."""
        if session_id not in self.conversations:
            self.conversations[session_id] = []
        
        self.conversations[session_id].append({
            'user_message': user_message,
            'assistant_response': assistant_response,
            'timestamp': time.time()
        })
        
        # Keep only last 10 exchanges to manage memory
        if len(self.conversations[session_id]) > 10:
            self.conversations[session_id] = self.conversations[session_id][-10:]
    
    def get_conversation_history(self, session_id: str) -> List[Dict]:
        """Get conversation history for a session."""
        return self.conversations.get(session_id, [])
    
    def clear_conversation_history(self, session_id: str):
        """Clear conversation history for a session."""
        if session_id in self.conversations:
            del self.conversations[session_id]