#!/usr/bin/env python3
"""
Email Conversation Bot for OpenClaw
Monitors inbox and handles two-way email conversations
"""

import json
import imaplib
import smtplib
import email
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from datetime import datetime, timedelta
import os
import sys
import time
import re
import sqlite3
from pathlib import Path

# Database for conversation tracking
DB_PATH = "/root/.openclaw/credentials/email_conversations.db"
CREDENTIALS_FILE = "/root/.openclaw/credentials/gmail.json"
LOG_FILE = "/var/log/email_bot.log"

class EmailConversationBot:
    def __init__(self):
        self.creds = self._load_credentials()
        self._init_database()
        
    def _load_credentials(self):
        with open(CREDENTIALS_FILE, 'r') as f:
            return json.load(f)
    
    def _init_database(self):
        """Initialize SQLite database for tracking conversations"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS conversations (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                thread_id TEXT UNIQUE,
                sender_email TEXT,
                sender_name TEXT,
                subject TEXT,
                last_message TEXT,
                last_received TIMESTAMP,
                auto_reply_enabled INTEGER DEFAULT 0,
                context TEXT,
                message_count INTEGER DEFAULT 0
            )
        ''')
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS message_history (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                conversation_id INTEGER,
                direction TEXT,  -- 'incoming' or 'outgoing'
                content TEXT,
                timestamp TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                FOREIGN KEY (conversation_id) REFERENCES conversations(id)
            )
        ''')
        
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS auto_reply_rules (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                pattern TEXT,
                response_template TEXT,
                sender_whitelist TEXT,
                enabled INTEGER DEFAULT 1
            )
        ''')
        
        conn.commit()
        conn.close()
        
    def log(self, message):
        """Log activity"""
        timestamp = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
        log_entry = f"[{timestamp}] {message}"
        print(log_entry)
        
        with open(LOG_FILE, 'a') as f:
            f.write(log_entry + '\n')
    
    def connect_imap(self):
        """Connect to Gmail IMAP"""
        mail = imaplib.IMAP4_SSL(self.creds['imap_server'], self.creds['imap_port'])
        mail.login(self.creds['email'], self.creds['app_password'])
        return mail
    
    def connect_smtp(self):
        """Connect to Gmail SMTP"""
        server = smtplib.SMTP(self.creds['smtp_server'], self.creds['smtp_port'])
        server.starttls()
        server.login(self.creds['email'], self.creds['app_password'])
        return server
    
    def get_email_body(self, msg):
        """Extract text body from email"""
        body = ""
        if msg.is_multipart():
            for part in msg.walk():
                if part.get_content_type() == "text/plain":
                    try:
                        body = part.get_payload(decode=True).decode('utf-8', errors='ignore')
                        break
                    except:
                        pass
        else:
            try:
                body = msg.get_payload(decode=True).decode('utf-8', errors='ignore')
            except:
                pass
        return body
    
    def check_new_emails(self, mark_as_read=False):
        """Check for new unread emails"""
        mail = self.connect_imap()
        mail.select('inbox')
        
        # Search for unread emails from last 24 hours
        date_str = (datetime.now() - timedelta(days=1)).strftime('%d-%b-%Y')
        status, messages = mail.search(None, f'(UNSEEN SINCE {date_str})')
        
        new_emails = []
        
        if status == 'OK' and messages[0]:
            email_ids = messages[0].split()
            
            for e_id in email_ids:
                status, msg_data = mail.fetch(e_id, '(RFC822)')
                if status == 'OK':
                    raw_email = msg_data[0][1]
                    msg = email.message_from_bytes(raw_email)
                    
                    # Skip emails from myself
                    sender = msg['From']
                    if self.creds['email'] in sender:
                        continue
                    
                    email_data = {
                        'id': e_id.decode(),
                        'sender': sender,
                        'sender_email': self._extract_email(sender),
                        'sender_name': self._extract_name(sender),
                        'subject': msg['Subject'] or '(No Subject)',
                        'date': msg['Date'],
                        'body': self.get_email_body(msg),
                        'message_id': msg['Message-ID'],
                        'references': msg.get('References', ''),
                        'in_reply_to': msg.get('In-Reply-To', '')
                    }
                    
                    new_emails.append(email_data)
                    
                    # Mark as read if requested
                    if mark_as_read:
                        mail.store(e_id, '+FLAGS', '\\Seen')
        
        mail.logout()
        return new_emails
    
    def _extract_email(self, from_field):
        """Extract email address from 'From' field"""
        match = re.search(r'<([^>]+)>', from_field)
        if match:
            return match.group(1)
        return from_field
    
    def _extract_name(self, from_field):
        """Extract name from 'From' field"""
        match = re.search(r'^"?([^"<]+)"?\s*<', from_field)
        if match:
            return match.group(1).strip()
        return from_field
    
    def get_or_create_conversation(self, sender_email, subject, thread_id=None):
        """Get existing conversation or create new one"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        # Try to find existing conversation
        cursor.execute(
            "SELECT id, auto_reply_enabled FROM conversations WHERE sender_email = ?",
            (sender_email,)
        )
        result = cursor.fetchone()
        
        if result:
            conn.close()
            return {'id': result[0], 'auto_reply_enabled': result[1]}
        
        # Create new conversation
        cursor.execute(
            """INSERT INTO conversations 
               (thread_id, sender_email, sender_name, subject, last_received, message_count)
               VALUES (?, ?, ?, ?, ?, ?)""",
            (thread_id or '', sender_email, '', subject, datetime.now(), 0)
        )
        
        conversation_id = cursor.lastrowid
        conn.commit()
        conn.close()
        
        self.log(f"New conversation started with {sender_email}")
        return {'id': conversation_id, 'auto_reply_enabled': 0}
    
    def store_message(self, conversation_id, direction, content):
        """Store message in history"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute(
            "INSERT INTO message_history (conversation_id, direction, content) VALUES (?, ?, ?)",
            (conversation_id, direction, content)
        )
        
        # Update conversation
        cursor.execute(
            "UPDATE conversations SET last_message = ?, message_count = message_count + 1 WHERE id = ?",
            (content[:200], conversation_id)
        )
        
        conn.commit()
        conn.close()
    
    def send_reply(self, to_email, subject, body, in_reply_to=None, references=None):
        """Send a reply email"""
        # Preview mode by default
        print("\n" + "="*60)
        print("📧 REPLY PREVIEW")
        print("="*60)
        print(f"To: {to_email}")
        print(f"Subject: Re: {subject}")
        print("-"*60)
        print(body)
        print("="*60)
        print("\n⚠️  This is a preview. Use --send flag to actually send.\n")
        return True
    
    def send_reply_actual(self, to_email, subject, body, in_reply_to=None, references=None):
        """Actually send the reply"""
        try:
            server = self.connect_smtp()
            
            msg = MIMEMultipart()
            msg['From'] = self.creds['email']
            msg['To'] = to_email
            msg['Subject'] = f"Re: {subject}" if not subject.startswith('Re:') else subject
            msg['Date'] = datetime.now().strftime('%a, %d %b %Y %H:%M:%S +0000')
            
            if in_reply_to:
                msg['In-Reply-To'] = in_reply_to
            if references:
                msg['References'] = references
            
            msg.attach(MIMEText(body, 'plain'))
            
            server.send_message(msg)
            server.quit()
            
            self.log(f"✅ Reply sent to {to_email}")
            return True
            
        except Exception as e:
            self.log(f"❌ Failed to send reply: {e}")
            return False
    
    def generate_reply(self, email_data, conversation_history):
        """Generate an appropriate reply (customize this!)"""
        sender_name = email_data['sender_name'] or 'there'
        
        # This is where you'd integrate with AI for smart replies
        # For now, using a simple acknowledgment template
        
        reply = f"""Hi {sender_name},

