#!/usr/bin/env python3
"""
Spam Manager UI - Web interface for managing spam emails.
"""
import os
import sys
import email
import base64
from datetime import datetime, timedelta
from pathlib import Path
from typing import List, Dict
from io import BytesIO

from flask import Flask, render_template_string, request, jsonify
from dotenv import load_dotenv

try:
    from playwright.sync_api import sync_playwright
    PLAYWRIGHT_AVAILABLE = True
except ImportError:
    PLAYWRIGHT_AVAILABLE = False

from llmass.imap import connect_imap
from llmass.organizer.parser import get_email_content, detect_imap_server

load_dotenv()

app = Flask(__name__)

# Global IMAP connection
imap_client = None
imap_session = None
email_address = None
password = None
imap_server = None

HTML_TEMPLATE = """
<!DOCTYPE html>
<html>
<head>
    <title>Spam Manager - LLMass</title>
    <meta charset="utf-8">
    <meta name="viewport" content="width=device-width, initial-scale=1">
    <style>
        * { box-sizing: border-box; margin: 0; padding: 0; }
        body {
            font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, Oxygen, Ubuntu, sans-serif;
            background: #f5f5f5;
            padding: 20px;
        }
        .container { max-width: 1400px; margin: 0 auto; }
        h1 { color: #333; margin-bottom: 20px; }
        .stats {
            background: white;
            padding: 20px;
            border-radius: 8px;
            margin-bottom: 20px;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
        }
        .stats-grid {
            display: grid;
            grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
            gap: 20px;
        }
        .stat-item {
            text-align: center;
        }
        .stat-value {
            font-size: 32px;
            font-weight: bold;
            color: #007bff;
        }
        .stat-label {
            color: #666;
            margin-top: 5px;
        }
        .email-grid {
            display: grid;
            grid-template-columns: repeat(auto-fill, minmax(350px, 1fr));
            gap: 20px;
        }
        .email-card {
            background: white;
            border-radius: 8px;
            overflow: hidden;
            box-shadow: 0 2px 4px rgba(0,0,0,0.1);
            transition: transform 0.2s;
        }
        .email-card:hover {
            transform: translateY(-5px);
            box-shadow: 0 4px 8px rgba(0,0,0,0.2);
        }
        .email-preview {
            width: 100%;
            height: 300px;
            border: none;
            background: #fff;
        }
        .email-screenshot {
            width: 100%;
            height: 300px;
            object-fit: cover;
            border-bottom: 1px solid #eee;
        }
        .email-info {
            padding: 15px;
        }
        .email-from {
            font-weight: bold;
            color: #333;
            margin-bottom: 5px;
            font-size: 14px;
            white-space: nowrap;
            overflow: hidden;
            text-overflow: ellipsis;
        }
        .email-subject {
            color: #666;
            margin-bottom: 10px;
            font-size: 13px;
            white-space: nowrap;
            overflow: hidden;
            text-overflow: ellipsis;
        }
        .email-date {
            color: #999;
            font-size: 12px;
            margin-bottom: 15px;
        }
        .email-actions {
            display: flex;
            gap: 10px;
        }
        .btn {
            flex: 1;
            padding: 10px;
            border: none;
            border-radius: 5px;
            cursor: pointer;
            font-weight: 500;
            transition: background 0.2s;
            font-size: 13px;
        }
        .btn-delete {
            background: #dc3545;
            color: white;
        }
        .btn-delete:hover {
            background: #c82333;
        }
        .btn-archive {
            background: #ffc107;
            color: #000;
        }
        .btn-archive:hover {
            background: #e0a800;
        }
        .btn-inbox {
            background: #28a745;
            color: white;
        }
        .btn-inbox:hover {
            background: #218838;
        }
        .loading {
            text-align: center;
            padding: 40px;
            color: #666;
        }
        .refresh-btn {
            background: #007bff;
            color: white;
            border: none;
            padding: 10px 20px;
            border-radius: 5px;
            cursor: pointer;
            font-size: 14px;
            margin-bottom: 20px;
        }
        .refresh-btn:hover {
            background: #0056b3;
        }
        .processing {
            opacity: 0.5;
            pointer-events: none;
        }
    </style>
</head>
<body>
    <div class="container">
        <h1>📧 Spam Manager</h1>
        
        <div class="stats">
            <div class="stats-grid">
                <div class="stat-item">
                    <div class="stat-value" id="total-count">-</div>
                    <div class="stat-label">Total Spam Emails</div>
                </div>
                <div class="stat-item">
                    <div class="stat-value" id="processed-count">0</div>
                    <div class="stat-label">Processed</div>
                </div>
            </div>
        </div>
        
        <button class="refresh-btn" onclick="refreshEmails()">🔄 Refresh</button>
        
        <div class="email-grid" id="email-grid">
            <div class="loading">Loading spam emails...</div>
        </div>
    </div>
    
    <script>
        let emails = [];
        let processedCount = 0;
        
        async function loadEmails() {
            try {
                const response = await fetch('/api/spam');
                const data = await response.json();
                emails = data.emails;
                document.getElementById('total-count').textContent = emails.length;
                renderEmails();
            } catch (error) {
                console.error('Error loading emails:', error);
                document.getElementById('email-grid').innerHTML = 
                    '<div class="loading">Error loading emails</div>';
            }
        }
        
        function renderEmails() {
            const grid = document.getElementById('email-grid');
            if (emails.length === 0) {
                grid.innerHTML = '<div class="loading">No spam emails found! 🎉</div>';
                return;
            }
            
            grid.innerHTML = emails.map((email, index) => `
                <div class="email-card" id="email-${index}">
                    ${email.screenshot ? 
                        `<img src="data:image/png;base64,${email.screenshot}" class="email-screenshot" alt="Email preview">` :
                        `<iframe srcdoc="${escapeHtmlForAttr(email.html)}" class="email-preview" sandbox="" referrerpolicy="no-referrer"></iframe>`
                    }
                    <div class="email-info">
                        <div class="email-from" title="${escapeHtmlForAttr(email.from)}">${escapeHtml(email.from)}</div>
                        <div class="email-subject" title="${escapeHtmlForAttr(email.subject)}">${escapeHtml(email.subject)}</div>
                        <div class="email-date">${escapeHtml(email.date)}</div>
                        <div class="email-actions">
                            <button class="btn btn-delete" onclick="performAction('${escapeJs(email.uid)}', 'delete', ${index})">
                                🗑️ Delete
                            </button>
                            <button class="btn btn-archive" onclick="performAction('${escapeJs(email.uid)}', 'archive', ${index})">
                                📦 Archive
                            </button>
                            <button class="btn btn-inbox" onclick="performAction('${escapeJs(email.uid)}', 'inbox', ${index})">
                                📥 To Inbox
                            </button>
                        </div>
                    </div>
                </div>
            `).join('');
        }
        
        async function performAction(uid, action, index) {
            const card = document.getElementById(`email-${index}`);
            card.classList.add('processing');
            
            try {
                const response = await fetch('/api/action', {
                    method: 'POST',
                    headers: { 'Content-Type': 'application/json' },
                    body: JSON.stringify({ uid, action })
                });
                
                const result = await response.json();
                if (result.success) {
                    card.remove();
                    processedCount++;
                    document.getElementById('processed-count').textContent = processedCount;
                    document.getElementById('total-count').textContent = 
                        parseInt(document.getElementById('total-count').textContent) - 1;
                } else {
                    alert('Error: ' + result.message);
                    card.classList.remove('processing');
                }
            } catch (error) {
                console.error('Error performing action:', error);
                alert('Error performing action');
                card.classList.remove('processing');
            }
        }
        
        function escapeHtml(text) {
            if (!text) return '';
            const div = document.createElement('div');
            div.textContent = text;
            return div.innerHTML;
        }
        
        function escapeHtmlForAttr(text) {
            if (!text) return '';
            return text
                .replace(/&/g, '&amp;')
                .replace(/"/g, '&quot;')
                .replace(/'/g, '&#39;')
                .replace(/</g, '&lt;')
                .replace(/>/g, '&gt;')
                .replace(/\n/g, '&#10;')
                .replace(/\r/g, '&#13;');
        }
        
        function escapeJs(text) {
            if (!text) return '';
            return text.replace(/'/g, "\\'").replace(/"/g, '\\"').replace(/\n/g, '\\n').replace(/\r/g, '\\r');
        }
        
        function refreshEmails() {
            document.getElementById('email-grid').innerHTML = '<div class="loading">Loading...</div>';
            processedCount = 0;
            document.getElementById('processed-count').textContent = '0';
            loadEmails();
        }
        
        // Load emails on page load
        loadEmails();
    </script>
</body>
</html>
"""


