from rag_pipeline import EnhancedLangChainRAG

rag = EnhancedLangChainRAG()

model = rag.initialize_embedding_model()

intent = rag.analyze_query_intent("Summarize the main points")
# Returns: 'summary'

from indexer import OptimizedLangChainIndexer

indexer = OptimizedLangChainIndexer()

index, processed_chunks = indexer.create_vectorstore_fast(
    chunks=["chunk1", "chunk2"],
    metadata=[{"source": "doc1"}, {"source": "doc2"}]
)

from model_loader import LLMManager

llm = LLMManager()

response = llm.generate_response(
    prompt="Explain quantum computing",
    max_tokens=500,
    temperature=0.7
)

from preprocess import load_and_split_texts

chunks, metadata = load_and_split_texts(strategy='hybrid')

from indexer import build_faiss_index

index, processed_chunks = build_faiss_index(chunks, metadata)

from rag_pipeline import run_rag_pipeline

response, sources = run_rag_pipeline(
    query="What are the main findings?",
    index=faiss_index,
    chunks=document_chunks,
    metadata=chunk_metadata
)

from preprocess import get_processing_stats

stats = get_processing_stats()
print(f"Processed {stats['files']} files with {stats['total_words']} words")

from indexer import get_index_info

info = get_index_info()
print(f"Index contains {info['chunk_count']} chunks")

# Base directories
BASE_DIR = Path(__file__).parent
DATA_DIR = BASE_DIR / "data"

# Data paths
PDF_DIR = str(DATA_DIR / "pdfs")
TEXT_DIR = str(DATA_DIR / "texts")
FAISS_INDEX_PATH = str(DATA_DIR / "faiss_index")
METADATA_PATH = str(DATA_DIR / "metadata.json")

# Model settings
EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"

# preprocess.py - Document loading and splitting
from langchain.document_loaders import PyMuPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

loader = PyMuPDFLoader(file_path)
documents = loader.load()

splitter = RecursiveCharacterTextSplitter(
    chunk_size=CHUNK_SIZE,
    chunk_overlap=CHUNK_OVERLAP,
    separators=["\n\n", "\n", ". ", " ", ""]
)
chunks = splitter.split_documents(documents)

# indexer.py - FAISS vector store
from langchain.vectorstores import FAISS
from langchain.embeddings import HuggingFaceEmbeddings

embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL_NAME)
vector_store = FAISS.from_documents(documents, embeddings)

# rag_pipeline.py - Complete RAG implementation
from langchain.chains import RetrievalQA
from langchain.prompts import PromptTemplate

# Create retriever
retriever = vector_store.as_retriever(search_kwargs={"k": TOP_K})

# Custom prompt template
prompt_template = PromptTemplate(
    template="Context: {context}\nQuestion: {question}\nAnswer:",
    input_variables=["context", "question"]
)

# Build QA chain
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    chain_type="stuff",
    retriever=retriever,
    chain_type_kwargs={"prompt": prompt_template}
)

def analyze_query_intent(query: str) -> str:
    """Classify query type for optimized processing"""
    intent_patterns = {
        'summary': ['summarize', 'overview', 'main points'],
        'analysis': ['analyze', 'evaluate', 'assessment'],
        'comparison': ['compare', 'contrast', 'versus'],
        'explanation': ['explain', 'how', 'why', 'what is']
    }
    
    query_lower = query.lower()
    for intent, patterns in intent_patterns.items():
        if any(pattern in query_lower for pattern in patterns):
            return intent
    return 'general'

def retrieve_chunks_advanced(query: str, index, strategy: str = "hybrid"):
    """Advanced retrieval with multiple strategies"""
    if strategy == "similarity":
        return similarity_search(query, index)
    elif strategy == "mmr":
        return mmr_search(query, index)
    elif strategy == "hybrid":
        # Combine multiple approaches
        sim_results = similarity_search(query, index)
        mmr_results = mmr_search(query, index)
        return merge_and_rank(sim_results, mmr_results)

# Create new loader in preprocess.py
from langchain.document_loaders.base import BaseLoader

class CustomDocumentLoader(BaseLoader):
    def __init__(self, file_path: str):
        self.file_path = file_path
    
    def load(self) -> List[Document]:
        # Implement custom loading logic
        pass

# Add to rag_pipeline.py
def create_custom_retriever(vector_store, strategy_config: Dict):
    """Create custom retrieval strategy"""
    if strategy_config["type"] == "custom":
        return CustomRetriever(
            vector_store=vector_store,
            **strategy_config["params"]
        )

# Add to app.py
def display_custom_component():
    """Custom Streamlit component"""
    with st.container():
        col1, col2 = st.columns(2)
        with col1:
            # Custom visualization
            pass
        with col2:
            # Interactive controls
            pass

