#!/usr/bin/env python3
"""
Rollback script for user authentication migration.

This script safely removes user_id fields and indexes added by the
user authentication migration, restoring the database to its previous state.

Usage:
    python rollback_user_authentication.py [--confirm] [--backup-first]

Options:
    --confirm       Skip confirmation prompt (for automated scripts)
    --backup-first  Create backup before rollback
"""

import argparse
import asyncio
from datetime import datetime
import json
import logging
import os
import os
import sys
from typing import Any, Dict

from dotenv import load_dotenv

from analytics.create_indexes import drop_user_id_indexes
from analytics.database import get_database, test_connection

        timestamp = datetime.utcnow().strftime("%Y%m%d_%H%M%S")
        backup_file = os.path.join(self.backup_dir, f"pre_rollback_backup_{timestamp}.json")
        
        logger.info(f"Creating pre-rollback backup: {backup_file}")
        
        try:
            backup_data = {
                "timestamp": timestamp,
                "backup_type": "pre_rollback",
                "collections": {}
            }
            
            # Backup user_id data from each collection
            for collection_name in ["sessions", "messages", "search_analytics"]:
                try:
                    collection = self.db[collection_name]
                    
                    # Only backup documents with user_id field
                    documents = await collection.find({"user_id": {"$exists": True}}).to_list(length=None)
                    
                    # Convert to serializable format
                    serializable_docs = []
                    for doc in documents:
                        serializable_doc = {}
                        for key, value in doc.items():
                            if isinstance(value, datetime):
                                serializable_doc[key] = value.isoformat()
                            else:
                                serializable_doc[key] = str(value) if hasattr(value, '__str__') else value
                        serializable_docs.append(serializable_doc)
                    
                    backup_data["collections"][collection_name] = {
                        "count": len(documents),
                        "documents": serializable_docs
                    }
                    
                    logger.info(f"Backed up {len(documents)} documents with user_id from {collection_name}")
                    
                except Exception as e:
                    logger.warning(f"Could not backup {collection_name}: {e}")
                    backup_data["collections"][collection_name] = {"error": str(e)}
            
            # Write backup file
            with open(backup_file, 'w') as f:
                json.dump(backup_data, f, indent=2)
            
            logger.info(f"Pre-rollback backup completed: {backup_file}")
            return backup_file
            
        except Exception as e:
            logger.error(f"Backup failed: {e}")
            raise
    
    async def execute_rollback(self) -> bool:
        """Execute the complete rollback process"""
        try:
            logger.info("Starting user authentication rollback...")
            
            # Create backup if requested
            backup_file = None
            if self.backup_first:
                backup_file = await self.create_rollback_backup()
            
            # Record rollback start in migration history
            await self._record_rollback_event("rollback_started", {
                "backup_file": backup_file
            })
            
            # Step 1: Drop user_id indexes
            logger.info("Step 1: Dropping user_id indexes...")
            index_success = await drop_user_id_indexes()
            
            if not index_success:
                logger.warning("Some indexes could not be dropped (they may not exist)")
            
            # Step 2: Remove user_id fields from documents
            logger.info("Step 2: Removing user_id fields from documents...")
            field_success = await self._remove_user_id_fields()
            
            if not field_success:
                await self._record_rollback_event("rollback_failed", {"step": "remove_fields"})
                return False
            
            # Step 3: Verify rollback
            logger.info("Step 3: Verifying rollback...")
            verification_success = await self._verify_rollback()
            
            if not verification_success:
                await self._record_rollback_event("rollback_failed", {"step": "verification"})
                return False
            
            # Record successful completion
            await self._record_rollback_event("rollback_completed", {
                "backup_file": backup_file,
                "verification_passed": True
            })
            
            logger.info("Rollback completed successfully!")
            return True
            
        except Exception as e:
            logger.error(f"Rollback failed: {e}")
            await self._record_rollback_event("rollback_failed", {"error": str(e)})
            return False
    
    async def _remove_user_id_fields(self) -> bool:
        """Remove user_id fields from all documents"""
        try:
            collections_to_update = ["sessions", "messages", "search_analytics"]
            
            for collection_name in collections_to_update:
                logger.info(f"Removing user_id field from {collection_name}...")
                
                collection = self.db[collection_name]
                
                # Remove user_id field from all documents
                result = await collection.update_many(
                    {"user_id": {"$exists": True}},
                    {"$unset": {"user_id": ""}}
                )
                
                logger.info(f"Removed user_id field from {result.modified_count} documents in {collection_name}")
            
            return True
            
        except Exception as e:
            logger.error(f"Error removing user_id fields: {e}")
            return False
    
    async def _verify_rollback(self) -> bool:
        """Verify that rollback was successful"""
        try:
            logger.info("Verifying rollback completion...")
            
            collections_to_check = ["sessions", "messages", "search_analytics"]
            
            for collection_name in collections_to_check:
                collection = self.db[collection_name]
                
                # Check that no documents have user_id field
                docs_with_user_id = await collection.count_documents({"user_id": {"$exists": True}})
                
                if docs_with_user_id > 0:
                    logger.error(f"Rollback verification failed: {docs_with_user_id} documents in {collection_name} still have user_id field")
                    return False
                else:
                    logger.info(f"✓ No user_id fields found in {collection_name}")
            
            # Check that user_id indexes are gone
            expected_user_indexes = [
                "user_id_sparse",
                "user_id_start_time_compound",
                "user_id_timestamp_compound"
            ]
            
            for collection_name in collections_to_check:
                collection = self.db[collection_name]
                indexes = await collection.list_indexes().to_list(length=None)
                index_names = [idx["name"] for idx in indexes]
                
                remaining_user_indexes = [name for name in expected_user_indexes if name in index_names]
                
                if remaining_user_indexes:
                    logger.warning(f"Some user_id indexes still exist in {collection_name}: {remaining_user_indexes}")
                    # This is a warning, not a failure
                else:
                    logger.info(f"✓ No user_id indexes found in {collection_name}")
            
            logger.info("✅ Rollback verification passed!")
            return True
            
        except Exception as e:
            logger.error(f"Rollback verification failed: {e}")
            return False
    
    async def _record_rollback_event(self, status: str, details: Dict[str, Any]):
        """Record rollback event in migration history"""
        try:
            # Check if migration_history collection exists
            collections = await self.db.list_collection_names()
            if "migration_history" not in collections:
                logger.warning("Migration history collection does not exist, cannot record rollback event")
                return
            
            migration_coll = self.db["migration_history"]
            
            event = {
                "migration_name": "user_authentication",
                "status": status,
                "timestamp": datetime.utcnow(),
                "details": details,
                "operation_type": "rollback"
            }
            
            await migration_coll.insert_one(event)
            
        except Exception as e:
            logger.warning(f"Could not record rollback event: {e}")

