#!/usr/bin/env python3
"""
Complete RAG pipeline: Document ingestion, chunking, embedding, and indexing.
"""

import argparse
from pathlib import Path
from typing import List
import json
from pygeai import Client


class RAGPipeline:
    def __init__(self, client: Client, assistant_id: str):
        self.client = client
        self.assistant_id = assistant_id
    
    def ingest_documents(self, directory: Path, recursive: bool = True):
        """Ingest all documents from a directory."""
        supported_formats = {".pdf", ".txt", ".md", ".docx", ".html"}
        
        if recursive:
            files = list(directory.rglob("*"))
        else:
            files = list(directory.glob("*"))
        
        documents = [f for f in files if f.suffix.lower() in supported_formats and f.is_file()]
        
        print(f"📂 Found {len(documents)} documents")
        
        # Upload in batches
        batch_size = 10
        uploaded = 0
        
        for i in range(0, len(documents), batch_size):
            batch = documents[i:i + batch_size]
            
            try:
                self.client.rag.upload_documents(
                    assistant_id=self.assistant_id,
                    files=[str(f) for f in batch]
                )
                uploaded += len(batch)
                print(f"  ✓ Uploaded {uploaded}/{len(documents)}")
            except Exception as e:
                print(f"  ✗ Error uploading batch: {e}")
        
        return uploaded
    
    def query_with_evaluation(self, queries: List[str]):
        """Query RAG and evaluate responses."""
        results = []
        
        for query in queries:
            print(f"\n❓ Query: {query}")
            
            try:
                response = self.client.rag.query(
                    assistant_id=self.assistant_id,
                    query=query,
                    include_sources=True
                )
                
                print(f"📝 Answer: {response.answer[:200]}...")
                print(f"📚 Sources: {len(response.sources)}")
                
                results.append({
                    "query": query,
                    "answer": response.answer,
                    "num_sources": len(response.sources),
                    "sources": [
                        {
                            "file": s.filename,
                            "score": s.score
                        }
                        for s in response.sources
                    ]
                })
            except Exception as e:
                print(f"✗ Error: {e}")
                results.append({"query": query, "error": str(e)})
        
        return results
    
    def cleanup(self):
        """Remove all documents from RAG assistant."""
        print("🗑️  Cleaning up...")
        try:
            self.client.rag.delete_all_documents(assistant_id=self.assistant_id)
            print("✓ All documents removed")
        except Exception as e:
            print(f"✗ Error: {e}")


def main():
    parser = argparse.ArgumentParser(description="RAG Pipeline")
    parser.add_argument("--assistant-id", required=True, help="RAG assistant ID")
    parser.add_argument("--docs-dir", type=Path, help="Documents directory")
    parser.add_argument("--query-file", type=Path, help="File with test queries (one per line)")
    parser.add_argument("--cleanup", action="store_true", help="Clean up after testing")
    
    args = parser.parse_args()
    
    client = Client()
    pipeline = RAGPipeline(client, args.assistant_id)
    
    # Ingest documents
    if args.docs_dir:
        print("\n📥 INGESTING DOCUMENTS")
        print("=" * 60)
        pipeline.ingest_documents(args.docs_dir)
    
    # Test queries
    if args.query_file:
        print("\n\n🔍 TESTING QUERIES")
        print("=" * 60)
        
        with open(args.query_file) as f:
            queries = [line.strip() for line in f if line.strip()]
        
        results = pipeline.query_with_evaluation(queries)
        
        # Save results
        output_file = Path("rag_results.json")
        with open(output_file, "w") as f:
            json.dump(results, f, indent=2)
        
        print(f"\n\n💾 Results saved to {output_file}")
    
    # Cleanup
    if args.cleanup:
        print("\n\n🧹 CLEANUP")
        print("=" * 60)
        pipeline.cleanup()
    
    print("\n\n✅ Pipeline complete!")


if __name__ == "__main__":
    main()