# Streamlit caching
@st.cache_data
def load_and_process_documents():
    """Cache document processing results"""
    return process_documents()

@st.cache_resource
def initialize_models():
    """Cache model initialization"""
    return load_embedding_model()

import asyncio
from concurrent.futures import ThreadPoolExecutor

async def process_documents_async(documents: List[str]):
    """Async document processing"""
    loop = asyncio.get_event_loop()
    with ThreadPoolExecutor() as executor:
        tasks = [
            loop.run_in_executor(executor, process_single_doc, doc)
            for doc in documents
        ]
        return await asyncio.gather(*tasks)

def robust_processing_pipeline(documents):
    """Processing with fallback mechanisms"""
    try:
        # Try advanced LangChain processing
        return langchain_processing(documents)
    except ImportError:
        # Fallback to basic processing
        return basic_processing(documents)
    except Exception as e:
        # Log error and use minimal processing
        logger.error(f"Processing error: {e}")
        return minimal_processing(documents)

def handle_api_errors(func):
    """Decorator for API error handling"""
    def wrapper(*args, **kwargs):
        try:
            return func(*args, **kwargs)
        except ConnectionError:
            return "🌐 Connection error. Please check your internet connection."
        except AuthenticationError:
            return "🔑 API key invalid. Please check your configuration."
        except RateLimitError:
            return "⏰ Rate limit exceeded. Please wait and try again."
    return wrapper

# Check environment loading
from dotenv import load_dotenv
load_dotenv()
print(os.getenv("OPENROUTER_API_KEY"))

# Verify .env file format
OPENROUTER_API_KEY=your_key_here  # No quotes, no spaces

# Check data directory structure
import os
print(os.listdir("data/"))

# Reset and rebuild index
python -c "
from indexer import get_index_info
print(get_index_info())
"

# Reduce chunk size for faster processing
# In .env file:
LANGCHAIN_CHUNK_SIZE=400
CHUNK_SIZE=300
TOP_K=3

# Or modify config.py temporarily:
CHUNK_SIZE = 300
CHUNK_OVERLAP = 30

# Monitor memory usage
import psutil
print(f"Memory usage: {psutil.virtual_memory().percent}%")

# Optimize settings in config.py:
DEVICE = "cpu"  # Ensure CPU-only processing
LANGCHAIN_ENABLE_CACHING = False  # Disable caching if needed

# Check file size limits
# In .streamlit/config.toml:
[server]
maxUploadSize = 200

# Verify file permissions
ls -la data/pdfs/
chmod 755 data/

# Check session state
import streamlit as st
print(st.session_state)

# Clear and reinitialize
if st.button("Reset Session"):
    for key in st.session_state.keys():
        del st.session_state[key]
    st.rerun()

# Test API connection
from model_loader import test_openrouter_connection
success, message = test_openrouter_connection()
print(f"API Status: {success}, Message: {message}")

# Check API key validity
import requests
headers = {"Authorization": f"Bearer {OPENROUTER_API_KEY}"}
response = requests.get("https://openrouter.ai/api/v1/models", headers=headers)
print(f"Status: {response.status_code}")

# Implement retry logic
import time
from functools import wraps

def retry_api_call(max_retries=3, delay=1):
    def decorator(func):
        @wraps(func)
        def wrapper(*args, **kwargs):
            for attempt in range(max_retries):
                try:
                    return func(*args, **kwargs)
                except Exception as e:
                    if "rate limit" in str(e).lower() and attempt < max_retries - 1:
                        time.sleep(delay * (2 ** attempt))  # Exponential backoff
                        continue
                    raise e
            return None
        return wrapper
    return decorator

# Add to config.py
DEBUG = True
LANGCHAIN_VERBOSE = True
LANGCHAIN_DEBUG = True

# Enable detailed logging
import logging
logging.basicConfig(level=logging.DEBUG)

# Install memory profiler
pip install memory-profiler

# Profile specific functions
from memory_profiler import profile

@profile
def process_large_document():
    # Your processing code
    pass

# Add timing decorators
import time
from functools import wraps

def timer(func):
    @wraps(func)
    def wrapper(*args, **kwargs):
        start = time.time()
        result = func(*args, **kwargs)
        end = time.time()
        print(f"{func.__name__} took {end - start:.2f} seconds")
        return result
    return wrapper

@timer
def slow_function():
    # Function to profile
    pass

# Implement graceful fallbacks
def robust_rag_pipeline(query, fallback_enabled=True):
    try:
        return advanced_rag_pipeline(query)
    except Exception as e:
        if fallback_enabled:
            print(f"Advanced pipeline failed: {e}")
            return simple_rag_pipeline(query)
        raise e

def simple_rag_pipeline(query):
    """Simplified RAG without advanced features"""
    # Basic similarity search without LangChain
    pass