def connect():
    """Connect to IMAP server."""
    global imap_client, imap_session
    
    try:
        imap_client, imap_session = connect_imap(
            imap_server, email_address, password,
            ssl=True, retries=2, backoff=0.5, verbose=False
        )
        return True
    except Exception as e:
        print(f"Error connecting: {e}")
        return False


# Global cache for detected spam folder
_detected_spam_folder = None

def detect_spam_folder():
    """Detect spam folder name. Returns None if not found."""
    global _detected_spam_folder
    
    # Return cached value if available
    if _detected_spam_folder:
        return _detected_spam_folder
    
    spam_folder_names = [
        'SPAM', 'Spam', 'spam',
        'Junk', 'JUNK', 'junk',
        'Junk Email', 'Junk E-mail',
        'Bulk', 'BULK',
        'INBOX.Spam', 'INBOX.Junk',
        '[Gmail]/Spam',
    ]
    
    try:
        result, folders = imap_client.safe_list('', '*')
        if result != 'OK' or not folders:
            print("Warning: Could not list folders")
            return None
        
        available_folders = []
        for folder_line in folders:
            if not folder_line:
                continue
            parts = folder_line.decode() if isinstance(folder_line, bytes) else folder_line
            import re
            match = re.search(r'"([^"]+)"$', parts)
            if match:
                available_folders.append(match.group(1))
            else:
                words = parts.split()
                if words:
                    available_folders.append(words[-1])
        
        # Check exact matches
        for spam_name in spam_folder_names:
            if spam_name in available_folders:
                _detected_spam_folder = spam_name
                print(f"Detected spam folder: {spam_name}")
                return spam_name
        
        # Check case-insensitive
        for folder in available_folders:
            folder_lower = folder.lower()
            if 'spam' in folder_lower or 'junk' in folder_lower or 'bulk' in folder_lower:
                _detected_spam_folder = folder
                print(f"Detected spam folder: {folder}")
                return folder
        
        print("Warning: No spam folder found, using INBOX")
        _detected_spam_folder = 'INBOX'
        return 'INBOX'  # Fallback to INBOX if no spam folder found
        
    except Exception as e:
        print(f"Error detecting spam folder: {e}")
        return None


