"""
Advanced Retrieval-Augmented Generation (RAG) System for Toxo.

This module implements state-of-the-art RAG patterns including:
- Multi-vector retrieval strategies
- Query decomposition and rewriting  
- Retrieval fusion and ranking
- Contextual compression and filtering
- Self-correcting retrieval loops
- Hierarchical retrieval
- Temporal and knowledge graph-aware retrieval
"""

import asyncio
import re
import json
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import Dict, List, Optional, Any, Union, Tuple, Callable
from enum import Enum
from datetime import datetime, timedelta
import logging
from pathlib import Path

import numpy as np
from sklearn.metrics.pairwise import cosine_similarity
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.cluster import KMeans

from ..core.config import ToxoConfig
from ..core.vector_store import VectorStore, VectorSearchResult
from ..integrations.gemini_client import GeminiClient
from ..utils.logger import get_logger
from ..utils.exceptions import ToxoError


class RetrievalStrategy(Enum):
    """Different retrieval strategies."""
    SEMANTIC = "semantic"
    HYBRID = "hybrid"  
    MULTI_VECTOR = "multi_vector"
    HIERARCHICAL = "hierarchical"
    CONTEXTUAL = "contextual"
    TEMPORAL = "temporal"
    GRAPH_AWARE = "graph_aware"


class FusionMethod(Enum):
    """Methods for fusing multiple retrieval results."""
    RRF = "reciprocal_rank_fusion"  # Reciprocal Rank Fusion
    WEIGHTED = "weighted_combination"
    VOTE = "vote_based"
    LEARNED = "learned_fusion"


@dataclass
class RetrievalQuery:
    """Enhanced query structure for retrieval."""
    original_query: str
    rewritten_queries: List[str] = field(default_factory=list)
    query_intent: Optional[str] = None
    query_type: Optional[str] = None  # factual, analytical, creative, etc.
    context: Dict[str, Any] = field(default_factory=dict)
    filters: Dict[str, Any] = field(default_factory=dict)
    temporal_constraints: Optional[Dict[str, Any]] = None
    user_profile: Optional[Dict[str, Any]] = None


@dataclass
class RetrievalResult:
    """Enhanced retrieval result with additional metadata."""
    content: str
    score: float
    metadata: Dict[str, Any]
    source: str
    retrieval_method: str
    confidence: float = 1.0
    relevance_explanation: Optional[str] = None
    chunk_position: Optional[int] = None
    document_hierarchy: Optional[List[str]] = None


@dataclass
class RAGResponse:
    """Complete RAG response with provenance."""
    answer: str
    retrieved_contexts: List[RetrievalResult]
    confidence_score: float
    retrieval_query: RetrievalQuery
    reasoning_trace: List[str] = field(default_factory=list)
    citations: List[Dict[str, Any]] = field(default_factory=list)
    alternative_answers: List[str] = field(default_factory=list)