# Backup and restore functions
def backup_index():
    """Backup FAISS index and metadata"""
    import shutil
    from datetime import datetime
    
    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
    backup_dir = f"backups/backup_{timestamp}"
    os.makedirs(backup_dir, exist_ok=True)
    
    shutil.copytree("data/", f"{backup_dir}/data/")
    print(f"Backup created: {backup_dir}")

def restore_index(backup_path):
    """Restore from backup"""
    import shutil
    if os.path.exists(backup_path):
        shutil.copytree(f"{backup_path}/data/", "data/", dirs_exist_ok=True)
        print("Restore completed")

# Use Black formatter with 88 character line length
black . --line-length 88

# Import sorting with isort
isort . --profile black

# Type hints for all functions
def process_document(file_path: str, chunk_size: int = 500) -> List[str]:
    """Process document and return chunks.
    
    Args:
        file_path: Path to the document file
        chunk_size: Size of text chunks
        
    Returns:
        List of text chunks
        
    Raises:
        FileNotFoundError: If file doesn't exist
        ValueError: If chunk_size is invalid
    """
    pass

class DocumentProcessor:
    """Advanced document processing with LangChain integration.
    
    This class handles document loading, text extraction, and chunking
    using various LangChain strategies for optimal performance.
    
    Attributes:
        strategy: Text splitting strategy ('recursive', 'semantic', etc.)
        chunk_size: Maximum size of text chunks
        
    Example:
        >>> processor = DocumentProcessor(strategy='hybrid')
        >>> chunks = processor.process_pdf('document.pdf')
        >>> print(f"Created {len(chunks)} chunks")
    """
    
    def __init__(self, strategy: str = 'hybrid', chunk_size: int = 500):
        """Initialize processor with specified strategy."""
        pass

# tests/test_rag_pipeline.py
import pytest
from unittest.mock import Mock, patch
from rag_pipeline import EnhancedLangChainRAG

class TestEnhancedLangChainRAG:
    def setup_method(self):
        """Setup test fixtures."""
        self.rag = EnhancedLangChainRAG()
    
    def test_analyze_query_intent_summary(self):
        """Test query intent analysis for summary queries."""
        query = "Summarize the main points of this document"
        intent = self.rag.analyze_query_intent(query)
        assert intent == "summary"
    
    @patch('rag_pipeline.generate_response')
    def test_run_rag_pipeline_success(self, mock_generate):
        """Test successful RAG pipeline execution."""
        mock_generate.return_value = "Test response"
        
        # Mock dependencies
        mock_index = Mock()
        chunks = ["chunk1", "chunk2"]
        metadata = [{"source": "doc1"}, {"source": "doc2"}]
        
        response, chunk_info = self.rag.run_rag_pipeline(
            "test query", mock_index, chunks, metadata
        )
        
        assert response == "Test response"
        assert len(chunk_info) > 0

# tests/test_integration.py
import tempfile
import os
from app import main
from indexer import build_faiss_index
from preprocess import load_and_split_texts

class TestIntegration:
    def test_end_to_end_pipeline(self):
        """Test complete document processing pipeline."""
        with tempfile.TemporaryDirectory() as temp_dir:
            # Create test PDF
            test_pdf_path = os.path.join(temp_dir, "test.pdf")
            self.create_test_pdf(test_pdf_path)
            
            # Process document
            chunks, metadata = load_and_split_texts()
            assert len(chunks) > 0
            
            # Build index
            index, processed_chunks = build_faiss_index(chunks, metadata)
            assert index is not None
            
            # Test query
            from rag_pipeline import run_rag_pipeline
            response, sources = run_rag_pipeline(
                "What is this document about?", 
                index, chunks, metadata
            )
            assert len(response) > 0

# Use Google-style docstrings
def retrieve_chunks_advanced(
    query: str, 
    index: Any, 
    chunks: List[str], 
    metadata: Optional[List[Dict]] = None,
    top_k: int = 5
) -> Tuple[List[str], List[Dict]]:
    """Retrieve document chunks using advanced strategies.
    
    Performs similarity search with multiple retrieval strategies
    including semantic search, MMR, and hybrid approaches.
    
    Args:
        query: User query string for semantic search
        index: FAISS index for similarity computation
        chunks: List of document text chunks
        metadata: Optional metadata for each chunk
        top_k: Number of top chunks to retrieve
        
    Returns:
        Tuple containing:
            - List of retrieved text chunks
            - List of chunk metadata with relevance scores
            
    Raises:
        ValueError: If query is empty or index is invalid
        IndexError: If top_k exceeds available chunks
        
    Example:
        >>> chunks, info = retrieve_chunks_advanced(
        ...     "What is machine learning?",
        ...     faiss_index,
        ...     document_chunks,
        ...     chunk_metadata,
        ...     top_k=5
        ... )
        >>> print(f"Retrieved {len(chunks)} relevant chunks")
    """