"""
BeastBullet + RAG: Personal Knowledge Assistant
Query your documents, emails, notes - 100% local
"""

import os
from pathlib import Path
from typing import List, Dict
import json

class DocumentRAG:
    """Document retrieval for personal knowledge."""
    
    def __init__(self, vector_db="chromadb", collection_name="my_knowledge"):
        self.collection_name = collection_name
        self.db = self._init_vector_db(vector_db)
    
    def _init_vector_db(self, db_type: str):
        """Initialize vector database."""
        if db_type == "chromadb":
            try:
                import chromadb
                client = chromadb.Client()
                collection = client.get_or_create_collection(
                    name=self.collection_name
                )
                print(f"✅ ChromaDB initialized: {self.collection_name}")
                return collection
            except ImportError:
                print("⚠️  Install chromadb: pip install chromadb")
                return None
        else:
            raise ValueError(f"Unsupported vector DB: {db_type}")
    
    def index_directory(self, directory: str, extensions: List[str] = None):
        """Index all documents in a directory."""
        if not self.db:
            return
        
        if extensions is None:
            extensions = ['.txt', '.md', '.pdf', '.docx', '.py', '.js', '.json']
        
        directory = Path(directory).expanduser()
        print(f"\n📁 Indexing: {directory}")
        
        indexed = 0
        for ext in extensions:
            for file_path in directory.rglob(f"*{ext}"):
                try:
                    # Read file
                    if ext == '.pdf':
                        content = self._read_pdf(file_path)
                    else:
                        with open(file_path, 'r', encoding='utf-8', errors='ignore') as f:
                            content = f.read()
                    
                    if not content.strip():
                        continue
                    
                    # Chunk content
                    chunks = self._chunk_text(content)
                    
                    # Add to vector DB
                    for i, chunk in enumerate(chunks):
                        self.db.add(
                            documents=[chunk],
                            ids=[f"{file_path}_{i}"],
                            metadatas=[{
                                "source": str(file_path),
                                "chunk": i,
                                "total_chunks": len(chunks)
                            }]
                        )
                    
                    indexed += 1
                    print(f"  ✅ {file_path.name} ({len(chunks)} chunks)")
                
                except Exception as e:
                    print(f"  ⚠️  Error indexing {file_path.name}: {e}")
        
        print(f"\n✅ Indexed {indexed} documents")
    
    def _chunk_text(self, text: str, chunk_size: int = 500) -> List[str]:
        """Split text into chunks."""
        words = text.split()
        chunks = []
        
        for i in range(0, len(words), chunk_size):
            chunk = ' '.join(words[i:i + chunk_size])
            if chunk.strip():
                chunks.append(chunk)
        
        return chunks
    
    def _read_pdf(self, file_path: Path) -> str:
        """Read PDF file."""
        try:
            import PyPDF2
            with open(file_path, 'rb') as f:
                reader = PyPDF2.PdfReader(f)
                text = ""
                for page in reader.pages:
                    text += page.extract_text()
                return text
        except ImportError:
            print("⚠️  Install PyPDF2: pip install PyPDF2")
            return ""
    
    def retrieve(self, query: str, n_results: int = 5) -> Dict:
        """Retrieve relevant documents."""
        if not self.db:
            return {"documents": [], "sources": []}
        
        results = self.db.query(
            query_texts=[query],
            n_results=n_results
        )
        
        # Format results
        documents = results['documents'][0] if results['documents'] else []
        metadatas = results['metadatas'][0] if results['metadatas'] else []
        
        sources = []
        for meta in metadatas:
            if meta['source'] not in [s['file'] for s in sources]:
                sources.append({
                    "file": meta['source'],
                    "chunks": 1
                })
            else:
                for s in sources:
                    if s['file'] == meta['source']:
                        s['chunks'] += 1
        
        return {
            "documents": documents,
            "sources": sources,
            "context": "\n\n".join(documents)
        }


class PersonalKnowledgeAssistant:
    """Personal knowledge assistant with BeastBullet + RAG."""
    
    def __init__(self):
        from code.integration_example import TinyLlamaIntegration, ExpertLoader
        
        print("🧠 Loading BeastBullet...")
        loader = ExpertLoader("experts/")
        loader.load_encoder()
        loader.load_all_experts()
        
        self.bb = TinyLlamaIntegration(
            experts=loader.experts,
            encoder=loader.encoder
        )
        
        self.rag = DocumentRAG()
        print("✅ Personal Knowledge Assistant ready!\n")
    
    def index_my_knowledge(self):
        """Index common personal knowledge locations."""
        print("📚 Indexing your knowledge base...")
        
        # Common locations
        locations = [
            "~/Documents",
            "~/Downloads",
            "~/Desktop",
            "~/Notes",
        ]
        
        for location in locations:
            path = Path(location).expanduser()
            if path.exists():
                self.rag.index_directory(path)
    
    def query(self, question: str) -> Dict:
        """Query your personal knowledge."""
        print("=" * 70)
        print(f"Question: {question}")
        print("=" * 70)
        
        # Retrieve relevant documents
        retrieval = self.rag.retrieve(question)
        
        if not retrieval["documents"]:
            print("\n⚠️  No relevant documents found in your knowledge base")
            result = self.bb.query(question)
            return result
        
        # Build context-aware query
        context_query = f"""Based on my personal documents, answer: {question}

My Documents:
{retrieval['context']}

Provide a clear answer based on my documents."""
        
        # Query BeastBullet with context
        result = self.bb.query(context_query)
        result["sources"] = retrieval["sources"]
        
        return result
    
    def display_result(self, result: Dict):
        """Display result."""
        print("\n" + "=" * 70)
        print("🧠 Answer:")
        print("=" * 70)
        print(result.get("answer", "No answer generated"))
        
        if "sources" in result and result["sources"]:
            print("\n" + "=" * 70)
            print("📁 Sources from your documents:")
            print("=" * 70)
            for source in result["sources"]:
                print(f"  • {source['file']} ({source['chunks']} relevant sections)")
        
        print("\n" + "=" * 70)
        print(f"Experts: {', '.join(result.get('experts_used', []))}")
        print(f"Confidence: {result.get('confidence', 0):.1%}")
        print("=" * 70)


# ============================================================================
# DEMO
# ============================================================================

def setup():
    """Setup and index documents."""
    assistant = PersonalKnowledgeAssistant()
    assistant.index_my_knowledge()
    return assistant


def interactive():
    """Interactive mode."""
    assistant = setup()
    
    print("\n🧠 Personal Knowledge Assistant - Interactive Mode")
    print("=" * 70)
    print("Ask questions about your documents. Type 'exit' to quit.")
    print("=" * 70)
    print()
    
    while True:
        try:
            question = input("You: ").strip()
            
            if question.lower() in ['exit', 'quit', 'q']:
                print("\n👋 Goodbye!")
                break
            
            if not question:
                continue
            
            result = assistant.query(question)
            assistant.display_result(result)
            print("\n")
            
        except KeyboardInterrupt:
            print("\n\n👋 Goodbye!")
            break


if __name__ == "__main__":
    interactive()