def get_spam_emails(limit=50):
    """Get spam emails with HTML preview or screenshot."""
    if not imap_client:
        return []
    
    try:
        spam_folder = detect_spam_folder()
        if not spam_folder:
            print("Error: No spam folder detected")
            return []
        
        print(f"Using spam folder: {spam_folder}")
        
        # Wrap SELECT in try-except for multiline response issues
        try:
            result, data = imap_client.safe_select(spam_folder, readonly=False)
        except Exception as select_error:
            if 'unexpected response' in str(select_error).lower():
                print(f"Error selecting folder (IMAP parsing issue): {spam_folder}")
                return []
            raise
        
        if result != 'OK':
            print(f"Error selecting folder {spam_folder}: {data}")
            return []
        
        result, data = imap_client.safe_uid('SEARCH', None, 'ALL')
        if result != 'OK' or not data or not data[0]:
            return []
        
        email_uids = data[0].split()[-limit:]  # Last N emails
        emails = []
        
        for uid in email_uids:
            try:
                # Some IMAP servers have issues with multiline responses
                # Wrap in try-except to handle "unexpected response" errors
                try:
                    result, data = imap_client.safe_uid('FETCH', uid, '(RFC822)')
                except Exception as fetch_error:
                    # Skip emails that cause IMAP parsing errors
                    if 'unexpected response' in str(fetch_error).lower():
                        continue
                    raise  # Re-raise other errors
                
                if result != 'OK' or not data:
                    continue
                
                # Parse IMAP response - FETCH returns data in specific format
                # Format: [(b'UID FLAGS...', b'email content'), b')'] 
                # Or: [(b'UID FLAGS... {size}', b'email content')]
                raw_email = None
                
                # Standard IMAP FETCH response has tuple as first element
                if data and isinstance(data[0], tuple):
                    # First tuple contains (metadata, content)
                    if len(data[0]) >= 2:
                        raw_email = data[0][1]
                
                # Validate we got bytes that look like an email
                if not raw_email or not isinstance(raw_email, bytes):
                    continue
                
                # Additional validation - check if it starts with email headers
                if not (raw_email.startswith(b'From:') or 
                        raw_email.startswith(b'Return-Path:') or
                        raw_email.startswith(b'Received:') or
                        raw_email.startswith(b'DKIM-Signature:') or
                        raw_email.startswith(b'ARC-')):
                    continue
                
                msg = email.message_from_bytes(raw_email)
                content = get_email_content(msg)
                
                # Get HTML body
                html_body = content.get('html_body', '')
                if not html_body:
                    html_body = f'<pre>{content.get("body", "No content")}</pre>'
                
                # Generate screenshot if Playwright available
                screenshot = None
                if PLAYWRIGHT_AVAILABLE and html_body:
                    screenshot = generate_screenshot(html_body)
                
                emails.append({
                    'uid': uid.decode() if isinstance(uid, bytes) else uid,
                    'from': content.get('from', 'Unknown'),
                    'subject': content.get('subject', 'No subject'),
                    'date': content.get('date', ''),
                    'html': html_body[:10000],  # Limit size
                    'screenshot': screenshot
                })
            
            except Exception as e:
                print(f"Error processing email {uid}: {e}")
                continue
        
        return emails
    
    except Exception as e:
        print(f"Error getting spam emails: {e}")
        return []


