#!/usr/bin/env python3
"""
Automated Data Consistency Fix
This script automatically fixes the data inconsistency issue.
"""

import os
import sys
from datetime import datetime
from pymongo import MongoClient
from bson import ObjectId

def get_mongodb_connection():
    """Get MongoDB connection using environment variables"""
    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 in environment variables")
            return None, None
            
        client = MongoClient(mongo_uri)
        db = client[mongo_db]
        
        # Test connection
        client.server_info()
        print(f"✅ Connected to MongoDB: {mongo_db}")
        return client, db
        
    except Exception as e:
        print(f"❌ Failed to connect to MongoDB: {e}")
        return None, None

def fix_data_consistency():
    """Fix the data consistency issue by migrating recordings to website user IDs"""
    print("🚀 Starting Automated Data Consistency Fix")
    print("=" * 60)
    
    # Get MongoDB connection
    client, db = get_mongodb_connection()
    if client is None or db is None:
        return
    
    try:
        # Get the desktop app user ID (the one with recordings)
        desktop_user_id = "d33268cde68beec8e654f40ec78cf9bba02a4383cb0159da7b49a88eaa4d1c55"
        
        # Get the most recent website user ID (the one with most test cases)
        website_user_id = "34d0f132-76ff-4adf-b98a-e8649016e747"
        
        print(f"🔄 Migrating data from desktop user ID to website user ID...")
        print(f"   Desktop User ID: {desktop_user_id}")
        print(f"   Website User ID: {website_user_id}")
        
        # Migrate codegen recordings from desktop user ID to website user ID
        recordings_result = db.codegen_recordings.update_many(
            {
                "$or": [
                    {"userId": desktop_user_id},
                    {"currentJWTUserId": desktop_user_id}
                ]
            },
            {
                "$set": {
                    "userId": website_user_id,
                    "currentJWTUserId": website_user_id,
                    "user_id": website_user_id
                }
            }
        )
        print(f"✅ Migrated {recordings_result.modified_count} codegen recordings")
        
        # Also add user_id field to existing test cases and scripts for consistency
        test_cases_result = db.test_cases.update_many(
            {"user_id": website_user_id, "userId": {"$exists": False}},
            {"$set": {"userId": website_user_id, "currentJWTUserId": website_user_id}}
        )
        print(f"✅ Updated {test_cases_result.modified_count} test cases with additional user ID fields")
        
        scripts_result = db.generated_scripts.update_many(
            {"user_id": website_user_id, "userId": {"$exists": False}},
            {"$set": {"userId": website_user_id, "currentJWTUserId": website_user_id}}
        )
        print(f"✅ Updated {scripts_result.modified_count} generated scripts with additional user ID fields")
        
        print("\n🎉 Data consistency fix completed!")
        print("\nNext steps:")
        print("1. Restart your website application")
        print("2. Restart your desktop application") 
        print("3. Log in again to both applications")
        print("4. Check if data is now consistent between platforms")
        
        # Verify the fix
        print("\n🔍 Verifying the fix...")
        recordings_count = db.codegen_recordings.count_documents({"userId": website_user_id})
        test_cases_count = db.test_cases.count_documents({"user_id": website_user_id})
        scripts_count = db.generated_scripts.count_documents({"user_id": website_user_id})
        
        print(f"📊 Data now associated with website user ID {website_user_id}:")
        print(f"   - Recordings: {recordings_count}")
        print(f"   - Test Cases: {test_cases_count}")
        print(f"   - Scripts: {scripts_count}")
        
    except Exception as e:
        print(f"❌ Error during data consistency fix: {e}")
    finally:
        if client:
            client.close()

if __name__ == "__main__":
    fix_data_consistency()