main.py

"""
Email RAG Assistant - Main Entry Point

A real-time agentic AI system that:
1. Monitors Gmail inbox for new emails
2. Converts emails to embeddings and stores them in Pinecone
3. Allows natural language queries against the email knowledge base

Usage:
    python main.py              # Run with interactive query mode
    python main.py --ingest     # Run ingestion only (once)
    python main.py --daemon     # Run as background daemon with scheduler
"""

import asyncio
import sys
import argparse
from datetime import datetime

from config import validate_config, INGESTION_INTERVAL_SECONDS
from ai_agents.gmail_ingestion_agent import GmailIngestionAgent
from ai_agents.embedding_agent import EmbeddingAgent
from ai_agents.rag_query_agent import RAGQueryAgent
from utils.scheduler import Scheduler, GracefulShutdown


class EmailRAGAssistant:
    """
    Main orchestrator for the Email RAG Assistant.
    Coordinates all agents and services.
    """
    
    def __init__(self):
        """Initialize the Email RAG Assistant."""
        self.gmail_agent = GmailIngestionAgent()
        self.embedding_agent = EmbeddingAgent()
        self.rag_agent = RAGQueryAgent()
        self.scheduler = None
        self._initialized = False
    
    async def initialize(self) -> bool:
        """
        Initialize all components.
        
        Returns:
            True if initialization successful.
        """
        print("\n" + "=" * 60)
        print("🚀 Email RAG Assistant - Initializing...")
        print("=" * 60)
        
        try:
            # Validate configuration
            validate_config()
            print("✓ Configuration validated")
            
            # Initialize Gmail authentication
            if not self.gmail_agent.authenticate():
                print("⚠ Gmail authentication pending - will prompt on first use")
            else:
                print("✓ Gmail authenticated")
            
            # Initialize embedding agent (loads model and Pinecone)
            if not self.embedding_agent.initialize():
                print("✗ Failed to initialize embedding agent")
                return False
            print("✓ Embedding agent initialized")
            
            # Initialize RAG agent
            if not self.rag_agent.initialize():
                print("✗ Failed to initialize RAG agent")
                return False
            print("✓ RAG query agent initialized")
            
            self._initialized = True
            print("\n✓ All components initialized successfully!")
            print("=" * 60)
            
            return True
            
        except Exception as e:
            print(f"\n✗ Initialization error: {e}")
            return False
    
    async def run_ingestion(self) -> int:
        """
        Run a single email ingestion cycle.
        
        Returns:
            Number of emails processed.
        """
        print(f"\n[{datetime.now().strftime('%H:%M:%S')}] Starting email ingestion...")
        
        try:
            # Check for new emails
            emails = await self.gmail_agent.check_inbox(max_emails=20)
            
            if not emails:
                print("No new emails found.")
                return 0
            
            print(f"Found {len(emails)} new email(s). Processing...")
            
            # Process emails through embedding agent
            processed = await self.embedding_agent.process_emails(emails)
            
            print(f"Successfully processed {processed}/{len(emails)} emails.")
            return processed
            
        except Exception as e:
            print(f"Error during ingestion: {e}")
            return 0
    
    async def query(self, question: str) -> str:
        """
        Query the email knowledge base.
        
        Args:
            question: The user's question.
            
        Returns:
            The assistant's answer.
        """
        return await self.rag_agent.query(question)
    
    async def interactive_mode(self):
        """Run the interactive query mode."""
        print("\n" + "=" * 60)
        print("📧 Email RAG Assistant - Interactive Mode")
        print("=" * 60)
        print("\nCommands:")
        print("  Type your question to search emails")
        print("  'ingest' - Run email ingestion")
        print("  'stats'  - Show system statistics")
        print("  'help'   - Show this help")
        print("  'quit'   - Exit the program")
        print("-" * 60)
        
        while True:
            try:
                print()
                user_input = input("You: ").strip()
                
                if not user_input:
                    continue
                
                command = user_input.lower()
                
                if command in ['quit', 'exit', 'q']:
                    print("\nGoodbye! 👋")
                    break
                
                elif command == 'help':
                    print("\nCommands:")
                    print("  - Type any question to search your emails")
                    print("  - 'ingest' - Check for new emails and process them")
                    print("  - 'stats'  - Show system statistics")
                    print("  - 'quit'   - Exit the program")
                
                elif command == 'ingest':
                    processed = await self.run_ingestion()
                    print(f"\nIngestion complete. Processed {processed} emails.")
                
                elif command == 'stats':
                    await self.show_stats()
                
                else:
                    # Treat as a question
                    print("\nSearching emails...")
                    response = await self.query(user_input)
                    print(f"\nAssistant: {response}")
                    
            except KeyboardInterrupt:
                print("\n\nInterrupted. Goodbye! 👋")
                break
            except EOFError:
                print("\n\nGoodbye! 👋")
                break
            except Exception as e:
                print(f"\nError: {e}")
    
    async def show_stats(self):
        """Display system statistics."""
        print("\n📊 System Statistics:")
        print("-" * 40)
        
        # Gmail stats
        gmail_processed = self.gmail_agent.get_processed_count()
        print(f"Emails processed (this session): {gmail_processed}")
        
        # Embedding/Pinecone stats
        pinecone_stats = self.embedding_agent.get_stats()
        print(f"Total emails in knowledge base: {pinecone_stats.get('total_vectors', 'N/A')}")
        
        # Scheduler stats if running
        if self.scheduler:
            status = self.scheduler.get_status()
            print(f"Scheduler running: {status['running']}")
            print(f"Ingestion interval: {status['interval_seconds']} seconds")
            print(f"Total ingestion runs: {status['tick_count']}")
    
    async def run_daemon(self):
        """Run as a background daemon with scheduled ingestion."""
        print("\n" + "=" * 60)
        print("📧 Email RAG Assistant - Daemon Mode")
        print("=" * 60)
        print(f"\nRunning email ingestion every {INGESTION_INTERVAL_SECONDS} seconds.")
        print("Press Ctrl+C to stop.\n")
        
        # Setup scheduler
        self.scheduler = Scheduler(
            interval_seconds=INGESTION_INTERVAL_SECONDS,
            on_tick=self.run_ingestion
        )
        
        # Setup graceful shutdown
        shutdown = GracefulShutdown()
        shutdown.register_callback(self.scheduler.stop)
        
        try:
            # Run initial ingestion
            await self.run_ingestion()
            
            # Start scheduler
            await self.scheduler.start()
            
            # Keep running until interrupted
            while True:
                await asyncio.sleep(1)
                
        except KeyboardInterrupt:
            print("\n\nShutting down gracefully...")
        finally:
            await self.scheduler.stop()
            print("Daemon stopped.")


async def main_async(args):
    """Async main function."""
    assistant = EmailRAGAssistant()
    
    # Initialize the system
    if not await assistant.initialize():
        print("\nFailed to initialize. Please check your configuration.")
        sys.exit(1)
    
    # Run based on mode
    if args.ingest:
        # Single ingestion run
        processed = await assistant.run_ingestion()
        print(f"\nIngestion complete. Processed {processed} emails.")
        
    elif args.daemon:
        # Daemon mode with scheduler
        await assistant.run_daemon()
        
    else:
        # Interactive mode (default)
        await assistant.interactive_mode()


def main():
    """Main entry point."""
    parser = argparse.ArgumentParser(
        description="Email RAG Assistant - Query your emails with AI"
    )
    parser.add_argument(
        "--ingest",
        action="store_true",
        help="Run email ingestion once and exit"
    )
    parser.add_argument(
        "--daemon",
        action="store_true",
        help="Run as daemon with scheduled ingestion"
    )
    
    args = parser.parse_args()
    
    # Run async main
    try:
        asyncio.run(main_async(args))
    except KeyboardInterrupt:
        print("\nInterrupted.")
        sys.exit(0)


if __name__ == "__main__":
    main()