def generate_screenshot(html):
    """Generate screenshot of HTML email."""
    try:
        with sync_playwright() as p:
            browser = p.chromium.launch(headless=True)
            page = browser.new_page(viewport={'width': 600, 'height': 800})
            page.set_content(html)
            
            # Take screenshot
            screenshot_bytes = page.screenshot(full_page=True, type='png')
            browser.close()
            
            # Convert to base64
            return base64.b64encode(screenshot_bytes).decode('utf-8')
    except Exception as e:
        print(f"Error generating screenshot: {e}")
        return None


@app.route('/')
def index():
    """Main page."""
    return render_template_string(HTML_TEMPLATE)


@app.route('/api/spam')
def api_spam():
    """API endpoint to get spam emails."""
    global imap_client, imap_session
    
    # Reconnect if connection lost
    if not imap_client or not imap_session:
        if not connect():
            return jsonify({'emails': [], 'error': 'Connection lost'}), 500
    
    try:
        emails = get_spam_emails(limit=50)
        return jsonify({'emails': emails})
    except Exception as e:
        print(f"Error in api_spam: {e}")
        # Try to reconnect
        if connect():
            try:
                emails = get_spam_emails(limit=50)
                return jsonify({'emails': emails})
            except Exception as e2:
                print(f"Error after reconnect: {e2}")
                return jsonify({'emails': [], 'error': str(e2)}), 500
        return jsonify({'emails': [], 'error': str(e)}), 500


