#!/usr/bin/env python3
"""
Direct Fix Data Consistency
This script directly fixes the data consistency issue by ensuring all data uses the same user ID format.
"""

import os
from pymongo import MongoClient

def direct_fix_data_consistency():
    """Directly fix the data consistency issue"""
    print("🔧 Direct Fix Data Consistency")
    print("=" * 50)
    
    try:
        mongo_uri = os.getenv('MONGODB_URI')
        mongo_db = os.getenv('MONGODB_DB')
        
        if not mongo_uri or not mongo_db:
            print("❌ MongoDB configuration not found")
            return
            
        client = MongoClient(mongo_uri)
        db = client[mongo_db]
        
        # Get the main user ID
        user_id = '34d0f132-76ff-4adf-b98a-e8649016e747'
        
        print(f"🎯 Fixing data for user: {user_id}")
        
        # Step 1: Ensure all recordings have all user ID fields
        print("\n📊 Fixing recordings...")
        recordings_result = db.codegen_recordings.update_many(
            {'userId': user_id},
            {
                '$set': {
                    'user_id': user_id,
                    'currentJWTUserId': user_id
                }
            }
        )
        print(f"✅ Updated {recordings_result.modified_count} recordings")
        
        # Step 2: Ensure all test cases have all user ID fields
        print("\n📋 Fixing test cases...")
        test_cases_result = db.test_cases.update_many(
            {'user_id': user_id},
            {
                '$set': {
                    'userId': user_id,
                    'currentJWTUserId': user_id
                }
            }
        )
        print(f"✅ Updated {test_cases_result.modified_count} test cases")
        
        # Step 3: Ensure all generated scripts have all user ID fields
        print("\n🔧 Fixing generated scripts...")
        scripts_result = db.generated_scripts.update_many(
            {'user_id': user_id},
            {
                '$set': {
                    'userId': user_id,
                    'currentJWTUserId': user_id
                }
            }
        )
        print(f"✅ Updated {scripts_result.modified_count} generated scripts")
        
        # Step 4: Verify the fix
        print("\n🔍 Verifying the fix...")
        
        # Check recordings
        recordings_count = db.codegen_recordings.count_documents({
            '$and': [
                {'userId': user_id},
                {'user_id': user_id},
                {'currentJWTUserId': user_id}
            ]
        })
        print(f"📊 Recordings with all user ID fields: {recordings_count}")
        
        # Check test cases
        test_cases_count = db.test_cases.count_documents({
            '$and': [
                {'user_id': user_id},
                {'userId': user_id},
                {'currentJWTUserId': user_id}
            ]
        })
        print(f"📊 Test cases with all user ID fields: {test_cases_count}")
        
        # Check generated scripts
        scripts_count = db.generated_scripts.count_documents({
            '$and': [
                {'user_id': user_id},
                {'userId': user_id},
                {'currentJWTUserId': user_id}
            ]
        })
        print(f"📊 Generated scripts with all user ID fields: {scripts_count}")
        
        # Step 5: Create a test query to verify API compatibility
        print("\n🧪 Testing API compatibility...")
        
        # Test the same query that the website API uses
        website_query = {
            '$or': [
                {'userId': user_id},
                {'currentJWTUserId': user_id},
                {'user_id': user_id},
                {'owner_id': user_id}
            ]
        }
        
        website_recordings = db.codegen_recordings.count_documents(website_query)
        print(f"📊 Recordings found by website API query: {website_recordings}")
        
        # Test the same query that the desktop app uses
        desktop_query = {
            '$or': [
                {'userId': user_id},
                {'currentJWTUserId': user_id}
            ]
        }
        
        desktop_recordings = db.codegen_recordings.count_documents(desktop_query)
        print(f"📊 Recordings found by desktop app query: {desktop_recordings}")
        
        if website_recordings == desktop_recordings:
            print("✅ Both applications should now see the same data!")
        else:
            print("⚠️ There might still be a query difference")
        
        print("\n🎉 Data consistency fix completed!")
        print("\nNext steps:")
        print("1. Restart both applications completely")
        print("2. Clear browser cache (Ctrl+Shift+R)")
        print("3. Log out and log back in to both applications")
        print("4. Check if data is now consistent")
        
        client.close()
        
    except Exception as e:
        print(f"❌ Error: {e}")

if __name__ == "__main__":
    direct_fix_data_consistency()