Thanks for your email! I've received your message and will get back to you soon.

Best regards,
Schlaus

---
This is an automated acknowledgment. I'll respond personally shortly.
"""
        return reply
    
    def process_new_emails(self, auto_reply=False, dry_run=True):
        """Main loop - check and process new emails"""
        self.log("Checking for new emails...")
        
        new_emails = self.check_new_emails(mark_as_read=False)
        
        if not new_emails:
            self.log("No new emails")
            return
        
        self.log(f"Found {len(new_emails)} new email(s)")
        
        for email_data in new_emails:
            sender = email_data['sender_email']
            subject = email_data['subject']
            
            self.log(f"📧 From: {sender} | Subject: {subject}")
            
            # Get or create conversation
            conversation = self.get_or_create_conversation(
                sender, subject, email_data.get('message_id')
            )
            
            # Store incoming message
            self.store_message(conversation['id'], 'incoming', email_data['body'])
            
            # Check if auto-reply is enabled for this conversation
            if auto_reply and conversation['auto_reply_enabled']:
                reply_body = self.generate_reply(email_data, conversation)
                
                if dry_run:
                    self.send_reply(
                        sender, subject, reply_body,
                        email_data.get('message_id'),
                        email_data.get('references')
                    )
                else:
                    success = self.send_reply_actual(
                        sender, subject, reply_body,
                        email_data.get('message_id'),
                        email_data.get('references')
                    )
                    
                    if success:
                        self.store_message(conversation['id'], 'outgoing', reply_body)
    
    def enable_auto_reply(self, sender_email):
        """Enable auto-reply for a specific sender"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute(
            "UPDATE conversations SET auto_reply_enabled = 1 WHERE sender_email = ?",
            (sender_email,)
        )
        
        conn.commit()
        conn.close()
        self.log(f"✅ Auto-reply enabled for {sender_email}")
    
    def list_conversations(self):
        """List all tracked conversations"""
        conn = sqlite3.connect(DB_PATH)
        cursor = conn.cursor()
        
        cursor.execute(
            "SELECT sender_email, subject, message_count, auto_reply_enabled, last_received FROM conversations ORDER BY last_received DESC"
        )
        
        conversations = cursor.fetchall()
        conn.close()
        
        if not conversations:
            print("📭 No conversations yet")
            return
        
        print(f"\n📬 {len(conversations)} Conversation(s):\n")
        for conv in conversations:
            email, subject, count, auto_reply, last = conv
            status = "🤖 Auto" if auto_reply else "👤 Manual"
            print(f"{status} | {email}")
            print(f"     Subject: {subject}")
            print(f"     Messages: {count} | Last: {last}")
            print()

