"""Enhanced RAG-Powered Analytics Agent with Community Insights.

Integrates retrieval-augmented generation for context-aware analytics,
versioning hooks per SCHEMA_VERSIONING.md, and rollback capabilities.
Enhanced with patterns from bziiit/AGENT_ANALYSE_RAG and fahmiaziz/rag-sql-agent.
"""

import os
import json
import logging
from datetime import datetime
from typing import Dict, List, Optional, Any, Tuple
from dataclasses import dataclass, asdict
import sys

# Import RAG pipeline for context retrieval
from smolagents.rag_pipeline import RAGPipeline, RetrievalConfig
from schema_registry import SchemaRegistry

logger = logging.getLogger(__name__)

@dataclass
class AnalyticsConfig:
    """Configuration for RAG-powered analytics."""
    rag_enabled: bool = True
    top_k_docs: int = 5
    similarity_threshold: float = 0.7
    enable_versioning: bool = True
    schema_version: str = "1.0.0"
    rollback_enabled: bool = True
    audit_trail: bool = True

@dataclass
class AnalyticsQuery:
    """Structured analytics query with RAG context."""
    query_id: str
    query_text: str
    context_docs: List[Dict[str, Any]]
    filters: Dict[str, Any]
    aggregations: List[str]
    timestamp: str
    schema_version: str

@dataclass
class AnalyticsResult:
    """Enhanced analytics result with citations and provenance."""
    query_id: str
    data: Dict[str, Any]
    citations: List[Dict[str, Any]]
    confidence_score: float
    schema_version: str
    rollback_checkpoint: Optional[str]
    timestamp: str

