client/main.py

#!/usr/bin/env python3

import requests
import sys
from typing import Dict, Any

API_BASE = "http://127.0.0.1:8000/api/v1"

class EmailQueryCLI:
    def __init__(self):
        self.session = requests.Session()
        
    def check_connection(self) -> bool:
        """Check if API server is running"""
        try:
            response = self.session.get(f"{API_BASE}/health")
            response.raise_for_status()
            return True
        except:
            return False
    
    def pretty_print_email(self, email: Dict) -> str:
        """Format email for display"""
        return f"""
📧 {email['subject']}
   📅 {email['date']} {email['time']}
   💬 {email['content'][:200]}...
   🆔 {email['message_id'][:20]}...
{"─" * 60}"""

    def handle_query(self, query: str):
        """Handle a natural language query"""
        print(f"\n🔍 Processing: '{query}'")
        
        try:
            # Try to get emails directly
            response = self.session.post(
                f"{API_BASE}/get_emails",
                json={"query": query}
            )
            
            if response.status_code == 200:
                data = response.json()
                self.display_email_results(data)
                return True
                
            elif response.status_code == 400:
                error_detail = response.json()["detail"]
                
                # Check if we need email mapping
                if isinstance(error_detail, dict) and error_detail.get("type") == "need_email_input":
                    mapping_success = self.handle_missing_mapping(error_detail)
                    if mapping_success and hasattr(self, '_retry_query'):
                        # Retry the query after successful mapping
                        print(f"🔄 Retrying query...")
                        delattr(self, '_retry_query')
                        return self.handle_query(query)  # Recursive retry
                    return mapping_success
                else:
                    print(f"❌ Error: {error_detail}")
                    return False
            else:
                print(f"❌ API Error: {response.status_code}")
                return False
                
        except Exception as e:
            print(f"❌ Connection Error: {e}")
            return False
    
    def handle_missing_mapping(self, error_detail: Dict) -> bool:
        """Handle case where email mapping is needed"""
        sender_intent = error_detail["sender_intent"]
        print(f"\n❓ {error_detail['message']}")
        
        try:
            email = input(f"📧 Enter email for '{sender_intent}': ").strip()
            if not email or "@" not in email:
                print("❌ Invalid email address")
                return False
            
            # Add the mapping
            mapping_response = self.session.post(
                f"{API_BASE}/add_email_mapping",
                json={"name": sender_intent, "email": email}
            )
            
            if mapping_response.status_code == 200:
                print(f"✅ Mapping saved: '{sender_intent}' → '{email}'")
                self._retry_query = True  # Flag to retry the original query
                return True
            else:
                print(f"❌ Failed to save mapping: {mapping_response.text}")
                return False
                
        except KeyboardInterrupt:
            print("\n❌ Cancelled")
            return False
    
    def display_email_results(self, data: Dict):
        """Display email search results"""
        print(f"\n✅ Found {data['total_emails']} emails")
        print(f"📤 From: {data['resolved_email']}")
        print(f"📅 Period: {data['start_date']} to {data['end_date']}")
        
        if data['emails']:
            print(f"\n📧 Emails:")
            for email in data['emails'][:10]:  # Show first 10
                print(self.pretty_print_email(email))
                
            if len(data['emails']) > 10:
                print(f"\n... and {len(data['emails']) - 10} more emails")
        else:
            print("\n📭 No emails found in this date range")
    
    def show_mappings(self):
        """Display all stored name-to-email mappings"""
        try:
            response = self.session.get(f"{API_BASE}/view_mappings")
            if response.status_code == 200:
                data = response.json()
                mappings = data["mappings"]
                
                print(f"\n📇 Stored Mappings ({len(mappings)}):")
                if mappings:
                    for name, email in mappings.items():
                        print(f"  👤 {name} → 📧 {email}")
                else:
                    print("  (No mappings stored)")
            else:
                print(f"❌ Failed to load mappings: {response.text}")
        except Exception as e:
            print(f"❌ Error: {e}")
    
    def run(self):
        """Main CLI loop"""
        if not self.check_connection():
            print("❌ Cannot connect to API server at http://127.0.0.1:8000")
            print("   Make sure to run: uvicorn main:app --reload")
            sys.exit(1)
        
        print("✅ Connected to Email Query System")
        print("💡 Try queries like:")
        print("   • 'emails from john last week'")
        print("   • 'show amazon emails from last month'")
        print("   • 'get dev@iitj.ac.in emails yesterday'")
        print("\n📋 Commands:")
        print("   • 'mappings' - View stored name-to-email mappings")
        print("   • 'quit' or Ctrl+C - Exit")
        print("=" * 60)
        
        while True:
            try:
                query = input("\n🗨️  You: ").strip()
                
                if not query:
                    continue
                    
                if query.lower() in ['quit', 'exit', 'q']:
                    break
                elif query.lower() in ['mappings', 'map', 'm']:
                    self.show_mappings()
                elif query.lower() in ['help', 'h']:
                    print("\n💡 Examples:")
                    print("   • emails from amazon last 5 days")
                    print("   • show john smith emails this week") 
                    print("   • get notifications from google yesterday")
                else:
                    self.handle_query(query)
                        
            except KeyboardInterrupt:
                break
            except Exception as e:
                print(f"❌ Unexpected error: {e}")
        
        print("\n👋 Goodbye!")

def main():
    """Entry point for CLI"""
    cli = EmailQueryCLI()
    cli.run()

if __name__ == "__main__":
    main()