def setup_logging():
    """Setup logging configuration"""
    logging.basicConfig(
        level=logging.INFO,
        format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
    )

async def main():
    """Main CLI function"""
    parser = argparse.ArgumentParser(description="Rollback user authentication migration")
    parser.add_argument("--confirm", action="store_true", help="Skip confirmation prompt")
    parser.add_argument("--backup-first", action="store_true", help="Create backup before rollback")
    
    args = parser.parse_args()
    
    setup_logging()
    
    print(f"🔄 User Authentication Migration Rollback")
    print(f"Backup first: {args.backup_first}")
    print("-" * 50)
    
    try:
        # Test database connection
        if not await test_connection():
            print("❌ Database connection failed. Check your MONGODB_URL.")
            sys.exit(1)
        
        # Initialize rollback manager
        manager = RollbackManager(backup_first=args.backup_first)
        await manager.initialize()
        
        # Check rollback feasibility
        print("Checking rollback feasibility...")
        feasibility = await manager.check_rollback_feasibility()
        
        if not feasibility.get("can_rollback", False):
            print(f"❌ Rollback not feasible: {feasibility.get('error', 'Unknown error')}")
            sys.exit(1)
        
        # Show warnings
        if feasibility.get("warnings"):
            print("\n⚠️  Rollback Warnings:")
            for warning in feasibility["warnings"]:
                print(f"  - {warning}")
        
        # Show data loss risk
        if feasibility.get("data_loss_risk"):
            print("\n🚨 DATA LOSS WARNING:")
            print("This rollback will permanently delete user_id associations.")
            print("User-specific analytics data will be lost.")
        
        # Show collection status
        print(f"\n📊 Collections Status:")
        for collection, status in feasibility.get("collections_status", {}).items():
            if isinstance(status, dict):
                print(f"  {collection}: {status['docs_with_user_id']}/{status['total_docs']} docs have user_id")
            else:
                print(f"  {collection}: {status}")
        
        # Confirmation prompt
        if not args.confirm:
            print(f"\n❓ Are you sure you want to proceed with rollback?")
            if feasibility.get("data_loss_risk"):
                print("   This will permanently delete user authentication data!")
            
            response = input("Type 'yes' to confirm: ").strip().lower()
            if response != 'yes':
                print("Rollback cancelled.")
                sys.exit(0)
        
        # Execute rollback
        print("\nExecuting rollback...")
        success = await manager.execute_rollback()
        
        if success:
            print("✅ Rollback completed successfully!")
            print("\nThe database has been restored to its pre-migration state.")
            print("All user_id fields and indexes have been removed.")
        else:
            print("❌ Rollback failed! Check logs for details.")
            sys.exit(1)
            
    except Exception as e:
        print(f"❌ Rollback failed: {e}")
        sys.exit(1)

if __name__ == "__main__":
    asyncio.run(main())