class QueryProcessor:
    """Advanced query processing and rewriting."""
    
    def __init__(self, gemini_client: GeminiClient):
        self.gemini_client = gemini_client
        self.logger = get_logger(f"{__name__}.query_processor")
    
    async def process_query(self, query: str, context: Optional[Dict[str, Any]] = None) -> RetrievalQuery:
        """Process and enhance the input query."""
        # Analyze query intent and type
        intent_analysis = await self._analyze_query_intent(query)
        
        # Generate query variations
        rewritten_queries = await self._generate_query_variations(query, intent_analysis)
        
        # Extract temporal constraints
        temporal_constraints = self._extract_temporal_constraints(query)
        
        # Create enhanced query object
        retrieval_query = RetrievalQuery(
            original_query=query,
            rewritten_queries=rewritten_queries,
            query_intent=intent_analysis.get('intent'),
            query_type=intent_analysis.get('type'),
            context=context or {},
            temporal_constraints=temporal_constraints
        )
        
        return retrieval_query
    
    async def _analyze_query_intent(self, query: str) -> Dict[str, Any]:
        """Analyze query intent and type using LLM."""
        prompt = f"""
        Analyze the following query and determine:
        1. Intent (informational, navigational, transactional, analytical)
        2. Type (factual, opinion, comparison, explanation, creative)
        3. Complexity (simple, medium, complex)
        4. Domain (if specific)
        5. Expected answer type (short, detailed, list, etc.)
        
        Query: "{query}"
        
        Return as JSON with keys: intent, type, complexity, domain, answer_type
        """
        
        try:
            response = await self.gemini_client.generate(prompt, temperature=0.1)
            return json.loads(response)
        except Exception as e:
            self.logger.warning(f"Failed to analyze query intent: {e}")
            return {"intent": "informational", "type": "factual", "complexity": "medium"}
    
    async def _generate_query_variations(self, query: str, intent_analysis: Dict[str, Any]) -> List[str]:
        """Generate multiple variations of the query for diverse retrieval."""
        prompt = f"""
        Generate 3-5 different variations of the following query to improve information retrieval.
        Each variation should:
        - Maintain the original intent
        - Use different keywords and phrasings
        - Focus on different aspects of the question
        - Include synonyms and related terms
        
        Original query: "{query}"
        Query type: {intent_analysis.get('type', 'factual')}
        
        Return only the variations, one per line.
        """
        
        try:
            response = await self.gemini_client.generate(prompt, temperature=0.7)
            variations = [line.strip() for line in response.strip().split('\n') if line.strip()]
            return variations[:5]  # Limit to 5 variations
        except Exception as e:
            self.logger.warning(f"Failed to generate query variations: {e}")
            return [query]  # Fallback to original query
    
    def _extract_temporal_constraints(self, query: str) -> Optional[Dict[str, Any]]:
        """Extract temporal constraints from the query."""
        temporal_patterns = {
            'recent': r'\b(recent|latest|new|current|today|this week|this month)\b',
            'past': r'\b(last year|previous|past|before|earlier|ago)\b',
            'future': r'\b(future|upcoming|next|will be|expected)\b',
            'specific_date': r'\b(\d{4}|\d{1,2}/\d{1,2}/\d{2,4})\b'
        }
        
        constraints = {}
        for constraint_type, pattern in temporal_patterns.items():
            if re.search(pattern, query, re.IGNORECASE):
                constraints[constraint_type] = True
        
        return constraints if constraints else None