@app.route('/api/action', methods=['POST'])
def api_action():
    """API endpoint to perform action on email."""
    data = request.json
    uid = data.get('uid')
    action = data.get('action')
    
    if not uid or not action:
        return jsonify({'success': False, 'message': 'Missing parameters'})
    
    try:
        # Convert uid to bytes if needed
        if isinstance(uid, str):
            uid = uid.encode()
        
        if action == 'delete':
            # Mark as deleted and expunge
            result, _ = imap_client.safe_uid('STORE', uid, '+FLAGS', '(\\Deleted)')
            if result == 'OK':
                imap_session.expunge()
                return jsonify({'success': True})
        
        elif action == 'archive':
            # Move to Archive folder (create if doesn't exist)
            archive_folder = 'INBOX.Archive'  # Standard hierarchical name
            
            # Check if Archive folder exists
            try:
                result, folders = imap_client.safe_list('', '*Archive*')
                folder_exists = False
                
                if result == 'OK' and folders:
                    # Parse folder list to see if Archive exists
                    for folder_line in folders:
                        if not folder_line:
                            continue
                        folder_str = folder_line.decode() if isinstance(folder_line, bytes) else str(folder_line)
                        if 'Archive' in folder_str:
                            folder_exists = True
                            # Extract actual folder name
                            import re
                            match = re.search(r'"([^"]*Archive[^"]*)"', folder_str)
                            if match:
                                archive_folder = match.group(1)
                            break
                
                # Create Archive folder if it doesn't exist
                if not folder_exists:
                    try:
                        imap_client.safe_create(archive_folder)
                        imap_client.safe_subscribe(archive_folder)
                    except Exception as create_error:
                        # Try simpler name if hierarchical fails
                        if 'INBOX.' in archive_folder:
                            archive_folder = 'Archive'
                            imap_client.safe_create(archive_folder)
                            imap_client.safe_subscribe(archive_folder)
                        else:
                            raise create_error
            except Exception as e:
                print(f"Error checking/creating Archive folder: {e}")
                # Try with simple name as fallback
                archive_folder = 'Archive'
            
            # Copy email to Archive
            result, _ = imap_client.safe_uid('COPY', uid, archive_folder)
            if result == 'OK':
                # Delete from current folder
                imap_client.safe_uid('STORE', uid, '+FLAGS', '(\\Deleted)')
                imap_session.expunge()
                return jsonify({'success': True, 'message': f'Archived to {archive_folder}'})
            else:
                return jsonify({'success': False, 'message': f'Failed to copy to {archive_folder}'})
        
        elif action == 'inbox':
            # Move to INBOX
            result, _ = imap_client.safe_uid('COPY', uid, 'INBOX')
            if result == 'OK':
                imap_client.safe_uid('STORE', uid, '+FLAGS', '(\\Deleted)')
                imap_session.expunge()
                return jsonify({'success': True})
        
        return jsonify({'success': False, 'message': 'Invalid action'})
    
    except Exception as e:
        return jsonify({'success': False, 'message': str(e)})


def main():
    """Main entry point."""
    global email_address, password, imap_server
    
    email_address = os.getenv('EMAIL_ADDRESS')
    password = os.getenv('EMAIL_PASSWORD')
    imap_server = os.getenv('IMAP_SERVER')
    
    if not email_address or not password:
        print("Error: EMAIL_ADDRESS and EMAIL_PASSWORD must be set in .env")
        sys.exit(1)
    
    if not imap_server:
        imap_server = detect_imap_server(email_address)
    
    if not connect():
        print("Error: Could not connect to IMAP server")
        sys.exit(1)
    
    print(f"✅ Connected to {imap_server}")
    print(f"🌐 Starting web UI on http://0.0.0.0:5001")
    print(f"📧 Managing SPAM folder for {email_address}")
    
    app.run(host='0.0.0.0', port=5001, debug=False)


if __name__ == '__main__':
    main()