def main():
    bot = EmailConversationBot()
    
    if len(sys.argv) < 2:
        print("Email Conversation Bot for OpenClaw")
        print("\nUsage:")
        print("  email_bot check              - Check new emails (preview mode)")
        print("  email_bot check --send       - Check and auto-reply (actually sends!)")
        print("  email_bot conversations      - List all conversations")
        print("  email_bot enable <email>     - Enable auto-reply for sender")
        print("  email_bot monitor            - Continuous monitoring mode")
        print("  email_bot reply <email_id>   - Reply to specific email")
        sys.exit(1)
    
    command = sys.argv[1]
    
    if command == "check":
        dry_run = "--send" not in sys.argv
        if not dry_run:
            print("⚠️  WARNING: This will actually send emails!")
            print("Press Ctrl+C within 3 seconds to cancel...")
            time.sleep(3)
        bot.process_new_emails(auto_reply=True, dry_run=dry_run)
    
    elif command == "conversations":
        bot.list_conversations()
    
    elif command == "enable" and len(sys.argv) > 2:
        bot.enable_auto_reply(sys.argv[2])
    
    elif command == "monitor":
        print("🔄 Starting email monitor (Ctrl+C to stop)...")
        while True:
            try:
                bot.process_new_emails(auto_reply=True, dry_run=True)
                time.sleep(60)  # Check every minute
            except KeyboardInterrupt:
                print("\n👋 Stopping monitor")
                break
    
    else:
        print(f"Unknown command: {command}")

if __name__ == "__main__":
    main()