"""
Fix ALL mail-related activities (including First Economy Mail, etc.)
"""
import psycopg2
import os
from dotenv import load_dotenv

load_dotenv()

# Try different connection methods
def get_connection():
    attempts = [
        lambda: psycopg2.connect(os.getenv("DATABASE_URL", "")),
        lambda: psycopg2.connect(host="localhost", database="timesheet", user="postgres", password="asdf1234"),
        lambda: psycopg2.connect(host="localhost", database="timesheet", user="timesheet_user", password="asdf1234"),
    ]
    
    for attempt in attempts:
        try:
            return attempt()
        except:
            continue
    return None

print("🔧 Fixing ALL Mail Activities in Productivity")
print("=" * 50)

conn = get_connection()
if not conn:
    print("❌ Could not connect to database!")
    exit(1)

try:
    cur = conn.cursor()
    
    # First, let's see what mail-related activities are in productive
    print("\n📧 Checking mail-related activities in productive category...")
    cur.execute("""
        SELECT DISTINCT window_title, COUNT(*) as count
        FROM activity_records
        WHERE category = 'productive'
        AND (
            LOWER(window_title) LIKE '%mail%'
            OR LOWER(window_title) LIKE '%gmail%'
            OR LOWER(window_title) LIKE '%outlook%'
            OR LOWER(window_title) LIKE '%inbox%'
            OR LOWER(window_title) LIKE '%email%'
        )
        GROUP BY window_title
        ORDER BY count DESC
        LIMIT 20
    """)
    
    mail_activities = cur.fetchall()
    if mail_activities:
        print("\nFound these mail activities in productive:")
        for title, count in mail_activities:
            print(f"  - {title}: {count} activities")
    
    # Now fix ALL mail-related activities
    print("\n🚀 Moving ALL mail activities to browser category...")
    
    # Comprehensive mail fix
    cur.execute("""
        UPDATE activity_records
        SET category = 'browser', subcategory = 'email'
        WHERE category = 'productive'
        AND (
            LOWER(window_title) LIKE '%mail%'
            OR LOWER(window_title) LIKE '%gmail%'
            OR LOWER(window_title) LIKE '%outlook%'
            OR LOWER(window_title) LIKE '%inbox%'
            OR LOWER(window_title) LIKE '%email%'
            OR LOWER(url) LIKE '%mail.google.com%'
            OR LOWER(url) LIKE '%outlook.com%'
            OR LOWER(url) LIKE '%mail.yahoo.com%'
        )
    """)
    mail_fixed = cur.rowcount
    print(f"  ✅ Moved {mail_fixed} mail activities to browser")
    
    # Also fix YouTube and other non-productive items
    fixes = [
        ("YouTube", """
            UPDATE activity_records
            SET category = 'browser', subcategory = 'entertainment'
            WHERE category = 'productive'
            AND (LOWER(window_title) LIKE '%youtube%' OR LOWER(url) LIKE '%youtube.com%' OR LOWER(url) LIKE '%youtu.be%')
        """),
        ("Social Media", """
            UPDATE activity_records
            SET category = 'browser', subcategory = 'social-media'
            WHERE category = 'productive'
            AND (
                LOWER(window_title) LIKE '%facebook%'
                OR LOWER(window_title) LIKE '%instagram%'
                OR LOWER(window_title) LIKE '%twitter%'
                OR LOWER(window_title) LIKE '%reddit%'
                OR LOWER(window_title) LIKE '%linkedin%'
                OR LOWER(window_title) LIKE '%whatsapp%'
            )
        """),
        ("Entertainment", """
            UPDATE activity_records
            SET category = 'browser', subcategory = 'entertainment'
            WHERE category = 'productive'
            AND (
                LOWER(window_title) LIKE '%netflix%'
                OR LOWER(window_title) LIKE '%spotify%'
                OR LOWER(window_title) LIKE '%twitch%'
                OR LOWER(window_title) LIKE '%amazon prime%'
                OR LOWER(window_title) LIKE '%disney%'
                OR LOWER(window_title) LIKE '%hulu%'
            )
        """),
        ("Shopping", """
            UPDATE activity_records
            SET category = 'browser', subcategory = 'shopping'
            WHERE category = 'productive'
            AND (
                LOWER(window_title) LIKE '%amazon.com%'
                OR LOWER(window_title) LIKE '%ebay%'
                OR LOWER(window_title) LIKE '%flipkart%'
                OR LOWER(window_title) LIKE '%shopping%'
                OR LOWER(window_title) LIKE '%myntra%'
            )
        """),
    ]
    
    for name, query in fixes:
        cur.execute(query)
        count = cur.rowcount
        if count > 0:
            print(f"  ✅ {name}: {count} activities moved to browser")
    
    conn.commit()
    
    # Verify First Economy Mail specifically
    print("\n🔍 Checking 'First Economy Mail' specifically...")
    cur.execute("""
        SELECT category, COUNT(*) 
        FROM activity_records 
        WHERE window_title LIKE '%First Economy Mail%'
        GROUP BY category
    """)
    for row in cur.fetchall():
        print(f"  First Economy Mail is now in: {row[0]} ({row[1]} activities)")
    
    # Show final counts
    print("\n📊 Final category distribution:")
    cur.execute("""
        SELECT category, COUNT(*) as count
        FROM activity_records
        WHERE timestamp > CURRENT_DATE - INTERVAL '7 days'
        GROUP BY category
        ORDER BY count DESC
    """)
    for row in cur.fetchall():
        emoji = {'productive': '💻', 'browser': '🌐', 'server': '☁️', 'non-work': '🎮'}.get(row[0], '❓')
        print(f"  {emoji} {row[0]}: {row[1]} activities")
    
    print("\n✅ Done! ALL mail activities (including 'First Economy Mail') have been moved to the browser category.")
    print("\n⚠️  IMPORTANT: You must RESTART your dashboard/browser to see the changes!")
    
except Exception as e:
    print(f"\n❌ Error: {e}")
finally:
    if conn:
        conn.close()