class AdvancedRetriever:
    """Advanced retrieval system with multiple strategies."""
    
    def __init__(
        self,
        vector_store: VectorStore,
        gemini_client: GeminiClient,
        config: ToxoConfig
    ):
        self.vector_store = vector_store
        self.gemini_client = gemini_client
        self.config = config
        self.logger = get_logger(f"{__name__}.retriever")
        
        # Initialize components
        self.query_processor = QueryProcessor(gemini_client)
        self.result_ranker = ResultRanker(gemini_client)
        self.context_compressor = ContextCompressor(gemini_client)
        
        # Retrieval strategies
        self.strategies = {
            RetrievalStrategy.SEMANTIC: self._semantic_retrieval,
            RetrievalStrategy.HYBRID: self._hybrid_retrieval,
            RetrievalStrategy.MULTI_VECTOR: self._multi_vector_retrieval,
            RetrievalStrategy.HIERARCHICAL: self._hierarchical_retrieval,
            RetrievalStrategy.CONTEXTUAL: self._contextual_retrieval
        }
    
    async def retrieve(
        self,
        query: str,
        strategy: RetrievalStrategy = RetrievalStrategy.HYBRID,
        top_k: int = 10,
        context: Optional[Dict[str, Any]] = None
    ) -> List[RetrievalResult]:
        """Main retrieval method."""
        # Process query
        retrieval_query = await self.query_processor.process_query(query, context)
        
        # Execute retrieval strategy
        retrieval_method = self.strategies.get(strategy, self._hybrid_retrieval)
        results = await retrieval_method(retrieval_query, top_k)
        
        # Re-rank results
        ranked_results = await self.result_ranker.rerank(results, retrieval_query)
        
        # Filter and validate results
        filtered_results = await self._filter_results(ranked_results, retrieval_query)
        
        return filtered_results[:top_k]
    
    async def _semantic_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Pure semantic vector retrieval."""
        # Generate embedding for original query
        query_embedding = await self.gemini_client.embed_text(query.original_query)
        
        # Search vector store
        search_results = await self.vector_store.search_similar(
            np.array(query_embedding),
            k=top_k,
            filter_metadata=query.filters
        )
        
        # Convert to RetrievalResult objects
        results = []
        for result in search_results:
            results.append(RetrievalResult(
                content=result.content,
                score=result.similarity_score,
                metadata=result.metadata,
                source=result.id,
                retrieval_method="semantic",
                confidence=result.similarity_score
            ))
        
        return results
    
    async def _hybrid_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Hybrid semantic + keyword retrieval with fusion."""
        # Semantic retrieval
        semantic_results = await self._semantic_retrieval(query, top_k * 2)
        
        # Keyword-based retrieval (simplified TF-IDF approach)
        keyword_results = await self._keyword_retrieval(query, top_k * 2)
        
        # Fuse results using RRF
        fused_results = await self._fuse_results(
            [semantic_results, keyword_results],
            method=FusionMethod.RRF
        )
        
        return fused_results[:top_k]
    
    async def _multi_vector_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Multi-vector retrieval using query variations."""
        all_results = []
        
        # Retrieve for original query
        original_results = await self._semantic_retrieval(query, top_k)
        all_results.append(original_results)
        
        # Retrieve for each query variation
        for variation in query.rewritten_queries:
            variation_query = RetrievalQuery(
                original_query=variation,
                context=query.context,
                filters=query.filters
            )
            variation_results = await self._semantic_retrieval(variation_query, top_k // 2)
            all_results.append(variation_results)
        
        # Fuse all results
        fused_results = await self._fuse_results(all_results, FusionMethod.RRF)
        return fused_results[:top_k]
    
    async def _hierarchical_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Hierarchical retrieval from document structure."""
        # First, retrieve broad document sections
        broad_results = await self._semantic_retrieval(query, top_k // 2)
        
        # Then, retrieve specific chunks within those documents
        detailed_results = []
        for result in broad_results:
            # Get document ID from metadata
            doc_id = result.metadata.get('document_id')
            if doc_id:
                # Search for chunks within this document
                doc_filter = {'document_id': doc_id}
                query_with_filter = RetrievalQuery(
                    original_query=query.original_query,
                    filters={**query.filters, **doc_filter}
                )
                chunk_results = await self._semantic_retrieval(query_with_filter, 5)
                detailed_results.extend(chunk_results)
        
        # Combine and deduplicate
        all_results = broad_results + detailed_results
        unique_results = self._deduplicate_results(all_results)
        
        return unique_results[:top_k]
    
    async def _contextual_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Context-aware retrieval considering user profile and session."""
        # Enhance query with context
        if query.user_profile or query.context:
            enhanced_query = await self._enhance_query_with_context(query)
            enhanced_query_obj = RetrievalQuery(
                original_query=enhanced_query,
                context=query.context,
                filters=query.filters
            )
            results = await self._semantic_retrieval(enhanced_query_obj, top_k)
        else:
            results = await self._semantic_retrieval(query, top_k)
        
        # Filter based on user preferences and context
        filtered_results = await self._apply_contextual_filters(results, query)
        
        return filtered_results
    
    async def _keyword_retrieval(
        self, 
        query: RetrievalQuery, 
        top_k: int
    ) -> List[RetrievalResult]:
        """Keyword-based retrieval using TF-IDF."""
        # This is a simplified implementation
        # In practice, you'd use a proper search engine like Elasticsearch
        
        # For now, return semantic results with modified scoring
        results = await self._semantic_retrieval(query, top_k)
        
        # Boost results that contain query keywords
        query_keywords = set(query.original_query.lower().split())
        for result in results:
            content_words = set(result.content.lower().split())
            keyword_overlap = len(query_keywords.intersection(content_words))
            keyword_boost = keyword_overlap / len(query_keywords) if query_keywords else 0
            result.score = result.score * 0.7 + keyword_boost * 0.3
            result.retrieval_method = "keyword"
        
        return sorted(results, key=lambda x: x.score, reverse=True)
    
    async def _fuse_results(
        self, 
        result_lists: List[List[RetrievalResult]],
        method: FusionMethod = FusionMethod.RRF
    ) -> List[RetrievalResult]:
        """Fuse multiple result lists."""
        if method == FusionMethod.RRF:
            return await self._reciprocal_rank_fusion(result_lists)
        elif method == FusionMethod.WEIGHTED:
            return await self._weighted_fusion(result_lists)
        else:
            # Simple concatenation fallback
            all_results = []
            for results in result_lists:
                all_results.extend(results)
            return self._deduplicate_results(all_results)
    
    async def _reciprocal_rank_fusion(
        self, 
        result_lists: List[List[RetrievalResult]]
    ) -> List[RetrievalResult]:
        """Implement Reciprocal Rank Fusion."""
        rrf_scores = {}
        k = 60  # RRF parameter
        
        for results in result_lists:
            for rank, result in enumerate(results, 1):
                key = result.source  # Use source as unique identifier
                if key not in rrf_scores:
                    rrf_scores[key] = {'result': result, 'score': 0}
                rrf_scores[key]['score'] += 1.0 / (k + rank)
        
        # Sort by RRF score
        sorted_results = sorted(
            rrf_scores.values(),
            key=lambda x: x['score'],
            reverse=True
        )
        
        # Update scores and return results
        final_results = []
        for item in sorted_results:
            result = item['result']
            result.score = item['score']
            result.retrieval_method = "rrf_fused"
            final_results.append(result)
        
        return final_results
    
    async def _weighted_fusion(
        self, 
        result_lists: List[List[RetrievalResult]],
        weights: Optional[List[float]] = None
    ) -> List[RetrievalResult]:
        """Weighted combination of result lists."""
        if weights is None:
            weights = [1.0] * len(result_lists)
        
        weighted_scores = {}
        
        for i, results in enumerate(result_lists):
            weight = weights[i]
            for result in results:
                key = result.source
                if key not in weighted_scores:
                    weighted_scores[key] = {'result': result, 'score': 0}
                weighted_scores[key]['score'] += result.score * weight
        
        # Sort by weighted score
        sorted_results = sorted(
            weighted_scores.values(),
            key=lambda x: x['score'],
            reverse=True
        )
        
        final_results = []
        for item in sorted_results:
            result = item['result']
            result.score = item['score']
            result.retrieval_method = "weighted_fused"
            final_results.append(result)
        
        return final_results
    
    def _deduplicate_results(self, results: List[RetrievalResult]) -> List[RetrievalResult]:
        """Remove duplicate results based on source."""
        seen_sources = set()
        unique_results = []
        
        for result in results:
            if result.source not in seen_sources:
                seen_sources.add(result.source)
                unique_results.append(result)
        
        return unique_results
    
    async def _filter_results(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery
    ) -> List[RetrievalResult]:
        """Filter results based on relevance and quality."""
        filtered_results = []
        
        for result in results:
            # Basic relevance threshold
            if result.score < 0.1:
                continue
            
            # Content quality check
            if len(result.content.strip()) < 10:
                continue
            
            # Temporal filtering if constraints exist
            if query.temporal_constraints:
                if not await self._matches_temporal_constraints(result, query.temporal_constraints):
                    continue
            
            filtered_results.append(result)
        
        return filtered_results
    
    async def _matches_temporal_constraints(
        self, 
        result: RetrievalResult, 
        constraints: Dict[str, Any]
    ) -> bool:
        """Check if result matches temporal constraints."""
        # Check document date in metadata
        doc_date = result.metadata.get('date') or result.metadata.get('created_at')
        if not doc_date:
            return True  # No date info, allow through
        
        try:
            if isinstance(doc_date, str):
                doc_date = datetime.fromisoformat(doc_date.replace('Z', '+00:00'))
            
            now = datetime.now()
            
            if constraints.get('recent'):
                # Recent means last 30 days
                if (now - doc_date).days > 30:
                    return False
            
            if constraints.get('past'):
                # Past means older than 1 year
                if (now - doc_date).days < 365:
                    return False
            
            return True
            
        except Exception:
            return True  # If date parsing fails, allow through
    
    async def _enhance_query_with_context(self, query: RetrievalQuery) -> str:
        """Enhance query with user context and profile."""
        context_parts = []
        
        if query.user_profile:
            context_parts.append(f"User context: {query.user_profile}")
        
        if query.context:
            context_parts.append(f"Session context: {query.context}")
        
        if context_parts:
            context_str = " ".join(context_parts)
            enhanced_query = f"{query.original_query} Context: {context_str}"
            return enhanced_query
        
        return query.original_query
    
    async def _apply_contextual_filters(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery
    ) -> List[RetrievalResult]:
        """Apply contextual filters based on user preferences."""
        # This is a placeholder for more sophisticated contextual filtering
        # In practice, you'd consider user preferences, expertise level, etc.
        return results


class ResultRanker:
    """Advanced result ranking and scoring."""
    
    def __init__(self, gemini_client: GeminiClient):
        self.gemini_client = gemini_client
        self.logger = get_logger(f"{__name__}.ranker")
    
    async def rerank(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery
    ) -> List[RetrievalResult]:
        """Re-rank results using advanced scoring."""
        if not results:
            return results
        
        # Calculate relevance scores
        relevance_scores = await self._calculate_relevance_scores(results, query)
        
        # Calculate diversity scores
        diversity_scores = self._calculate_diversity_scores(results)
        
        # Calculate authority scores
        authority_scores = self._calculate_authority_scores(results)
        
        # Combine scores
        final_scores = []
        for i, result in enumerate(results):
            combined_score = (
                0.6 * relevance_scores[i] +
                0.2 * diversity_scores[i] +
                0.2 * authority_scores[i]
            )
            final_scores.append(combined_score)
            result.score = combined_score
        
        # Sort by combined score
        sorted_results = sorted(
            zip(results, final_scores),
            key=lambda x: x[1],
            reverse=True
        )
        
        return [result for result, score in sorted_results]
    
    async def _calculate_relevance_scores(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery
    ) -> List[float]:
        """Calculate relevance scores using LLM evaluation."""
        relevance_scores = []
        
        for result in results:
            try:
                # Use LLM to evaluate relevance
                prompt = f"""
                Rate the relevance of the following content to the query on a scale of 0-1.
                
                Query: "{query.original_query}"
                Content: "{result.content[:500]}..."
                
                Consider:
                - Direct answer to the question
                - Topical relevance
                - Completeness of information
                
                Return only a number between 0 and 1.
                """
                
                response = await self.gemini_client.generate(prompt, temperature=0.1)
                score = float(response.strip())
                relevance_scores.append(max(0, min(1, score)))
                
            except Exception as e:
                self.logger.warning(f"Failed to calculate relevance score: {e}")
                relevance_scores.append(result.score)  # Fallback to original score
        
        return relevance_scores
    
    def _calculate_diversity_scores(self, results: List[RetrievalResult]) -> List[float]:
        """Calculate diversity scores to avoid redundant results."""
        if len(results) <= 1:
            return [1.0] * len(results)
        
        # Simple diversity based on content similarity
        contents = [result.content for result in results]
        
        try:
            # Use TF-IDF to calculate content similarity
            vectorizer = TfidfVectorizer(max_features=1000, stop_words='english')
            tfidf_matrix = vectorizer.fit_transform(contents)
            similarity_matrix = cosine_similarity(tfidf_matrix)
            
            diversity_scores = []
            for i in range(len(results)):
                # Diversity is inverse of average similarity to other results
                similarities = similarity_matrix[i]
                avg_similarity = np.mean([sim for j, sim in enumerate(similarities) if i != j])
                diversity_score = 1.0 - avg_similarity
                diversity_scores.append(max(0, diversity_score))
            
            return diversity_scores
            
        except Exception:
            # Fallback: uniform diversity scores
            return [1.0] * len(results)
    
    def _calculate_authority_scores(self, results: List[RetrievalResult]) -> List[float]:
        """Calculate authority scores based on source metadata."""
        authority_scores = []
        
        for result in results:
            score = 0.5  # Base score
            
            # Boost score based on source type
            source_type = result.metadata.get('source_type', '').lower()
            if source_type in ['academic', 'research', 'official']:
                score += 0.3
            elif source_type in ['news', 'documentation']:
                score += 0.2
            elif source_type in ['blog', 'forum']:
                score += 0.1
            
            # Boost score based on recency
            doc_date = result.metadata.get('date')
            if doc_date:
                try:
                    if isinstance(doc_date, str):
                        doc_date = datetime.fromisoformat(doc_date.replace('Z', '+00:00'))
                    
                    days_old = (datetime.now() - doc_date).days
                    if days_old < 30:
                        score += 0.2
                    elif days_old < 365:
                        score += 0.1
                    
                except Exception:
                    pass
            
            authority_scores.append(min(1.0, score))
        
        return authority_scores


class ContextCompressor:
    """Compress and optimize retrieved context for generation."""
    
    def __init__(self, gemini_client: GeminiClient):
        self.gemini_client = gemini_client
        self.logger = get_logger(f"{__name__}.compressor")
    
    async def compress_context(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery,
        max_tokens: int = 4000
    ) -> str:
        """Compress retrieved context to fit within token limits."""
        if not results:
            return ""
        
        # Calculate current token count (rough estimation)
        total_content = " ".join([result.content for result in results])
        estimated_tokens = len(total_content.split()) * 1.3  # Rough estimation
        
        if estimated_tokens <= max_tokens:
            return self._format_context(results)
        
        # Need compression
        compressed_content = await self._compress_with_llm(results, query, max_tokens)
        return compressed_content
    
    async def _compress_with_llm(
        self, 
        results: List[RetrievalResult], 
        query: RetrievalQuery,
        max_tokens: int
    ) -> str:
        """Use LLM to compress context while preserving relevance."""
        # Combine all content
        combined_content = "\n\n".join([
            f"Source {i+1}: {result.content}"
            for i, result in enumerate(results)
        ])
        
        prompt = f"""
        Compress the following retrieved content to approximately {max_tokens // 4} words while preserving all information relevant to the query.
        
        Query: "{query.original_query}"
        
        Content to compress:
        {combined_content[:8000]}  # Limit input to avoid token limits
        
        Requirements:
        - Preserve key facts and information relevant to the query
        - Maintain source attribution where possible
        - Remove redundant information
        - Keep the most important details
        
        Compressed content:
        """
        
        try:
            compressed = await self.gemini_client.generate(prompt, temperature=0.1)
            return compressed
        except Exception as e:
            self.logger.warning(f"Failed to compress context: {e}")
            # Fallback: truncate content
            return combined_content[:max_tokens * 4]  # Rough character estimation
    
    def _format_context(self, results: List[RetrievalResult]) -> str:
        """Format context with proper source attribution."""
        formatted_parts = []
        
        for i, result in enumerate(results, 1):
            source_info = f"Source {i}"
            if 'title' in result.metadata:
                source_info += f" ({result.metadata['title']})"
            
            formatted_parts.append(f"{source_info}:\n{result.content}")
        
        return "\n\n".join(formatted_parts)


class AdvancedRAGSystem:
    """Complete Advanced RAG system."""
    
    def __init__(
        self,
        vector_store: VectorStore,
        gemini_client: GeminiClient,
        config: ToxoConfig
    ):
        self.vector_store = vector_store
        self.gemini_client = gemini_client
        self.config = config
        self.logger = get_logger(f"{__name__}.rag_system")
        
        # Initialize components
        self.retriever = AdvancedRetriever(vector_store, gemini_client, config)
        self.context_compressor = ContextCompressor(gemini_client)
        
        # Self-correction settings
        self.max_correction_iterations = 2
        self.confidence_threshold = 0.7
    
    async def query(
        self,
        question: str,
        strategy: RetrievalStrategy = RetrievalStrategy.HYBRID,
        top_k: int = 10,
        context: Optional[Dict[str, Any]] = None,
        self_correct: bool = True
    ) -> RAGResponse:
        """Main RAG query method with self-correction."""
        iteration = 0
        response = None
        
        while iteration <= self.max_correction_iterations:
            try:
                # Retrieve relevant context
                retrieved_results = await self.retriever.retrieve(
                    question, strategy, top_k, context
                )
                
                if not retrieved_results:
                    return RAGResponse(
                        answer="I couldn't find relevant information to answer your question.",
                        retrieved_contexts=[],
                        confidence_score=0.0,
                        retrieval_query=await self.retriever.query_processor.process_query(question, context)
                    )
                
                # Compress context
                compressed_context = await self.context_compressor.compress_context(
                    retrieved_results,
                    await self.retriever.query_processor.process_query(question, context)
                )
                
                # Generate answer
                answer, reasoning_trace = await self._generate_answer(
                    question, compressed_context, retrieved_results
                )
                
                # Calculate confidence
                confidence = await self._calculate_confidence(
                    question, answer, retrieved_results
                )
                
                # Create response
                response = RAGResponse(
                    answer=answer,
                    retrieved_contexts=retrieved_results,
                    confidence_score=confidence,
                    retrieval_query=await self.retriever.query_processor.process_query(question, context),
                    reasoning_trace=reasoning_trace,
                    citations=self._generate_citations(retrieved_results)
                )
                
                # Check if self-correction is needed
                if not self_correct or confidence >= self.confidence_threshold or iteration >= self.max_correction_iterations:
                    break
                
                # Self-correction iteration
                question = await self._refine_query_for_correction(question, answer, retrieved_results)
                iteration += 1
                
            except Exception as e:
                self.logger.error(f"RAG query failed on iteration {iteration}: {e}")
                break
        
        return response or RAGResponse(
            answer="I encountered an error while processing your question.",
            retrieved_contexts=[],
            confidence_score=0.0,
            retrieval_query=await self.retriever.query_processor.process_query(question, context)
        )
    
    async def _generate_answer(
        self,
        question: str,
        context: str,
        retrieved_results: List[RetrievalResult]
    ) -> Tuple[str, List[str]]:
        """Generate answer using retrieved context."""
        prompt = f"""
        Answer the following question using the provided context. Be comprehensive and accurate.
        
        Question: {question}
        
        Context:
        {context}
        
        Instructions:
        - Use only information from the provided context
        - If the context doesn't contain enough information, acknowledge this
        - Provide specific details and examples when available
        - Mention source numbers when referencing specific information
        - Be clear and well-structured
        
        Answer:
        """
        
        try:
            answer = await self.gemini_client.generate(prompt, temperature=0.3)
            
            # Generate reasoning trace
            reasoning_trace = [
                f"Retrieved {len(retrieved_results)} relevant documents",
                f"Generated answer based on retrieved context",
                f"Answer length: {len(answer)} characters"
            ]
            
            return answer, reasoning_trace
            
        except Exception as e:
            self.logger.error(f"Failed to generate answer: {e}")
            return "I encountered an error while generating the answer.", []
    
    async def _calculate_confidence(
        self,
        question: str,
        answer: str,
        retrieved_results: List[RetrievalResult]
    ) -> float:
        """Calculate confidence score for the generated answer."""
        try:
            # Use LLM to evaluate answer quality
            prompt = f"""
            Evaluate the quality and confidence of this answer on a scale of 0-1.
            
            Question: {question}
            Answer: {answer}
            
            Consider:
            - Completeness of the answer
            - Accuracy based on available context
            - Relevance to the question
            - Clarity and coherence
            
            Return only a number between 0 and 1.
            """
            
            response = await self.gemini_client.generate(prompt, temperature=0.1)
            confidence = float(response.strip())
            return max(0, min(1, confidence))
            
        except Exception as e:
            self.logger.warning(f"Failed to calculate confidence: {e}")
            # Fallback: base confidence on retrieval scores
            if retrieved_results:
                avg_retrieval_score = np.mean([r.score for r in retrieved_results])
                return min(0.8, avg_retrieval_score)
            return 0.5
    
    async def _refine_query_for_correction(
        self,
        original_question: str,
        current_answer: str,
        retrieved_results: List[RetrievalResult]
    ) -> str:
        """Refine query for self-correction iteration."""
        prompt = f"""
        The following answer may be incomplete or inaccurate. Refine the original question to retrieve better information.
        
        Original question: {original_question}
        Current answer: {current_answer}
        
        Create a refined question that might retrieve more relevant information. Focus on:
        - Missing aspects of the original question
        - Areas where the current answer seems weak
        - More specific terminology or concepts
        
        Refined question:
        """
        
        try:
            refined_question = await self.gemini_client.generate(prompt, temperature=0.5)
            return refined_question.strip()
        except Exception as e:
            self.logger.warning(f"Failed to refine query: {e}")
            return original_question
    
    def _generate_citations(self, retrieved_results: List[RetrievalResult]) -> List[Dict[str, Any]]:
        """Generate citations from retrieved results."""
        citations = []
        
        for i, result in enumerate(retrieved_results, 1):
            citation = {
                'id': i,
                'source': result.source,
                'title': result.metadata.get('title', f'Source {i}'),
                'score': result.score,
                'retrieval_method': result.retrieval_method
            }
            
            # Add additional metadata if available
            if 'author' in result.metadata:
                citation['author'] = result.metadata['author']
            if 'date' in result.metadata:
                citation['date'] = result.metadata['date']
            if 'url' in result.metadata:
                citation['url'] = result.metadata['url']
            
            citations.append(citation)
        
        return citations
    
    def get_state(self) -> Dict[str, Any]:
        """Get the current state of the Advanced RAG system for serialization."""
        return {
            "vector_store_available": self.vector_store is not None,
            "gemini_client_available": self.gemini_client is not None,
            "config_available": self.config is not None,
            "retriever_available": hasattr(self, 'retriever') and self.retriever is not None,
            "context_compressor_available": hasattr(self, 'context_compressor') and self.context_compressor is not None,
            "max_correction_iterations": getattr(self, 'max_correction_iterations', 2),
            "confidence_threshold": getattr(self, 'confidence_threshold', 0.7),
            "system_type": "AdvancedRAGSystem"
        }


# Factory function
def create_advanced_rag_system(
    vector_store: VectorStore,
    gemini_client: GeminiClient,
    config: ToxoConfig
) -> AdvancedRAGSystem:
    """Create an advanced RAG system with all components."""
    return AdvancedRAGSystem(vector_store, gemini_client, config) 