┌─────────────────────────────────────────────────────────────────┐
│                    RAG ARCHITECTURE PIPELINE                    │
├─────────────────────────────────────────────────────────────────┤
│                                                                 │
│  ┌──────────────┐     ┌──────────────┐     ┌──────────────┐   │
│  │   INGESTION  │────▶│   INDEXING   │────▶│   RETRIEVAL  │   │
│  │              │     │              │     │              │   │
│  │ • Documents  │     │ • Chunking   │     │ • Query      │   │
│  │ • PDFs       │     │ • Embedding  │     │   Embedding  │   │
│  │ • Emails     │     │ • Vector     │     │ • Similarity │   │
│  │ • Databases  │     │   Storage    │     │   Search     │   │
│  └──────────────┘     └──────────────┘     └──────────────┘   │
│                                                                 │
│  ┌──────────────┐     ┌──────────────┐     ┌──────────────┐   │
│  │   PROMPT     │────▶│   GENERATION │────▶│   GUARDRAILS │   │
│  │   ENGINEERING│     │              │     │              │   │
│  │              │     │ • LLM        │     │ • PII Masking│   │
│  │ • System     │     │   Inference  │     │ • Toxicity   │   │
│  │   Prompts    │     │ • Context    │     │ • Hallucination│  │
│  │ • Few-shot   │     │   Injection  │     │ • Audit Log  │   │
│  │   Examples   │     │ • Response   │     │              │   │
│  └──────────────┘     └──────────────┘     └──────────────┘   │
│                                                                 │
└─────────────────────────────────────────────────────────────────┘

Text: "Wire transfers over $10,000 require AML review."
      ↓
Vector: [0.234, -0.567, 0.891, 0.123, ...]  (1536 dimensions for ada-002)
      ↓
Text: "Large international payments need compliance checks."
      → Similar vector → Close in semantic space

{
  "audit_id": "audit-12345",
  "timestamp": "2025-08-07T14:32:17Z",
  "user_id": "user_001",
  "user_role": "Loan_Officer",
  "query": "What is the maximum loan amount for a first-time home buyer?",
  "retrieved_sources": [
    {"doc": "Home_Loan_Policy_2025.pdf", "page": 12, "relevance": 0.89},
    {"doc": "Lending_Guidelines_v3.pdf", "page": 45, "relevance": 0.76}
  ],
  "context": "For first-time home buyers, the maximum loan amount is $726,200...",
  "response": "Based on the current policy, first-time home buyers can qualify for loans up to $726,200...",
  "hallucination_check": {"passed": true, "score": 0.92},
  "pii_redacted": false,
  "model": "gpt-4-azure-2025-07-01",
  "latency_ms": 1245,
  "user_feedback": null
}

# ===================================================================
# MODULE 10, LESSON 4: ENTERPRISE GENERATIVE AI & RAG
# ===================================================================

import pandas as pd
import numpy as np
import hashlib
import re
import json
import pickle
import os
from datetime import datetime, timedelta
from typing import List, Dict, Tuple, Optional
from dataclasses import dataclass, field
from enum import Enum
import warnings
warnings.filterwarnings('ignore')

print("="*70)
print("ENTERPRISE RAG SYSTEM – COMPLETE IMPLEMENTATION")
print("="*70)

# ----------------------------------------------------------------
# PART A: DATA CLASSES AND ENUMS
# ----------------------------------------------------------------

class RiskLevel(Enum):
    """Risk levels for AI interactions."""
    CRITICAL = "critical"  # Financial decisions, customer PII
    HIGH = "high"          # Policy interpretation
    MEDIUM = "medium"      # General guidance
    LOW = "low"            # FAQ, general information

@dataclass
class Document:
    """Represents a document in the knowledge base."""
    id: str
    content: str
    source: str
    doc_type: str
    created_date: str
    metadata: Dict = field(default_factory=dict)
    
@dataclass
class DocumentChunk:
    """Represents a chunked document for embedding."""
    doc_id: str
    chunk_index: int
    content: str
    metadata: Dict = field(default_factory=dict)

@dataclass
class RetrievalResult:
    """Result from the retrieval stage."""
    chunk: DocumentChunk
    similarity: float
    score: float

@dataclass
class GenerationResult:
    """Complete result from the RAG pipeline."""
    query: str
    redacted_query: str
    response: str
    sources: List[str]
    context: str
    retrieved_chunks: List[RetrievalResult]
    guardrail_results: Dict
    timestamp: str
    model: str
    audit_id: str