class RAGAnalyticsAgent:
    """Enhanced Analytics Agent with RAG, versioning, and rollback."""
    
    def __init__(
        self,
        config: Optional[AnalyticsConfig] = None,
        rag_pipeline: Optional[RAGPipeline] = None,
        schema_registry: Optional[SchemaRegistry] = None
    ):
        self.config = config or AnalyticsConfig()
        self.rag_pipeline = rag_pipeline
        self.schema_registry = schema_registry
        self.checkpoints: Dict[str, Dict] = {}
        
        logger.info(f"Initialized RAGAnalyticsAgent with config: {self.config}")
    
    def create_checkpoint(self, checkpoint_id: str, state: Dict[str, Any]) -> None:
        """Create rollback checkpoint per SCHEMA_VERSIONING.md."""
        self.checkpoints[checkpoint_id] = {
            "state": state,
            "timestamp": datetime.utcnow().isoformat(),
            "schema_version": self.config.schema_version
        }
        logger.info(f"Created checkpoint: {checkpoint_id}")
    
    def rollback_to_checkpoint(self, checkpoint_id: str) -> Optional[Dict[str, Any]]:
        """Rollback to a previous checkpoint."""
        if checkpoint_id in self.checkpoints:
            checkpoint = self.checkpoints[checkpoint_id]
            logger.info(f"Rolling back to checkpoint: {checkpoint_id}")
            return checkpoint["state"]
        logger.warning(f"Checkpoint not found: {checkpoint_id}")
        return None
    
    def retrieve_context(
        self,
        query_text: str,
        top_k: Optional[int] = None
    ) -> List[Dict[str, Any]]:
        """Retrieve relevant context using RAG pipeline."""
        if not self.config.rag_enabled or not self.rag_pipeline:
            return []
        
        top_k = top_k or self.config.top_k_docs
        
        try:
            results = self.rag_pipeline.retrieve(
                query=query_text,
                top_k=top_k,
                threshold=self.config.similarity_threshold
            )
            
            context_docs = [
                {
                    "content": r.get("content", ""),
                    "source": r.get("metadata", {}).get("source", "unknown"),
                    "score": r.get("score", 0.0),
                    "metadata": r.get("metadata", {})
                }
                for r in results
            ]
            
            logger.info(f"Retrieved {len(context_docs)} context documents")
            return context_docs
            
        except Exception as e:
            logger.error(f"Context retrieval failed: {e}")
            return []
    
    def validate_schema(
        self,
        data: Dict[str, Any],
        schema_version: Optional[str] = None
    ) -> Tuple[bool, List[str]]:
        """Validate data against schema version."""
        if not self.config.enable_versioning or not self.schema_registry:
            return True, []
        
        schema_version = schema_version or self.config.schema_version
        
        try:
            is_valid = self.schema_registry.validate(data, schema_version)
            errors = [] if is_valid else ["Schema validation failed"]
            return is_valid, errors
        except Exception as e:
            logger.error(f"Schema validation error: {e}")
            return False, [str(e)]
    
    def analyze(
        self,
        query_text: str,
        data_source: Dict[str, Any],
        filters: Optional[Dict[str, Any]] = None,
        aggregations: Optional[List[str]] = None
    ) -> AnalyticsResult:
        """Perform RAG-enhanced analytics with versioning and rollback."""
        query_id = f"analytics_{datetime.utcnow().timestamp()}"
        
        # Create checkpoint before analysis
        if self.config.rollback_enabled:
            self.create_checkpoint(
                f"{query_id}_pre",
                {"data_source": data_source, "filters": filters}
            )
        
        # Retrieve context
        context_docs = self.retrieve_context(query_text)
        
        # Build query
        query = AnalyticsQuery(
            query_id=query_id,
            query_text=query_text,
            context_docs=context_docs,
            filters=filters or {},
            aggregations=aggregations or [],
            timestamp=datetime.utcnow().isoformat(),
            schema_version=self.config.schema_version
        )
        
        # Perform analysis (simplified - extend with actual analytics logic)
        try:
            analysis_data = self._execute_analysis(query, data_source)
            
            # Validate result schema
            is_valid, errors = self.validate_schema(analysis_data)
            if not is_valid:
                logger.error(f"Schema validation failed: {errors}")
                # Attempt rollback
                if self.config.rollback_enabled:
                    rollback_state = self.rollback_to_checkpoint(f"{query_id}_pre")
                    if rollback_state:
                        logger.info("Successfully rolled back to pre-analysis state")
            
            # Build citations from context
            citations = [
                {
                    "source": doc["source"],
                    "relevance_score": doc["score"],
                    "excerpt": doc["content"][:200]
                }
                for doc in context_docs
            ]
            
            # Calculate confidence based on context quality
            avg_score = sum(d["score"] for d in context_docs) / len(context_docs) if context_docs else 0.5
            
            result = AnalyticsResult(
                query_id=query_id,
                data=analysis_data,
                citations=citations,
                confidence_score=avg_score,
                schema_version=self.config.schema_version,
                rollback_checkpoint=f"{query_id}_pre" if self.config.rollback_enabled else None,
                timestamp=datetime.utcnow().isoformat()
            )
            
            # Audit trail
            if self.config.audit_trail:
                self._log_audit(query, result)
            
            return result
            
        except Exception as e:
            logger.error(f"Analysis failed: {e}")
            # Return error result with rollback info
            return AnalyticsResult(
                query_id=query_id,
                data={"error": str(e)},
                citations=[],
                confidence_score=0.0,
                schema_version=self.config.schema_version,
                rollback_checkpoint=f"{query_id}_pre" if self.config.rollback_enabled else None,
                timestamp=datetime.utcnow().isoformat()
            )
    
    def _execute_analysis(
        self,
        query: AnalyticsQuery,
        data_source: Dict[str, Any]
    ) -> Dict[str, Any]:
        """Execute the actual analysis logic (extend as needed)."""
        # Simplified analytics - enhance with real logic
        results = {
            "query_id": query.query_id,
            "metrics": {},
            "insights": [],
            "context_applied": len(query.context_docs) > 0
        }
        
        # Apply filters
        filtered_data = data_source
        for key, value in query.filters.items():
            if key in filtered_data:
                filtered_data = {k: v for k, v in filtered_data.items() if v == value}
        
        # Apply aggregations
        for agg in query.aggregations:
            if agg == "count":
                results["metrics"]["count"] = len(filtered_data)
            elif agg == "sum":
                results["metrics"]["sum"] = sum(v for v in filtered_data.values() if isinstance(v, (int, float)))
        
        # Extract insights from RAG context
        if query.context_docs:
            for doc in query.context_docs[:3]:  # Top 3 docs
                results["insights"].append({
                    "source": doc["source"],
                    "insight": doc["content"][:150]
                })
        
        return results
    
    def _log_audit(
        self,
        query: AnalyticsQuery,
        result: AnalyticsResult
    ) -> None:
        """Log audit trail for compliance."""
        audit_entry = {
            "timestamp": datetime.utcnow().isoformat(),
            "query_id": query.query_id,
            "query_text": query.query_text,
            "schema_version": query.schema_version,
            "confidence_score": result.confidence_score,
            "citations_count": len(result.citations)
        }
        logger.info(f"Audit: {json.dumps(audit_entry)}")
    
    def get_version_info(self) -> Dict[str, Any]:
        """Get current version and rollback information."""
        return {
            "schema_version": self.config.schema_version,
            "checkpoints": list(self.checkpoints.keys()),
            "versioning_enabled": self.config.enable_versioning,
            "rollback_enabled": self.config.rollback_enabled
        }

# Example usage
if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    
    # Initialize agent
    config = AnalyticsConfig(
        rag_enabled=True,
        top_k_docs=5,
        enable_versioning=True,
        rollback_enabled=True
    )
    
    agent = RAGAnalyticsAgent(config=config)
    
    # Example query
    result = agent.analyze(
        query_text="What are the top engagement metrics for Q4?",
        data_source={"engagement": 100, "reach": 500},
        aggregations=["count", "sum"]
    )
    
    print(f"\nAnalysis Result:")
    print(f"Query ID: {result.query_id}")
    print(f"Confidence: {result.confidence_score}")
    print(f"Citations: {len(result.citations)}")
    print(f"Rollback checkpoint: {result.rollback_checkpoint}")
    print(f"\nVersion Info: {json.dumps(agent.get_version_info(), indent=2)}")