# ----------------------------------------------------------------
# PART B: PII REDACTION ENGINE
# ----------------------------------------------------------------

class PIIRedactor:
    """
    Comprehensive PII detection and redaction using regex patterns.
    In production, use Presidio or spaCy NER.
    """
    def __init__(self):
        self.patterns = {
            'email': re.compile(r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'),
            'phone': re.compile(r'\b(?:(\+?1[-.]?)?\(?[0-9]{3}\)?[-.]?[0-9]{3}[-.]?[0-9]{4})\b'),
            'ssn': re.compile(r'\b\d{3}-\d{2}-\d{4}\b'),
            'credit_card': re.compile(r'\b(?:\d{4}[- ]?){3}\d{4}\b'),
            'account_number': re.compile(r'\b\d{10,17}\b'),
            'passport': re.compile(r'\b[A-Z]{1,2}\d{6,9}\b'),
            'address': re.compile(r'\b\d{1,5}\s+[A-Za-z]+\s+(?:Street|St|Avenue|Ave|Road|Rd|Boulevard|Blvd|Lane|Ln|Drive|Dr|Court|Ct)\b', re.IGNORECASE)
        }
        self.redaction_map = {
            'email': '[EMAIL_REDACTED]',
            'phone': '[PHONE_REDACTED]',
            'ssn': '[SSN_REDACTED]',
            'credit_card': '[CREDIT_CARD_REDACTED]',
            'account_number': '[ACCOUNT_REDACTED]',
            'passport': '[PASSPORT_REDACTED]',
            'address': '[ADDRESS_REDACTED]'
        }
        self.detected_pii = []
    
    def redact(self, text: str) -> Tuple[str, Dict]:
        """
        Redact PII from text and return redacted text and detection summary.
        """
        redacted_text = text
        detections = {}
        
        for pii_type, pattern in self.patterns.items():
            matches = pattern.findall(text)
            if matches:
                detections[pii_type] = len(matches)
                redacted_text = pattern.sub(self.redaction_map[pii_type], redacted_text)
                self.detected_pii.append({'type': pii_type, 'count': len(matches)})
        
        return redacted_text, detections
    
    def get_summary(self) -> Dict:
        """Get summary of all PII detections."""
        summary = {}
        for detection in self.detected_pii:
            summary[detection['type']] = summary.get(detection['type'], 0) + detection['count']
        return summary

# ----------------------------------------------------------------
# PART C: VECTOR DATABASE (ChromaDB Simulation)
# ----------------------------------------------------------------

class VectorDatabase:
    """
    Simulates a vector database for storing and retrieving document embeddings.
    In production: use ChromaDB, Pinecone, or Weaviate.
    """
    def __init__(self, embedding_dimension: int = 384):
        self.chunks: List[DocumentChunk] = []
        self.embeddings = None
        self.embedding_dimension = embedding_dimension
        self.similarity_threshold = 0.3  # Minimum similarity for retrieval
        
    def _generate_embedding(self, text: str) -> np.ndarray:
        """
        Simulates embedding generation.
        In production: use OpenAI, Cohere, or open-source embedding models.
        """
        # Simulate meaningful embeddings using TF-IDF-like approach
        import hashlib
        hash_bytes = hashlib.md5(text.encode()).digest()
        # Use hash to generate deterministic pseudo-random embedding
        seed = int.from_bytes(hash_bytes, 'big') % 1000000
        np.random.seed(seed)
        embedding = np.random.randn(self.embedding_dimension)
        embedding = embedding / np.linalg.norm(embedding)  # Normalize
        return embedding
    
    def add_documents(self, chunks: List[DocumentChunk]) -> None:
        """Add document chunks to the vector database."""
        self.chunks.extend(chunks)
        
        # Generate embeddings for all chunks
        embeddings = []
        for chunk in chunks:
            embedding = self._generate_embedding(chunk.content)
            embeddings.append(embedding)
        
        # Convert to numpy array
        if self.embeddings is None:
            self.embeddings = np.array(embeddings)
        else:
            self.embeddings = np.vstack([self.embeddings, np.array(embeddings)])
        
        print(f"✅ Added {len(chunks)} chunks to vector DB. Total: {len(self.chunks)}")
    
    def similarity_search(self, query: str, top_k: int = 5) -> List[RetrievalResult]:
        """
        Search for similar documents using cosine similarity.
        """
        if self.embeddings is None or len(self.chunks) == 0:
            return []
        
        # Generate query embedding
        query_embedding = self._generate_embedding(query)
        
        # Calculate cosine similarity
        similarities = np.dot(self.embeddings, query_embedding) / (
            np.linalg.norm(self.embeddings, axis=1) * np.linalg.norm(query_embedding)
        )
        
        # Sort by similarity
        indices = np.argsort(similarities)[::-1]
        
        results = []
        for idx in indices[:top_k]:
            if similarities[idx] >= self.similarity_threshold:
                result = RetrievalResult(
                    chunk=self.chunks[idx],
                    similarity=float(similarities[idx]),
                    score=float(similarities[idx])  # Simple score = similarity
                )
                results.append(result)
        
        return results

# ----------------------------------------------------------------
# PART D: DOCUMENT PROCESSOR
# ----------------------------------------------------------------

class DocumentProcessor:
    """
    Processes documents for ingestion into the RAG system.
    """
    def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
        self.chunk_size = chunk_size
        self.chunk_overlap = chunk_overlap
        self.documents: List[Document] = []
    
    def add_document(self, content: str, source: str, doc_type: str, 
                     metadata: Optional[Dict] = None) -> Document:
        """Add a document to the processor."""
        doc_id = f"doc_{hashlib.md5((source + datetime.now().isoformat()).encode()).hexdigest()[:8]}"
        doc = Document(
            id=doc_id,
            content=content,
            source=source,
            doc_type=doc_type,
            created_date=datetime.now().isoformat(),
            metadata=metadata or {}
        )
        self.documents.append(doc)
        return doc
    
    def chunk_document(self, doc: Document) -> List[DocumentChunk]:
        """
        Chunk a document into smaller pieces for embedding.
        """
        chunks = []
        content = doc.content
        
        # Split by paragraphs first
        paragraphs = content.split('\n\n')
        
        current_chunk = ""
        chunk_index = 0
        
        for paragraph in paragraphs:
            if len(current_chunk) + len(paragraph) <= self.chunk_size:
                current_chunk += paragraph + "\n\n"
            else:
                # Save current chunk
                if current_chunk.strip():
                    chunk = DocumentChunk(
                        doc_id=doc.id,
                        chunk_index=chunk_index,
                        content=current_chunk.strip(),
                        metadata={
                            'source': doc.source,
                            'doc_type': doc.doc_type,
                            'created_date': doc.created_date,
                            **doc.metadata
                        }
                    )
                    chunks.append(chunk)
                    chunk_index += 1
                    
                    # Start new chunk with overlap
                    overlap_content = current_chunk.split('\n\n')[-2:] if current_chunk.count('\n\n') > 1 else []
                    current_chunk = "\n\n".join(overlap_content) + "\n\n" if overlap_content else ""
                    current_chunk += paragraph + "\n\n"
        
        # Add the last chunk
        if current_chunk.strip():
            chunk = DocumentChunk(
                doc_id=doc.id,
                chunk_index=chunk_index,
                content=current_chunk.strip(),
                metadata={
                    'source': doc.source,
                    'doc_type': doc.doc_type,
                    'created_date': doc.created_date,
                    **doc.metadata
                }
            )
            chunks.append(chunk)
        
        return chunks
    
    def process_all_documents(self) -> List[DocumentChunk]:
        """Process all documents into chunks."""
        all_chunks = []
        for doc in self.documents:
            chunks = self.chunk_document(doc)
            all_chunks.extend(chunks)
        return all_chunks

# ----------------------------------------------------------------
# PART E: HALLUCINATION DETECTOR
# ----------------------------------------------------------------

class HallucinationDetector:
    """
    Detects potential hallucinations in LLM responses.
    """
    def __init__(self):
        self.threshold = 0.6
    
    def check_grounding(self, response: str, context: str) -> Dict:
        """
        Check if the response is grounded in the provided context.
        Uses semantic similarity and keyword overlap.
        """
        # Compute keyword overlap
        response_words = set(response.lower().split())
        context_words = set(context.lower().split())
        
        # Remove stopwords (simplified)
        stopwords = {'the', 'a', 'an', 'of', 'to', 'for', 'with', 'on', 'at', 'from', 'by', 'in', 'is', 'are', 'was', 'were'}
        response_filtered = response_words - stopwords
        context_filtered = context_words - stopwords
        
        # Calculate overlap ratio
        overlap = response_filtered.intersection(context_filtered)
        overlap_ratio = len(overlap) / max(len(response_filtered), 1)
        
        # Simulate semantic similarity (in production: use embedding similarity)
        semantic_similarity = min(1.0, overlap_ratio * 1.2)
        
        # Determine if grounded
        is_grounded = semantic_similarity >= self.threshold
        
        return {
            'is_grounded': is_grounded,
            'semantic_similarity': semantic_similarity,
            'overlap_ratio': overlap_ratio,
            'overlap_count': len(overlap),
            'word_match_ratio': overlap_ratio
        }

# ----------------------------------------------------------------
# PART F: COMPLETE RAG PIPELINE
# ----------------------------------------------------------------

class EnterpriseRAGPipeline:
    """
    Complete RAG pipeline with all components and guardrails.
    """
    def __init__(self, vector_db: VectorDatabase, redactor: PIIRedactor, 
                 hallucination_detector: HallucinationDetector):
        self.vector_db = vector_db
        self.redactor = redactor
        self.hallucination_detector = hallucination_detector
        self.audit_log: List[Dict] = []
        self.audit_id_counter = 0
        
        # Configuration
        self.llm_model = "Azure-OpenAI-GPT-4"  # In production: actual deployment
        self.retrieval_top_k = 5
        self.hallucination_threshold = 0.6
        self.max_response_length = 2000
        
        # System prompt template
        self.system_prompt = """
        You are a banking expert assistant. Your task is to answer questions based ONLY on the provided context.
        
        Context:
        {context}
        
        Instructions:
        1. Answer only based on the context provided above.
        2. If the context doesn't contain the answer, say "I don't have that information in my knowledge base."
        3. If you need to provide advice, clearly state that it is based on the available information.
        4. Never make up information or speculate.
        5. When referencing specific policies, mention the source document.
        """
    
    def ingest_documents(self, documents: List[Document]) -> None:
        """Ingest documents into the knowledge base."""
        processor = DocumentProcessor()
        for doc in documents:
            processor.add_document(doc.content, doc.source, doc.doc_type, doc.metadata)
        
        chunks = processor.process_all_documents()
        self.vector_db.add_documents(chunks)
        print(f"📚 Ingested {len(documents)} documents into knowledge base.")
    
    def retrieve(self, query: str) -> Tuple[str, List[RetrievalResult]]:
        """
        Retrieve relevant document chunks for the query.
        """
        results = self.vector_db.similarity_search(query, top_k=self.retrieval_top_k)
        
        if not results:
            return "No relevant documents found in the knowledge base.", []
        
        # Combine chunks into context
        context_parts = []
        for r in results:
            context_parts.append(f"From {r.chunk.metadata.get('source', 'unknown')}:\n{r.chunk.content}")
        
        context = "\n\n---\n\n".join(context_parts)
        return context, results
    
    def generate_response(self, query: str, context: str) -> str:
        """
        Generate a response using the LLM.
        In production: Call Azure OpenAI, AWS Bedrock, or similar.
        """
        # In production: This would be an API call to an LLM
        # Here we simulate the response based on context
        
        # Simulate LLM response with simple rule-based generation
        response = self._simulate_llm(query, context)
        
        return response
    
    def _simulate_llm(self, query: str, context: str) -> str:
        """
        Simulate LLM response for demonstration.
        In production: Replace with actual LLM API call.
        """
        query_lower = query.lower()
        response = ""
        
        # Check if context has relevant information
        if "no relevant documents" in context or not context.strip():
            return "I don't have that information in my knowledge base. Please consult the official documentation or contact your compliance officer."
        
        # Simulate responses based on query patterns
        if "wire transfer" in query_lower or "aml" in query_lower:
            if "10,000" in context or "ten thousand" in context.lower():
                response = "Based on the AML policy, wire transfers exceeding $10,000 require manual review. This applies to both domestic and international transfers, and the review must be completed within 24 hours."
            else:
                response = "Based on the available context, wire transfers have specific thresholds that trigger AML review. Please check the specific policy documents for the exact limit."
        
        elif "loan" in query_lower or "credit" in query_lower or "mortgage" in query_lower:
            if "credit score" in context.lower() or "680" in context:
                response = "Based on the underwriting guidelines, a credit score above 680 qualifies for the standard loan rate for amounts under $50,000. For amounts above this, additional income verification is required."
            elif "first-time" in query_lower and "home" in query_lower:
                response = "Based on the home loan policy, first-time home buyers are eligible for special programs. Please refer to the Home Loan Policy document for specific details on rates and eligibility."
            else:
                response = "Based on the lending guidelines, loan approvals are subject to credit assessment, income verification, and the specific product terms. Please refer to the Loan Underwriting Guide for complete details."
        
        elif "onboarding" in query_lower or "kyc" in query_lower:
            if "biometric" in context.lower() or "15 minutes" in context.lower():
                response = "Based on the onboarding procedure, digital onboarding requires biometric verification (facial match) against the submitted government ID. The process should be completed within 15 minutes for a seamless customer experience."
            else:
                response = "Based on the available context, the digital onboarding process includes identity verification and biometric checks. Please refer to the Digital Onboarding Procedure document for complete details."
        
        elif "encryption" in query_lower or "pii" in query_lower or "security" in query_lower:
            if "encrypt" in context.lower():
                response = "Based on the data protection policy, all customer PII must be encrypted at rest and in transit. Access to sensitive data requires multi-factor authentication and is logged in an immutable audit trail."
            else:
                response = "Based on the available context, data security is a priority with encryption requirements for PII. Please refer to the Customer Data Protection document for specific technical requirements."
        
        elif "fraud" in query_lower or "risk" in query_lower:
            response = "Based on the available context, the bank has fraud detection mechanisms in place. For specific details on fraud prevention, please refer to the Fraud Risk Management policy."
        
        else:
            # Generic response if no specific match
            if len(context) > 100:
                # Extract a summary-like response
                sentences = context.split('.')
                summary = '. '.join(sentences[:3])
                response = f"Based on the available context: {summary}."
            else:
                response = "Based on the available context, I can provide information on this topic. Please refer to the specific policy documents for detailed guidance."
        
        return response
    
    def apply_guardrails(self, query: str, response: str, context: str) -> Dict:
        """
        Apply all guardrails to the response.
        """
        guardrail_results = {}
        
        # 1. PII detection and redaction
        redacted_response, pii_detections = self.redactor.redact(response)
        guardrail_results['pii_redacted'] = len(pii_detections) > 0
        guardrail_results['pii_detections'] = pii_detections
        
        # 2. Hallucination detection
        hallucination_check = self.hallucination_detector.check_grounding(redacted_response, context)
        guardrail_results['hallucination_check'] = hallucination_check
        
        # 3. Response length check
        guardrail_results['response_length'] = len(redacted_response)
        guardrail_results['response_truncated'] = len(redacted_response) > self.max_response_length
        
        return guardrail_results
    
    def process_query(self, query: str, user_id: str = "anonymous", 
                      risk_level: RiskLevel = RiskLevel.MEDIUM) -> GenerationResult:
        """
        Process a query through the complete RAG pipeline.
        """
        print(f"\n{'='*60}")
        print(f"📝 Processing Query: '{query[:100]}...'")
        print(f"   User: {user_id} | Risk Level: {risk_level.value}")
        print(f"{'='*60}")
        
        # 1. PII Redaction on Query
        redacted_query, query_pii = self.redactor.redact(query)
        print(f"   🔒 PII Redacted: {'Yes' if query_pii else 'No'}")
        
        # 2. Retrieval
        context, results = self.retrieve(redacted_query)
        print(f"   📚 Retrieved {len(results)} document chunks")
        if results:
            for r in results[:2]:
                print(f"      - Source: {r.chunk.metadata.get('source', 'unknown')} (Similarity: {r.similarity:.3f})")
        
        # 3. Generation
        response = self.generate_response(redacted_query, context)
        print(f"   💬 Generated Response: '{response[:150]}...'")
        
        # 4. Guardrails
        guardrail_results = self.apply_guardrails(redacted_query, response, context)
        print(f"   🛡️ Guardrails Applied: {list(guardrail_results.keys())}")
        
        # 5. Final response (use redacted response)
        final_response = guardrail_results.get('redacted_response', response)
        
        # If hallucination detected, override response
        if guardrail_results.get('hallucination_check', {}