"""Core migration logic for embedding models."""

import itertools
import random
import string
import threading
from collections import defaultdict
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import List, Optional, Tuple
from datetime import datetime

from h2ogpte import H2OGPTE

from h2ogpte_migration.config import MigrationConfig, MODEL_MAPPINGS
from h2ogpte_migration.models import CollectionMigration, UserMigration, MigrationReport
from h2ogpte_migration.logger import MigrationLogger
from h2ogpte_migration.database import MigrationDatabase


def random_suffix(length: int = 5) -> str:
    return ''.join(random.choices(string.ascii_uppercase + string.digits, k=length))


class CollectionMigrator:
    def __init__(self, config: MigrationConfig, logger: MigrationLogger):
        self.config = config
        self.log = logger
        self.admin_client: Optional[H2OGPTE] = None
        self.user_client: Optional[H2OGPTE] = None
        self.current_username: Optional[str] = None
        self.current_user_id: Optional[str] = None
        self.users_with_temp_keys: set = set()  # track users for API key cleanup
        self.created_key_names: set = set()  # track exact key names we create
        self._key_lock = threading.Lock()  # protects users_with_temp_keys and created_key_names
        self._executor: Optional[ThreadPoolExecutor] = None  # track active executor for graceful shutdown
        self._cached_users: Optional[list] = None  # cached user list to avoid redundant fetches
        self.db: Optional[MigrationDatabase] = None
        # Suppress SSL warnings once if no cert provided
        # TODO: revisit
        if not config.cert_path:
            import urllib3
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
        safe_config = {k: v for k, v in (config.__dict__ if hasattr(config, '__dict__') else {}).items()
                       if k not in ('admin_api_key', 'user_api_key')}
        self.report = MigrationReport(
            start_time=datetime.now().isoformat(),
            config=safe_config
        )
    
    @property
    def is_self_service(self) -> bool:
        return bool(self.config.user_api_key)
    
    def _is_guest_user(self, username: str) -> bool:
        return username.startswith("Guest_")
    
    def _get_ssl_verify_config(self):
        """Get SSL verification configuration.

        Returns:
            SSL verification setting (cert path or False for no verification)
        """
        return self.config.cert_path if self.config.cert_path else False

    def _determine_job_status(self, job) -> Tuple[str, Optional[str]]:
        """Determine job status and error message from a job object.
        
        Args:
            job: H2OGPTE Job object
            
        Returns:
            Tuple of (status_type, error_message) where status_type is one of:
            'canceled', 'failed', 'completed', 'running'
        """
        if job.canceled:
            return ('canceled', 'Job was canceled')
        elif job.failed > 0 or len(job.errors) > 0:
            if job.errors:
                error_msg = f"Job had errors: {'; '.join(job.errors)}"
            else:
                error_msg = "Job failed."
            if job.completed:
                return ('failed', error_msg)
            else:
                # still running, let it complete
                return ('running', error_msg)
        elif job.completed:
            # success (no failures, no errors)
            return ('completed', None)
        else:
            # still running
            return ('running', None)
    
    def run(self) -> MigrationReport:
        try:
            self.log.header("h2oGPTe Embedding Model Migration")
            self.log.info(str(self.config))
            
            if self.config.migrate_chats and not self.config.wait_for_completion and not self.config.verify:
                self.log.warning("--migrate-chats without --wait-for-completion will be deferred to --verify --migrate-chats")

            if self.config.max_concurrent_jobs > 1 and not self.config.wait_for_completion:
                self.log.warning("--max-concurrent-jobs has no effect without --wait-for-completion")

            if self.config.force_remigrate:
                self.log.warning("--force-remigrate: if connectors/chats were previously moved to another collection, "
                                 "use --move-connectors/--move-chats to recover them after re-migration completes")

            if self.is_self_service:
                self._connect_user()
            else:
                self._connect_admin()

            # Move mode — manual connector/chat operations (no DB needed)
            if self.config.move_connectors or self.config.move_chats:
                self._run_move_operations()
                self.log.info("")
                self.log.info(f"Detailed log: {self.log.get_log_file()}")
                self.report.complete()
                return self.report

            # Initialize database
            if not self.config.dry_run:
                self.log.info(f"Initializing migration tracking database: {self.config.db_path}")
                self.db = MigrationDatabase(self.config.db_path)
                self.log.success("Database initialized")

            # Check if this is verify mode (no migration model flags set)
            # Note: --users, --all-users, --collections are scope filters, not migration flags
            has_migration_flags = (
                self.config.use_model_mappings or
                self.config.target_model or
                self.config.source_model
            )

            is_verify_mode = self.config.verify and not has_migration_flags

            if is_verify_mode:
                if self.is_self_service:
                    self.log.info("Running in verify mode (your migrations)")
                elif self.config.users:
                    self.log.info(f"Running in verify mode for user(s): {', '.join(self.config.users)}")
                elif self.config.collection_ids:
                    self.log.info(f"Running in verify mode for {len(self.config.collection_ids)} collection(s)")
                else:
                    self.log.info("Running in verify mode (all users)")

                if self.db:
                    verify_counters = self._verify_import_jobs()
                    self.report.verify_counters = verify_counters

                self.log.info("")
                self.log.info(f"Detailed log: {self.log.get_log_file()}")
                self.report.complete()
                return self.report
            elif self.is_self_service:
                if self.config.collection_ids:
                    self._migrate_own_collections_by_id(self.config.collection_ids)
                else:
                    # migrate all of user's collections
                    self._migrate_single_user(
                        username=self.current_username,
                        user_client=self.user_client,
                        user_id=self.current_user_id
                    )
            else:
                # Admin migration mode
                if self.config.collection_ids:
                    self._migrate_collections_by_id(self.config.collection_ids)
                elif self.config.users:
                    self._migrate_users(self.config.users)
                elif self.config.all_users:
                    self._migrate_all_users()
            
            self.report.complete()
            self._print_summary()
            
            return self.report
            
        except Exception as e:
            self.log.failure(f"Migration failed: {e}")
            self.report.complete()
            raise
        finally:
            # Wait for any running workers to finish before cleanup
            if self._executor:
                self.log.info("Waiting for running workers to finish...")
                self._executor.shutdown(wait=True)
                self._executor = None

            # Clean up API keys FIRST (requires admin client)
            if self.admin_client:
                self._cleanup_api_keys()

            # Then close database connection
            if self.db:
                self.db.close()
                self.log.info("Database connection closed")
    
    def _connect_admin(self):
        self.log.info("Connecting to h2oGPTe as admin...")
        self.log.info(f"URL: {self.config.url}")
        
        self.admin_client = H2OGPTE(
            address=self.config.url,
            api_key=self.config.admin_api_key,
            verify=self._get_ssl_verify_config()
        )
        
        # Verify connection
        meta = self.admin_client.get_meta()
        self.log.success(f"Connected as: {meta.username}")
    
    def _connect_user(self):
        # self service mode
        self.log.info("Connecting to h2oGPTe...")
        self.log.info(f"URL: {self.config.url}")
        
        self.user_client = H2OGPTE(
            address=self.config.url,
            api_key=self.config.user_api_key,
            verify=self._get_ssl_verify_config()
        )
        
        # Get current user info
        meta = self.user_client.get_meta()
        self.current_username = meta.username
        self.current_user_id = meta.user_id
        
        self.log.success(f"Connected as: {self.current_username}")
    
    def _migrate_all_users(self):
        self.log.section("Fetching all users")

        users = self.admin_client.list_users(offset=0, limit=-1)
        self.log.info(f"Found {len(users)} total users")
        
        # Filter users with collections (and cache their clients)
        users_with_collections = []
        user_clients = {}  # Cache to avoid duplicate API key creation
        guest_user_count = 0
        
        for user in users:
            # Skip guest users in bulk migration mode
            if self._is_guest_user(user.username):
                guest_user_count += 1
                self.log.info(f"  Skipping guest user: {user.username}")
                continue
            
            try:
                user_client = self._get_user_client(user.id, user.username)
                collection_infos = user_client.list_recent_collections(offset=0, limit=-1)
                owned_collections = [c for c in collection_infos if c.username == user.username]

                if owned_collections:
                    users_with_collections.append((user, len(owned_collections)))
                    user_clients[user.username] = user_client  # Cache the client
                    self.log.info(f"  {user.username}: {len(owned_collections)} collections")

            except Exception as e:
                self.log.warning(f"Could not fetch collections for {user.username}: {e}")
        
        # Log guest user skip summary
        if guest_user_count > 0:
            self.log.info(f"Skipped checking {guest_user_count} guest user(s)")
        
        self.log.info("")
        self.log.info(f"Collected users with collections ({len(users_with_collections)}):")
        for user, count in users_with_collections:
            self.log.info(f"  - {user.username} ({count} collections)")
        
        # Migrate each user (reuse cached clients)
        for user, _ in users_with_collections:
            try:
                cached_client = user_clients[user.username]
                self._migrate_single_user(user.username, user_client=cached_client, user_id=user.id)
            except Exception as e:
                self.log.error(f"Failed to migrate user {user.username}: {e}")
    
    def _migrate_users(self, usernames: List[str]):
        """Migrate collections for a list of users."""
        if len(usernames) == 1:
            self.log.header(f"Migrating User: {usernames[0]}")
        else:
            self.log.header(f"Migrating {len(usernames)} Users")
            self.log.info(f"Users: {', '.join(usernames)}")
        
        for username in usernames:
            try:
                self._migrate_single_user(username)
            except Exception as e:
                self.log.error(f"Failed to migrate user {username}: {e}")
    
    def _migrate_single_user(self, username: str, user_client: Optional[H2OGPTE] = None, user_id: Optional[str] = None):
        """Migrate all collections for a single user."""
        if len(self.config.users or []) > 1:
            # Multi-user mode - use section instead of header
            self.log.section(f"User: {username}")
        else:
            # Single user mode - already has header from _migrate_users
            pass
        
        # Find user (only if user_id not provided)
        if not user_id:
            target_user = self._find_user_by_username(username)
            
            if not target_user:
                self.log.failure(f"User '{username}' not found")
                return
            
            user_id = target_user.id
        
        self.log.info(f"User ID: {user_id}")
        
        # Create API key for user (only if not provided)
        if not user_client:
            user_client = self._get_user_client(user_id, username)
        
        # Get user's collections (lightweight list)
        self.log.info("Fetching user's collections...")
        collection_infos = user_client.list_recent_collections(offset=0, limit=-1)
        
        # Filter to owned collections
        owned_infos = [c for c in collection_infos if c.username == username]
        self.log.info(f"User owns {len(owned_infos)} collections")

        # Group collections based on migration mode
        if self.config.use_model_mappings:
            # Model mappings mode: group by embedding_model
            collections_by_model = defaultdict(list)
            for collection in owned_infos:
                if collection.embedding_model in MODEL_MAPPINGS:
                    collections_by_model[collection.embedding_model].append(collection)
            
            total_to_migrate = sum(len(colls) for colls in collections_by_model.values())
            skipped_count = len(owned_infos) - total_to_migrate
            self.log.info(f"Found {total_to_migrate} collections with deprecated models:")
            for model, colls in collections_by_model.items():
                target = MODEL_MAPPINGS[model]
                self.log.info(f"  {model} ({len(colls)} collections) → {target}")
            if skipped_count > 0:
                self.log.info(f"Skipped {skipped_count} collection(s) with unmapped models")
        else:
            # Manual mode: filter by source model
            collections_to_migrate = []
            if self.config.source_model:
                collections_to_migrate = [c for c in owned_infos
                                         if c.embedding_model == self.config.source_model]
                self.log.info(f"Collections with model '{self.config.source_model}': {len(collections_to_migrate)}")
            else:
                collections_to_migrate = owned_infos
                self.log.info(f"Migrating all owned collections")
            
            # For manual mode, create simple grouping
            collections_by_model = {self.config.source_model or "all": collections_to_migrate}
        
        # Check if anything to migrate
        if not any(collections_by_model.values()):
            self.log.warning("No collections to migrate for this user")
            return
        
        # Create user migration tracker
        user_migration = UserMigration(
            username=username,
            user_id=user_id
        )
        
        # Migrate collections
        total_count = sum(len(colls) for colls in collections_by_model.values())
        self.log.info(f"Collections to migrate: {total_count}")
        self.log.info("")
        
        if self.config.wait_for_completion and self.config.max_concurrent_jobs > 1:
            # Build flat work list for parallel processing
            work_items = []
            for source_model, collections in collections_by_model.items():
                target_model = self._get_target_model_for_source(source_model)
                for collection in collections:
                    work_items.append((user_client, collection, target_model, username, user_id))
            user_migration.collections = self._parallel_migrate_collections(work_items, total_count)
        else:
            current_idx = 0
            for source_model, collections in collections_by_model.items():
                # Determine target model using helper
                target_model = self._get_target_model_for_source(source_model)

                for collection in collections:
                    current_idx += 1
                    migration_result = self._migrate_collection(
                        user_client=user_client,
                        collection=collection,
                        target_model=target_model,
                        current=current_idx,
                        total=total_count,
                        username=username,
                        user_id=user_id
                    )
                    user_migration.collections.append(migration_result)

        self.report.users.append(user_migration)
        
        # Print user summary
        self.log.info("")
        self.log.info(f"User Summary:")
        self.log.info(f"  Total: {user_migration.total_collections}")
        self.log.success(f"Successful: {user_migration.successful_migrations}")
        if user_migration.failed_migrations > 0:
            self.log.failure(f"Failed: {user_migration.failed_migrations}")
        if user_migration.skipped_migrations > 0:
            self.log.warning(f"Skipped: {user_migration.skipped_migrations}")
    
    def _migrate_own_collections_by_id(self, collection_ids: List[str]):
        """Migrate specific collections with ownership verification (self-service)."""
        self.log.section("Migrating Your Collections")
        self.log.info(f"Checking {len(collection_ids)} collection(s)...")
        
        collections_to_migrate = []
        not_owned = []
        not_found = []
        
        for col_id in collection_ids:
            try:
                collection = self.user_client.get_collection(col_id)
                
                # Verify ownership
                if collection.username != self.current_username:
                    not_owned.append({
                        'id': col_id,
                        'name': collection.name,
                        'owner': collection.username
                    })
                    continue
                
                # Check model mapping if applicable
                if self.config.use_model_mappings and collection.embedding_model not in MODEL_MAPPINGS:
                    self.log.info(f"  Skipping {collection.name} — model '{collection.embedding_model}' has no model mapping")
                    continue

                collections_to_migrate.append(collection)

            except Exception as e:
                not_found.append({'id': col_id, 'error': str(e)})
                self.log.warning(f"Could not access collection {col_id}: {e}")

        # Report non-owned or not-found collections
        if not_owned:
            self.log.warning(f"Cannot migrate {len(not_owned)} collection(s) you don't own:")
            for item in not_owned:
                self.log.info(f"  - {item['name']} (ID: {item['id']}, owned by: {item['owner']})")
        
        if not_found:
            self.log.warning(f"Could not find {len(not_found)} collection(s)")
        
        if not collections_to_migrate:
            self.log.failure("No collections to migrate")
            return
        
        # Create user migration tracker
        user_migration = UserMigration(
            username=self.current_username,
            user_id=self.current_user_id
        )
        
        # Migrate each collection
        total = len(collections_to_migrate)
        if self.config.wait_for_completion and self.config.max_concurrent_jobs > 1:
            work_items = []
            for collection in collections_to_migrate:
                target_model = self._get_target_model_for_source(collection.embedding_model)
                work_items.append((self.user_client, collection, target_model, self.current_username, self.current_user_id))
            user_migration.collections = self._parallel_migrate_collections(work_items, total)
        else:
            for idx, collection in enumerate(collections_to_migrate, 1):
                target_model = self._get_target_model_for_source(collection.embedding_model)

                migration_result = self._migrate_collection(
                    user_client=self.user_client,
                    collection=collection,
                    target_model=target_model,
                    current=idx,
                    total=total,
                    username=self.current_username,
                    user_id=self.current_user_id
                )
                user_migration.collections.append(migration_result)

        self.report.users.append(user_migration)
    
    def _migrate_collections_by_id(self, collection_ids: List[str]):
        """Migrate specific collections by ID, auto-detecting owners."""
        
        self.log.section(f"Migrating Collections by ID")
        self.log.info(f"Processing {len(collection_ids)} collection(s)...")
        
        collections_by_owner = defaultdict(list)
        skipped_collections = []
        failed_to_fetch = []
        
        # Step 1: Look up each collection using admin filter (h2ogpte 1.6.56+)
        # CollectionInfo now includes embedding_model, so we get owner + model in one call per ID
        self.log.info("Fetching collection info (admin function)...")
        collection_info_map = {}  # col_id -> CollectionInfo

        for col_id in collection_ids:
            try:
                results = self.admin_client.list_all_collections_sort(
                    offset=0, limit=1, sort_column="updated_at", ascending=False, filter=col_id
                )
                if results and results[0].id == col_id:
                    col_info = results[0]
                    collection_info_map[col_info.id] = col_info
                elif results:
                    self.log.error(f"Collection lookup returned wrong ID: expected {col_id}, got {results[0].id}")
                    failed_to_fetch.append({'id': col_id, 'error': 'Collection lookup returned wrong result'})
                else:
                    self.log.error(f"Collection not found: {col_id}")
                    failed_to_fetch.append({'id': col_id, 'error': 'Collection not found'})
            except Exception as e:
                self.log.error(f"Failed to fetch collection {col_id}: {e}")
                failed_to_fetch.append({'id': col_id, 'error': str(e)})

        # Step 2: Filter by embedding model and group by owner
        collections_by_owner_id = defaultdict(list)

        for col_id, col_info in collection_info_map.items():
            embedding_model = col_info.embedding_model
            if self.config.use_model_mappings:
                if embedding_model not in MODEL_MAPPINGS:
                    skipped_collections.append({
                        'id': col_info.id,
                        'name': col_info.name,
                        'model': embedding_model,
                        'reason': 'Model has no migration mapping defined'
                    })
                    continue
            collections_by_owner_id[col_info.username].append(col_info)

        self.log.info(f"Found {sum(len(v) for v in collections_by_owner_id.values())} collection(s) grouped by {len(collections_by_owner_id)} owner(s)")

        # Step 3: Create API keys for each owner
        owner_clients = {}
        for username in collections_by_owner_id.keys():
            try:
                user = self._find_user_by_username(username)
                if user:
                    user_client = self._get_user_client(user.id, username)
                    owner_clients[username] = (user_client, user.id)
                else:
                    self.log.error(f"Could not find user: {username}")
            except Exception as e:
                self.log.error(f"Failed to create client for {username}: {e}")

        # Check if we successfully created any clients
        if not owner_clients and collections_by_owner_id:
            self.log.info("")
            self.log.failure("Failed to create API clients for any collection owners")

            for username, col_infos in collections_by_owner_id.items():
                for col_info in col_infos:
                    failed_to_fetch.append({
                        'id': col_info.id,
                        'error': f'Could not create API client for owner {username}'
                    })

            self.log.info("")
            self.log.info(f"Collection Summary:")
            self.log.info(f"  Total requested: {len(collection_ids)}")
            self.log.failure(f"  Failed to fetch: {len(failed_to_fetch)}")
            return

        # Step 4: Build collections_by_owner using CollectionInfo objects
        for username, col_infos in collections_by_owner_id.items():
            if username not in owner_clients:
                self.log.warning(f"Skipping {len(col_infos)} collection(s) for {username} (no client)")
                for col_info in col_infos:
                    failed_to_fetch.append({'id': col_info.id, 'error': f'Could not create client for owner {username}'})
                continue
            collections_by_owner[username].extend(col_infos)
        
        # Report what we found
        self.log.info("")
        self.log.info(f"Collection Summary:")
        self.log.info(f"  Total requested: {len(collection_ids)}")
        self.log.info(f"  Can migrate: {sum(len(colls) for colls in collections_by_owner.values())}")
        if skipped_collections:
            self.log.warning(f"  Skipped (unmapped models): {len(skipped_collections)}")
        if failed_to_fetch:
            self.log.failure(f"  Failed to fetch: {len(failed_to_fetch)}")
        
        # Show skipped collections details
        if skipped_collections:
            self.log.info("")
            self.log.warning("Skipped Collections:")
            for skip in skipped_collections:
                self.log.info(f"  - {skip['name']} (ID: {skip['id']})")
                self.log.info(f"    Model: {skip['model']}")
                self.log.info(f"    Reason: {skip['reason']}")
        
        # Check if we have anything to migrate
        if not collections_by_owner:
            self.log.info("")
            if self.config.use_model_mappings:
                self.log.failure("No collections to migrate!")
                self.log.info("None of the specified collections use models with defined migration mappings.")
                self.log.info("")
                self.log.info("Available model mappings:")
                for source, target in MODEL_MAPPINGS.items():
                    self.log.info(f"  {source} → {target}")
                self.log.info("")
                self.log.info("Tip: Use --target-model to manually specify the target embedding model.")
            else:
                self.log.failure("No collections available for migration!")
            return
        
        # Group by owner and log
        self.log.info("")
        self.log.info(f"Collections grouped by owner ({len(collections_by_owner)} owner(s)):")
        for username, colls in collections_by_owner.items():
            self.log.info(f"  - {username}: {len(colls)} collection(s)")
        
        # Migrate collections grouped by owner (reuse owner_clients from Step 3)
        self.log.info("")
        total_collections = sum(len(colls) for colls in collections_by_owner.values())

        if self.config.wait_for_completion and self.config.max_concurrent_jobs > 1:
            # Build flat work list across all owners
            work_items = []
            work_owners = []  # track (username, user_id) per item for report grouping
            for username, collections in collections_by_owner.items():
                if username not in owner_clients:
                    self.log.warning(f"Skipping {len(collections)} collection(s) for {username} (no client)")
                    continue
                user_client, user_id = owner_clients[username]
                for collection in collections:
                    target_model = self._get_target_model_for_source(collection.embedding_model)
                    work_items.append((user_client, collection, target_model, username, user_id))
                    work_owners.append((username, user_id))

            results = self._parallel_migrate_collections(work_items, total_collections)

            # Group results by owner for the report
            results_by_owner = defaultdict(list)
            for result, (username, user_id) in zip(results, work_owners):
                results_by_owner[username].append((result, user_id))
            for username, user_results in results_by_owner.items():
                user_migration = UserMigration(username=username, user_id=user_results[0][1])
                user_migration.collections = [r for r, _ in user_results]
                self.report.users.append(user_migration)
        else:
            current_idx = 0

            for username, collections in collections_by_owner.items():
                if username not in owner_clients:
                    self.log.warning(f"Skipping {len(collections)} collection(s) for {username} (no client)")
                    continue

                user_client, user_id = owner_clients[username]

                # Create user migration tracker
                user_migration = UserMigration(username=username, user_id=user_id)

                for collection in collections:
                    current_idx += 1
                    target_model = self._get_target_model_for_source(collection.embedding_model)

                    # Migrate the collection
                    migration_result = self._migrate_collection(
                        user_client=user_client,
                        collection=collection,
                        target_model=target_model,
                        current=current_idx,
                        total=total_collections,
                        username=username,
                        user_id=user_id
                    )
                    user_migration.collections.append(migration_result)

                self.report.users.append(user_migration)

    def _clone_client(self, client: H2OGPTE) -> H2OGPTE:
        """Create a new H2OGPTE client with the same credentials.

        H2OGPTE clients share a urllib3.PoolManager and ApiClient internally,
        which are not guaranteed to be thread-safe. Each worker thread gets its
        own client instance sharing the same API key to avoid shared mutable state.
        """
        return H2OGPTE(
            address=self.config.url,
            api_key=client._api_key,
            verify=self._get_ssl_verify_config(),
        )

    def _parallel_migrate_collections(self, work_items, total):
        """Process collections concurrently using a thread pool.

        Args:
            work_items: list of (user_client, collection, target_model, username, user_id)
            total: total count for progress display

        Returns:
            list of CollectionMigration results (same order as work_items)
        """
        results = [None] * len(work_items)
        progress_counter = itertools.count(1)

        self.log.info(f"Processing {total} collections with {self.config.max_concurrent_jobs} concurrent workers...")

        def worker(idx, user_client, collection, target_model, username, user_id):
            # Create a thread-local client to avoid sharing HTTP sessions across threads
            thread_client = self._clone_client(user_client)
            current = next(progress_counter)
            self.log.set_prefix(f"[{collection.name} ({collection.id})]")
            try:
                return self._migrate_collection(
                    user_client=thread_client,
                    collection=collection,
                    target_model=target_model,
                    current=current,
                    total=total,
                    username=username,
                    user_id=user_id,
                )
            finally:
                self.log.clear_prefix()

        self._executor = ThreadPoolExecutor(max_workers=self.config.max_concurrent_jobs)
        try:
            futures = {}
            for idx, (user_client, collection, target_model, username, user_id) in enumerate(work_items):
                future = self._executor.submit(worker, idx, user_client, collection, target_model, username, user_id)
                futures[future] = idx

            for future in as_completed(futures):
                idx = futures[future]
                try:
                    results[idx] = future.result()
                except Exception as e:
                    self.log.error(f"Unexpected worker failure: {e}")
                    user_client, collection, target_model, username, user_id = work_items[idx]
                    migration = CollectionMigration(
                        old_collection_id=collection.id,
                        old_collection_name=collection.name,
                        old_model=collection.embedding_model,
                        new_model=target_model,
                    )
                    migration.mark_failed(f"Worker thread error: {e}")
                    results[idx] = migration
        except KeyboardInterrupt:
            self.log.warning("Interrupted — canceling pending workers and waiting for running ones to finish...")
            for future in futures:
                future.cancel()
            raise
        finally:
            self._executor.shutdown(wait=True)
            self._executor = None

        return results

    def _migrate_collection(
        self, 
        user_client: H2OGPTE, 
        collection, 
        target_model: str,
        current: int, 
        total: int,
        username: str = None,
        user_id: str = None
    ) -> CollectionMigration:
        """Migrate a single collection."""
        
        migration = CollectionMigration(
            old_collection_id=collection.id,
            old_collection_name=collection.name,
            old_model=collection.embedding_model,
            new_model=target_model
        )
        
        try:
            self.log.progress(current, total, collection.name)
            self.log.info(f"  Model: {collection.embedding_model} → {target_model}")
            self.log.info(f"  Document count: {collection.document_count}")
            self.log.info(f"  Collection ID: {collection.id}")

            # --- Resume / skip logic ---
            resuming = False
            new_collection_id = None
            new_collection_name = None

            if self.db and not self.config.force_remigrate:
                existing = self.db.check_migration_exists(collection.id)
                if existing:
                    # Fully done — skip
                    chats_done = existing['chats_migrated'] or not self.config.migrate_chats
                    if existing['import_completed'] and existing['connectors_migrated'] and chats_done:
                        self.log.warning(f"  Already migrated (status: {existing['job_status']})")
                        self.log.info(f"    New collection: {existing['new_collection_name']} (ID: {existing['new_collection_id']})")
                        migration.mark_skipped("Already migrated")
                        return migration

                    # Import failed — need a new collection (can't reuse, would duplicate docs)
                    if existing['job_status'] == 'failed' and existing['import_submitted']:
                        if not self.config.retry_failed:
                            self.log.warning(f"  Previously failed: {existing['error']}")
                            self.log.info(f"    Use --retry-failed to retry")
                            migration.mark_skipped("Previously failed (use --retry-failed)")
                            return migration
                        self.log.info(f"  Retrying failed import (will create new collection)")
                        self.log.warning(f"    Previous new collection {existing['new_collection_id']} may need manual cleanup")
                        # Fall through to create a new collection

                    # Collection created but import never started — reuse it
                    elif existing['collection_created'] and not existing['import_submitted']:
                        self.log.info(f"  Resuming from existing collection: {existing['new_collection_id']}")
                        new_collection_id = existing['new_collection_id']
                        new_collection_name = existing['new_collection_name']
                        resuming = True

                    # Import submitted but not completed — tell user to verify
                    elif existing['import_submitted'] and not existing['import_completed']:
                        self.log.info(f"  Import job already submitted, use --verify to check status")
                        migration.mark_skipped("Import in progress")
                        return migration

                    # Import completed but post-import steps pending — handle them
                    elif existing['import_completed'] and (not existing['connectors_migrated'] or (not existing['chats_migrated'] and self.config.migrate_chats)):
                        self.log.info(f"  Import completed, running post-import steps...")
                        new_collection_id = existing['new_collection_id']
                        post_ok = self._run_post_import_steps(user_client, collection.id, new_collection_id, existing, indent="  ")
                        if post_ok:
                            migration.mark_success(
                                new_collection_id=new_collection_id,
                                new_collection_name=existing['new_collection_name'],
                                job_id=existing['job_id']
                            )
                        else:
                            migration.mark_failed("Post-import steps failed (connectors/chats)")
                        return migration

            # --- Fetch full collection details (needed for collection_settings) ---
            full_collection = user_client.get_collection(collection.id)
            old_settings = full_collection.collection_settings or {}

            # --- Pre-migration document status check ---
            pre_migration_bad_docs = self._check_source_doc_statuses(
                user_client, collection.id, collection.document_count, indent="  "
            )

            # Override ocr_model with CLI flag
            old_settings["ocr_model"] = self.config.ocr_model
            if self.config.ocr_model != "auto":
                self.log.info(f"  OCR model override: {self.config.ocr_model}")

            # --- Dry run ---
            if self.config.dry_run:
                new_name = f"(Migrated) {collection.name} - {datetime.now().strftime('%Y%m%d_%H%M%S')}"
                self.log.info(f"  [DRY-RUN] Would create: '{new_name}'")

                try:
                    user_perms = user_client.list_collection_permissions(collection.id)
                    group_perms = user_client.list_collection_group_permissions(collection.id)

                    if full_collection.is_public:
                        public_perms = user_client.list_collection_public_permissions(collection.id)
                        self.log.info(f"  [DRY-RUN] Would make collection public with permissions: {public_perms}")

                    non_owner_users = [p for p in user_perms if p.username != username]
                    if len(non_owner_users) > 0:
                        self.log.info(f"  [DRY-RUN] Would copy {len(non_owner_users)} user permission(s):")
                        for perm in non_owner_users:
                            self.log.info(f"    - {perm.username}: {', '.join(perm.permissions)}")

                    if len(group_perms) > 0:
                        self.log.info(f"  [DRY-RUN] Would copy {len(group_perms)} group permission(s):")
                        for perm in group_perms:
                            self.log.info(f"    - {perm.group_id}: {', '.join(perm.permissions)}")

                except Exception as e:
                    self.log.info(f"  Could not fetch permissions for dry-run preview: {e}")

                # Preview lifecycle settings that would be copied (server-side)
                lifecycle_items = []
                if full_collection.expiry_date:
                    lifecycle_items.append(f"expiry_date={full_collection.expiry_date}")
                if full_collection.inactivity_interval:
                    lifecycle_items.append(f"inactivity_interval={full_collection.inactivity_interval} days")
                if full_collection.size_limit:
                    lifecycle_items.append(f"size_limit={full_collection.size_limit} bytes")
                if lifecycle_items:
                    self.log.info(f"  [DRY-RUN] Would copy lifecycle settings: {', '.join(lifecycle_items)}")

                # Preview collection properties that would be copied
                if full_collection.prompt_template_id:
                    self.log.info(f"  [DRY-RUN] Would copy prompt_template_id: {full_collection.prompt_template_id}")
                if full_collection.workspace:
                    self.log.info(f"  [DRY-RUN] Would copy workspace: {full_collection.workspace}")
                if full_collection.rag_type:
                    self.log.info(f"  [DRY-RUN] Would copy rag_type: {full_collection.rag_type}")
                if full_collection.metadata_dict:
                    self.log.info(f"  [DRY-RUN] Would copy metadata: {full_collection.metadata_dict}")
                if old_settings:
                    self.log.info(f"  [DRY-RUN] Would copy collection_settings: {old_settings}")
                if full_collection.chat_settings:
                    self.log.info(f"  [DRY-RUN] Would copy chat_settings: {full_collection.chat_settings}")

                self.log.info(f"  [DRY-RUN] Would import {collection.document_count} documents with:")
                self.log.info(f"    - copy_document: {self.config.copy_document}")
                self.log.info(f"    - skip_reparse: {self.config.skip_reparse}")
                self.log.info(f"    - collection_settings used as default import job configs")
                if self.config.migrate_chats:
                    self.log.info(f"  [DRY-RUN] Would migrate chat sessions to new collection")
                migration.mark_skipped("Dry run mode")
                return migration

            # --- Step 1: Create new collection (or reuse from resume) ---
            if not resuming:
                self.log.info(f"  Creating new collection...")
                if old_settings:
                    self.log.info(f"  Collection settings: {old_settings}")
                if full_collection.chat_settings:
                    self.log.info(f"  Chat settings: {full_collection.chat_settings}")
                if full_collection.prompt_template_id:
                    self.log.info(f"  Prompt template: {full_collection.prompt_template_id}")
                if full_collection.workspace:
                    self.log.info(f"  Workspace: {full_collection.workspace}")
                new_collection_name = f"(Migrated) {collection.name} - {datetime.now().strftime('%Y%m%d_%H%M%S')}"
                try:
                    new_collection_id = user_client.create_collection(
                        name=new_collection_name,
                        description=full_collection.description or "",
                        embedding_model=target_model,
                        prompt_template_id=full_collection.prompt_template_id,
                        collection_settings=old_settings or None,
                        chat_settings=full_collection.chat_settings or None,
                        workspace=full_collection.workspace,
                    )
                except Exception as e:
                    if full_collection.workspace and "Invalid workspace" in str(e):
                        self.log.warning(f"  Could not use workspace '{full_collection.workspace}': {e}")
                        self.log.info(f"  Retrying without workspace, falling back to user's default workspace...")
                        new_collection_id = user_client.create_collection(
                            name=new_collection_name,
                            description=full_collection.description or "",
                            embedding_model=target_model,
                            prompt_template_id=full_collection.prompt_template_id,
                            collection_settings=old_settings or None,
                            chat_settings=full_collection.chat_settings or None,
                        )
                    else:
                        raise
                self.log.success(f"  New collection created: {new_collection_name} (ID: {new_collection_id})")

                # Copy rag_type (not available in create_collection)
                if full_collection.rag_type:
                    try:
                        user_client.update_collection_rag_type(
                            collection_id=new_collection_id,
                            name=new_collection_name,
                            description=full_collection.description or "",
                            rag_type=full_collection.rag_type,
                        )
                    except Exception as e:
                        self.log.warning(f"  Could not copy rag_type: {e}")

                # Copy custom metadata
                if full_collection.metadata_dict:
                    try:
                        user_client.update_collection_metadata(
                            collection_id=new_collection_id,
                            collection_metadata=full_collection.metadata_dict,
                        )
                    except Exception as e:
                        self.log.warning(f"  Could not copy metadata: {e}")

                # Save to DB IMMEDIATELY to prevent orphaned collections
                # Log recovery commands if force-remigrate is overwriting a previous migration
                if self.db and self.config.force_remigrate:
                    existing = self.db.check_migration_exists(collection.id)
                    if existing and existing.get('new_collection_id'):
                        prev_id = existing['new_collection_id']
                        if existing.get('connectors_migrated') or existing.get('chats_migrated'):
                            self.log.info(f"  Any connectors and/or chats were previously moved to {prev_id}")
                            self.log.info(f"  To recover: --move-connectors --move-chats --from {prev_id} --to {new_collection_id}")

                if self.db:
                    self.db.insert_migration({
                        'old_collection_id': collection.id,
                        'old_collection_name': collection.name,
                        'new_collection_id': new_collection_id,
                        'new_collection_name': new_collection_name,
                        'old_model': collection.embedding_model,
                        'new_model': target_model,
                        'job_status': 'pending',
                        'user_id': user_id,
                        'username': username,
                        'created_at': datetime.now(),
                        'collection_created': 1,
                    })

            # --- Step 2: Copy permissions (idempotent — safe to re-run on resume) ---
            self.log.info(f"  Copying permissions and sharing settings...")
            self._copy_collection_permissions(user_client, full_collection, new_collection_id)
            if self.db:
                self.db.mark_step_completed(collection.id, 'permissions_copied')

            # --- Step 3: Import documents + server-side settings (lifecycle, connectors, metadata) ---
            if self.config.wait_for_completion:
                self.log.info(f"  Importing documents (copy_document={self.config.copy_document}, skip_reparse={self.config.skip_reparse})...")

                job = user_client.import_collection_into_collection(
                    collection_id=new_collection_id,
                    src_collection_id=collection.id,
                    copy_document=self.config.copy_document,
                    skip_reparse=self.config.skip_reparse,
                    preserve_document_status=True,
                    preserve_metadata=True,
                    copy_lifecycle_settings=True,
                )
                self.log.info(f"  Import job finished: {job.id}")

                # Mark submitted + store job_id AFTER the call returns successfully
                if self.db:
                    self.db.update_job_status(collection.id, 'submitted', job_id=job.id)
                    self.db.mark_step_completed(collection.id, 'import_submitted')

                # Determine final status
                if job.canceled:
                    final_status = 'failed'
                    error = 'Job was canceled'
                elif job.failed or len(job.errors) > 0:
                    final_status = 'failed'
                    error = "; ".join(job.errors) if job.errors else f"{job.failed} document(s) failed"
                else:
                    final_status = 'completed'
                    error = None

                if final_status == 'completed':
                    new_doc_count = self._validate_migration(user_client, collection.id, new_collection_id, target_model, indent="  ", pre_migration_bad_docs=pre_migration_bad_docs)

                    if self.config.verify_query:
                        if new_doc_count > 0:
                            rag_ok = self._verify_rag(user_client, new_collection_id, new_collection_name, indent="  ")
                            if not rag_ok:
                                migration.rag_failed = True
                        else:
                            self.log.info(f"  Skipping RAG verification — no documents in collection")

                    if self.db:
                        self.db.mark_completed(collection.id)

                    # Post-import steps: connectors + chats (only after import succeeds)
                    job_record = {'connectors_migrated': 0, 'chats_migrated': 0}
                    post_ok = self._run_post_import_steps(user_client, collection.id, new_collection_id, job_record, indent="  ")

                    if post_ok:
                        migration.mark_success(
                            new_collection_id=new_collection_id,
                            new_collection_name=new_collection_name,
                            job_id=job.id
                        )
                    else:
                        migration.mark_failed("Import succeeded but post-import steps failed (connectors/chats)")
                else:
                    self.log.failure(f"  Import job failed: {error}")
                    if self.db:
                        self.db.update_job_status(collection.id, final_status, job_id=job.id, error=error)
                    migration.mark_failed(error)

            else:
                # Submit job without waiting (direct REST API call)
                self.log.info(f"  Submitting import job (copy_document={self.config.copy_document}, skip_reparse={self.config.skip_reparse})...")

                try:
                    header = user_client._get_auth_header()
                    with user_client._RESTClient(user_client) as rest_client:
                        response = rest_client.collection_api.create_import_collection_to_collection_job(
                            collection_id=new_collection_id,
                            create_import_collection_to_collection_job_request={
                                'source_collection_id': collection.id
                            },
                            copy_document=self.config.copy_document,
                            skip_reparse=self.config.skip_reparse,
                            preserve_document_status=True,
                            preserve_metadata=True,
                            copy_lifecycle_settings=True,
                            timeout=None,
                            _headers=header,
                        )

                    job_id = response.id
                    self.log.success(f"  Import job submitted: {job_id}")
                    if self.config.force_remigrate:
                        self.log.info(f"  Run --verify to check completion")
                    else:
                        self.log.info(f"  Run --verify to check completion and migrate connectors" +
                                      (" and chats" if self.config.migrate_chats else ""))

                    if self.db:
                        self.db.update_job_status(collection.id, 'submitted', job_id=job_id)
                        self.db.mark_step_completed(collection.id, 'import_submitted')

                except Exception as e:
                    self.log.error(f"  Failed to submit import job: {e}")
                    self.log.info(f"  REST API error details: {str(e)}")
                    raise

                migration.mark_success(
                    new_collection_id=new_collection_id,
                    new_collection_name=new_collection_name,
                    job_id=job_id
                )

        except Exception as e:
            self.log.failure(f"  Migration failed: {e}")
            self.log.info(f"  Error details: {str(e)}")
            migration.mark_failed(str(e))

        return migration
    
    def _verify_single_job(self, user_client: H2OGPTE, job_record: dict, counters: dict) -> None:
        """Verify a single import job and update database.
        
        Args:
            user_client: Client with access to the job
            job_record: Database record for the job
            counters: Dictionary of counters to update (completed, failed, etc.)
        """
        job_id = job_record['job_id']
        collection_name = job_record['old_collection_name']
        old_collection_id = job_record['old_collection_id']

        try:
            # Already-completed job pulled in for post-import steps
            if job_record['job_status'] == 'completed':
                self.log.info(f"  '{collection_name}' — import already completed")
                post_ok = self._run_post_import_steps(user_client, old_collection_id, job_record['new_collection_id'], job_record)
                counters['completed'] += 1
                if not post_ok:
                    counters['post_steps_failed'] += 1
                return

            self.log.info(f"  Checking job {job_id} for '{collection_name}'...")

            job = user_client.get_job(job_id)

            # Determine job status using helper method
            status_type, error_msg = self._determine_job_status(job)

            if status_type == 'canceled':
                self.db.update_job_status(old_collection_id, 'failed', error=error_msg)
                self.log.failure(f"    Job was canceled")
                counters['canceled'] += 1

            elif status_type == 'failed':
                self.db.update_job_status(old_collection_id, 'failed', error=error_msg)
                self.log.failure(f"    Completed with failures: {error_msg}")
                counters['failed'] += 1

            elif status_type == 'completed':
                self.db.mark_completed(old_collection_id)
                self.log.success(f"    Completed successfully!")

                new_doc_count = self._validate_migration(
                    user_client,
                    old_collection_id,
                    job_record['new_collection_id'],
                    job_record['new_model'],
                    indent="    "
                )

                if self.config.verify_query:
                    if new_doc_count > 0:
                        rag_ok = self._verify_rag(
                            user_client, job_record['new_collection_id'],
                            job_record['new_collection_name'], indent="    "
                        )
                        if not rag_ok:
                            counters['rag_failed'] += 1
                    else:
                        self.log.info(f"    Skipping RAG verification — no documents in collection")

                post_ok = self._run_post_import_steps(user_client, old_collection_id, job_record['new_collection_id'], job_record)
                counters['completed'] += 1
                if not post_ok:
                    counters['post_steps_failed'] += 1

            else:  # running
                progress_pct = int(job.progress * 100)
                if job.failed > 0:
                    self.log.warning(f"    Running ({progress_pct}%), {job.passed} passed, {job.failed} failed so far")
                else:
                    self.log.warning(f"    Running ({progress_pct}% complete)")
                self.db.update_job_status(old_collection_id, 'running', error=error_msg)
                counters['running'] += 1

        except Exception as e:
            self.log.warning(f"    Could not check job {job_id}: {e}")
            counters['not_checked'] += 1
    
    def _print_verification_summary(self, counters: dict):
        """Print verification results summary."""
        self.log.info("")
        self.log.info("Job Verification Summary:")
        if counters['completed'] > 0:
            self.log.success(f"  Completed: {counters['completed']}")
        if counters['failed'] > 0:
            self.log.failure(f"  Failed: {counters['failed']}")
        if counters['canceled'] > 0:
            self.log.failure(f"  Canceled: {counters['canceled']}")
        if counters['running'] > 0:
            self.log.warning(f"  Still Running: {counters['running']}")
        if counters['post_steps_failed'] > 0:
            self.log.warning(f"  Post-Import Steps Failed: {counters['post_steps_failed']} (connectors/chats — will retry on next --verify)")
        if counters.get('rag_failed', 0) > 0:
            self.log.warning(f"  RAG Verification Failed: {counters['rag_failed']}")
        if counters['not_checked'] > 0:
            self.log.warning(f"  Could Not Check: {counters['not_checked']}")
    
    def _verify_import_jobs(self) -> dict:
        """Verify import jobs and run post-import steps.

        Returns:
            Dictionary of counters (completed, failed, running, canceled, etc.)
        """
        self.log.section("Verifying Import Jobs")

        # Get pending jobs with appropriate filtering
        # Always include completed jobs needing post-import steps (connectors, chats)
        if self.is_self_service:
            pending_jobs = self.db.get_pending_jobs(
                usernames=[self.current_username],
                collection_ids=self.config.collection_ids,
                include_completed_needing_post_steps=True,
            )
        else:
            pending_jobs = self.db.get_pending_jobs(
                usernames=self.config.users,
                collection_ids=self.config.collection_ids,
                include_completed_needing_post_steps=True,
            )

        if not pending_jobs:
            self.log.info("No pending jobs to verify")
            return {'completed': 0, 'failed': 0, 'running': 0, 'canceled': 0, 'not_checked': 0, 'post_steps_failed': 0, 'rag_failed': 0}
        
        self.log.info(f"Checking status of {len(pending_jobs)} import job(s)...")
        self.log.info("")
        
        # Initialize counters
        counters = {
            'completed': 0,
            'failed': 0,
            'running': 0,
            'canceled': 0,
            'not_checked': 0,
            'post_steps_failed': 0,
            'rag_failed': 0,
        }
        
        if self.is_self_service:
            # Self-service: use existing client
            self.log.info(f"Verifying {len(pending_jobs)} job(s) for your account")
            
            for job_record in pending_jobs:
                self._verify_single_job(self.user_client, job_record, counters)
        else:
            # Admin: group by user and create clients
            jobs_by_user = defaultdict(list)
            for job_record in pending_jobs:
                jobs_by_user[job_record['username']].append(job_record)
            
            for username, user_jobs in jobs_by_user.items():
                self.log.info(f"Verifying {len(user_jobs)} job(s) for user: {username}")
                
                try:
                    user_id = user_jobs[0]['user_id']
                    user_client = self._get_user_client(user_id, username)
                    
                    for job_record in user_jobs:
                        self._verify_single_job(user_client, job_record, counters)
                        
                except Exception as e:
                    self.log.error(f"  Could not create client for user {username}: {e}")
                    self.log.warning(f"  Skipping verification for {len(user_jobs)} job(s)")
                    counters['not_checked'] += len(user_jobs)
        
        # Print verification summary
        self._print_verification_summary(counters)
        return counters

    def _copy_collection_permissions(
        self,
        user_client: H2OGPTE,
        old_collection,
        new_collection_id: str
    ):
        """Copy permissions from old collection to new collection.

        Args:
            user_client: H2OGPTE client with access to both collections.
            old_collection: Full Collection object (has is_public, username, id).
            new_collection_id: ID of the new collection.
        """
        old_collection_id = old_collection.id
        owner_username = old_collection.username
        stats = {'users': 0, 'groups': 0, 'is_public': False, 'errors': 0}

        # Copy public permissions
        if old_collection.is_public:
            try:
                public_perms = user_client.list_collection_public_permissions(old_collection_id)
                self.log.info(f"    Collection is public, copying permissions: {public_perms}")
                user_client.make_collection_public(new_collection_id, permissions=public_perms)
                stats['is_public'] = True
            except Exception as e:
                self.log.warning(f"    Failed to make collection public: {e}")
                stats['errors'] += 1

        # Copy user permissions
        try:
            self.log.info("    Copying user permissions...")
            user_permissions = user_client.list_collection_permissions(old_collection_id)
            non_owner_perms = [p for p in user_permissions if p.username != owner_username]

            if not non_owner_perms:
                self.log.info("      No user permissions to copy")

            for perm in non_owner_perms:
                try:
                    self.log.info(f"      User: {perm.username} - {', '.join(perm.permissions)}")
                    user_client.share_collection(new_collection_id, perm)
                    stats['users'] += 1
                except Exception as e:
                    self.log.warning(f"      Failed to copy permission for {perm.username}: {e}")
                    stats['errors'] += 1

        except Exception as e:
            self.log.warning(f"    Could not fetch user permissions: {e}")
            stats['errors'] += 1

        # Copy group permissions
        try:
            self.log.info("    Copying group permissions...")
            group_permissions = user_client.list_collection_group_permissions(old_collection_id)

            if not group_permissions:
                self.log.info("      No group permissions to copy")

            for perm in group_permissions:
                try:
                    self.log.info(f"      Group: {perm.group_id} - {', '.join(perm.permissions)}")
                    user_client.share_collection_with_group(new_collection_id, perm)
                    stats['groups'] += 1
                except Exception as e:
                    self.log.warning(f"      Failed to copy permission for {perm.group_id}: {e}")
                    stats['errors'] += 1

        except Exception as e:
            self.log.warning(f"    Could not fetch group permissions: {e}")
            stats['errors'] += 1

        # Log summary
        if stats['is_public']:
            self.log.info(f"    Made collection public")
        if stats['users'] > 0:
            self.log.info(f"    Copied permissions for {stats['users']} user(s)")
        if stats['groups'] > 0:
            self.log.info(f"    Copied permissions for {stats['groups']} group(s)")
        if stats['errors'] > 0:
            self.log.warning(f"    {stats['errors']} permission(s) failed to copy")

    def _get_all_users_including_self(self) -> List:
        """Get all users including the calling user.

        Note: list_users() excludes the calling user, so we add them manually.
        Results are cached to avoid redundant API calls.
        """
        if self._cached_users is not None:
            return self._cached_users

        users = self.admin_client.list_users(offset=0, limit=-1)

        # Add calling user (list_users excludes them)
        meta = self.admin_client.get_meta()
        calling_user = type('User', (), {
            'id': meta.user_id,
            'username': meta.username,
            'email': getattr(meta, 'email', '')
        })()
        users.append(calling_user)

        self._cached_users = users
        return self._cached_users
    
    def _find_user_by_username(self, username: str):
        """Find a user by username (includes calling user).
        
        Returns:
            User object if found, None otherwise.
        """
        users = self._get_all_users_including_self()
        return next((u for u in users if u.username == username), None)
    
    def _get_target_model_for_source(self, source_model: str) -> str:
        """Determine target model based on config and source model.
        
        Args:
            source_model: The source embedding model
            
        Returns:
            Target embedding model to migrate to
        """
        if self.config.use_model_mappings:
            return MODEL_MAPPINGS.get(source_model, self.config.target_model)
        return self.config.target_model

    def _run_move_operations(self):
        """Manually move connectors and/or chats between collections."""
        self.log.header("Manual Move Operations")
        source_id = self.config.move_from
        target_id = self.config.move_to
        all_ok = True

        if self.is_self_service:
            client = self.user_client
            self.log.info(f"Source: {source_id}")
            self.log.info(f"Target: {target_id}")
        else:
            # Admin: look up source collection to find owner, create API key
            try:
                results = self.admin_client.list_all_collections_sort(
                    offset=0, limit=1, sort_column="updated_at", ascending=False, filter=source_id
                )
                if not results or results[0].id != source_id:
                    self.log.failure(f"Source collection not found: {source_id}")
                    self.report.move_failed = True
                    return
                source_col = results[0]
                owner = source_col.username
                self.log.info(f"Source: {source_col.name} (ID: {source_id}, owner: {owner})")
                self.log.info(f"Target: {target_id}")

                user = self._find_user_by_username(owner)
                if not user:
                    self.log.failure(f"Could not find user: {owner}")
                    self.report.move_failed = True
                    return
                client = self._get_user_client(user.id, owner)
            except Exception as e:
                self.log.failure(f"Failed to set up move operation: {e}")
                self.report.move_failed = True
                return

        self.log.info("")

        if self.config.move_connectors:
            if not self._migrate_connectors(client, source_id, target_id, indent="  "):
                all_ok = False

        if self.config.move_chats:
            if not self._migrate_chat_sessions(client, source_id, target_id, indent="  "):
                all_ok = False

        if all_ok:
            self.log.info("Move operation completed")
        else:
            self.log.failure("Move operation completed with errors")
            self.report.move_failed = True

    def _run_post_import_steps(self, user_client: H2OGPTE, old_collection_id: str, new_collection_id: str, job_record: dict, indent: str = "    ") -> bool:
        """Run post-import steps (connectors then chats) for a completed migration.

        Chats are only migrated if connectors succeeded — ensures we don't
        move chats to a collection with broken/missing connectors.

        Returns:
            True if all required steps completed successfully, False otherwise.
        """
        connectors_done = job_record.get('connectors_migrated')

        if not connectors_done:
            connectors_ok = self._migrate_connectors(
                user_client, old_collection_id, new_collection_id, indent=indent
            )
            if connectors_ok and self.db:
                self.db.mark_step_completed(old_collection_id, 'connectors_migrated')
                connectors_done = True

        if not connectors_done:
            return False

        if self.config.migrate_chats and not job_record.get('chats_migrated'):
            chats_ok = self._migrate_chat_sessions(
                user_client, old_collection_id, new_collection_id, indent=indent
            )
            if chats_ok and self.db:
                self.db.mark_step_completed(old_collection_id, 'chats_migrated')
            if not chats_ok:
                return False

        return True

    def _migrate_connectors(self, user_client: H2OGPTE, old_collection_id: str, new_collection_id: str, indent: str = "  ") -> bool:
        """Move scheduled connectors from old collection to new collection.

        Args:
            user_client: H2OGPTE client (must be collection owner)
            old_collection_id: Source collection ID
            new_collection_id: Target collection ID
            indent: Indent string for log messages

        Returns:
            True if migration succeeded, False otherwise.
        """
        try:
            connectors = user_client.list_scheduled_connectors_for_collection(old_collection_id)
            if not connectors:
                self.log.info(f"{indent}No scheduled connectors to migrate")
                return True

            self.log.info(f"{indent}Migrating {len(connectors)} scheduled connector(s)...")
            connector_job = user_client.migrate_scheduled_connectors_to_collection(
                collection_id=new_collection_id,
                src_collection_id=old_collection_id,
            )
            if connector_job.completed and len(connector_job.errors) == 0:
                self.log.success(f"{indent}Scheduled connectors migrated successfully")
                return True
            else:
                error_detail = '; '.join(connector_job.errors) if connector_job.errors else 'Unknown error'
                self.log.warning(f"{indent}Connector migration had issues: {error_detail}")
                return False
        except Exception as e:
            self.log.warning(f"{indent}Could not migrate connectors: {e}")
            return False

    def _migrate_chat_sessions(self, user_client: H2OGPTE, old_collection_id: str, new_collection_id: str, indent: str = "  ") -> bool:
        """Migrate chat sessions from old collection to new collection.

        Args:
            user_client: H2OGPTE client (must be collection owner)
            old_collection_id: Source collection ID
            new_collection_id: Target collection ID
            indent: Indent string for log messages

        Returns:
            True if migration succeeded, False otherwise.
        """
        try:
            chat_count = user_client.count_chat_sessions_for_collection(old_collection_id)
            if chat_count == 0:
                self.log.info(f"{indent}No chat sessions to migrate")
                return True

            self.log.info(f"{indent}Migrating {chat_count} chat session(s)...")
            chat_job = user_client.migrate_chat_sessions_to_collection(
                source_collection_id=old_collection_id,
                target_collection_id=new_collection_id,
            )
            if chat_job.completed and len(chat_job.errors) == 0:
                self.log.success(f"{indent}Chat sessions migrated successfully")
                return True
            else:
                error_detail = '; '.join(chat_job.errors) if chat_job.errors else 'Unknown error'
                self.log.warning(f"{indent}Chat session migration had issues: {error_detail}")
                return False
        except Exception as e:
            self.log.warning(f"{indent}Could not migrate chat sessions: {e}")
            return False

    def _check_source_doc_statuses(self, user_client, collection_id, doc_count, indent="  "):
        """Check source collection for documents with non-standard statuses.

        Returns set of document names with non-standard status (not completed/agent_only).
        """
        if doc_count == 0:
            return set()
        try:
            docs = user_client.list_documents_in_collection(collection_id, offset=0, limit=doc_count or 1000)
            bad_docs = {}
            for doc in docs:
                if doc.status not in ('completed', 'agent_only'):
                    bad_docs[doc.name] = doc.status
            if bad_docs:
                self.log.warning(f"{indent}Source collection has {len(bad_docs)} document(s) with non-standard status:")
                for name, status in list(bad_docs.items())[:10]:
                    self.log.warning(f"{indent}  - {name}: {status}")
                if len(bad_docs) > 10:
                    self.log.warning(f"{indent}  ... and {len(bad_docs) - 10} more")
            return set(bad_docs.keys())
        except Exception as e:
            self.log.warning(f"{indent}Could not check source document statuses: {e}")
            return set()

    def _validate_migration(
        self,
        user_client: H2OGPTE,
        old_id: str,
        new_id: str,
        target_model: str,
        indent: str = "  ",
        pre_migration_bad_docs: set = None
    ) -> int:
        """Validate a completed migration by comparing old and new collections.

        Checks: document count, embedding model, lifecycle settings, and document statuses.
        All data comes from two get_collection() calls plus one list_documents_in_collection() call.

        Args:
            user_client: H2OGPTE client with access to both collections
            old_id: Old collection ID
            new_id: New collection ID
            target_model: Expected embedding model on the new collection
            indent: Indent string for log messages
            pre_migration_bad_docs: Set of document names that had non-standard status before migration

        Returns:
            The new collection's document count (-1 if validation failed).
        """
        if pre_migration_bad_docs is None:
            pre_migration_bad_docs = set()

        try:
            old_coll = user_client.get_collection(old_id)
            new_coll = user_client.get_collection(new_id)

            # Document count
            if old_coll.document_count == new_coll.document_count:
                self.log.success(f"{indent}Documents: {old_coll.document_count} → {new_coll.document_count}")
            else:
                self.log.warning(f"{indent}Document count mismatch: {old_coll.document_count} → {new_coll.document_count}")

            # Document statuses (check each doc is completed or agent_only)
            if new_coll.document_count > 0:
                try:
                    new_docs = user_client.list_documents_in_collection(new_id, offset=0, limit=new_coll.document_count or 1000)
                    bad_statuses = [
                        (doc.name, doc.status) for doc in new_docs
                        if doc.status not in ('completed', 'agent_only')
                    ]
                    if bad_statuses:
                        # Separate inherited (pre-existing) from new issues
                        inherited = [(n, s) for n, s in bad_statuses if n in pre_migration_bad_docs]
                        new_issues = [(n, s) for n, s in bad_statuses if n not in pre_migration_bad_docs]
                        valid_count = len(new_docs) - len(bad_statuses)

                        if new_issues:
                            self.log.warning(f"{indent}  Documents with unexpected status (new since migration):")
                            for name, status in new_issues[:10]:
                                self.log.warning(f"{indent}    - {name}: {status}")
                            if len(new_issues) > 10:
                                self.log.warning(f"{indent}    ... and {len(new_issues) - 10} more")
                        self.log.success(f"{indent}  {valid_count} of {len(new_docs)} documents have valid status (completed/agent_only)")
                        if inherited:
                            self.log.info(f"{indent}  {len(inherited)} document(s) were already non-standard before migration (unchanged)")
                    else:
                        self.log.success(f"{indent}All documents have valid status (completed/agent_only)")
                except Exception as e:
                    self.log.warning(f"{indent}Could not validate document statuses: {e}")

            # Embedding model
            if new_coll.embedding_model == target_model:
                self.log.success(f"{indent}Embedding model: {new_coll.embedding_model}")
            else:
                self.log.warning(f"{indent}Embedding model mismatch: expected {target_model}, got {new_coll.embedding_model}")

            # Lifecycle settings
            lifecycle_ok = True
            if old_coll.expiry_date != new_coll.expiry_date:
                self.log.warning(f"{indent}Expiry date mismatch: {old_coll.expiry_date} → {new_coll.expiry_date}")
                lifecycle_ok = False
            if old_coll.inactivity_interval != new_coll.inactivity_interval:
                self.log.warning(f"{indent}Inactivity interval mismatch: {old_coll.inactivity_interval} → {new_coll.inactivity_interval}")
                lifecycle_ok = False
            if old_coll.size_limit != new_coll.size_limit:
                self.log.warning(f"{indent}Size limit mismatch: {old_coll.size_limit} → {new_coll.size_limit}")
                lifecycle_ok = False
            if lifecycle_ok:
                self.log.success(f"{indent}Lifecycle settings match")

            return new_coll.document_count

        except Exception as e:
            self.log.warning(f"{indent}Could not validate migration: {e}")
            return -1

    def _verify_rag(self, user_client: H2OGPTE, collection_id: str, collection_name: str, indent: str = "    ") -> bool:
        """Verify a collection works with RAG by sending a test query.

        Creates a temporary chat session, sends the verify query,
        checks that the response has document references, and cleans up.

        Returns:
            True if RAG verification passed, False otherwise.
        """
        chat_session_id = None
        try:
            self.log.info(f"{indent}RAG verification: sending query to '{collection_name}'...")

            chat_session_id = user_client.create_chat_session(collection_id=collection_id)

            with user_client.connect(chat_session_id) as session:
                reply = session.query(
                    self.config.verify_query,
                    timeout=300,
                )

            if reply is None:
                self.log.warning(f"{indent}RAG verification: no response received")
                return False

            if reply.error:
                self.log.warning(f"{indent}RAG verification failed: {reply.error}")
                return False

            references = user_client.list_chat_message_references(reply.id)

            if len(references) == 0:
                self.log.warning(f"{indent}RAG verification: response had no document references")
                content_preview = (reply.content or "")[:200]
                self.log.info(f"{indent}  Response preview: {content_preview}...")
                return False

            wrong_collection_refs = [r for r in references
                                     if r.collection_id and r.collection_id != collection_id]
            if wrong_collection_refs:
                self.log.warning(f"{indent}RAG verification: {len(wrong_collection_refs)} reference(s) from wrong collection")
                return False

            self.log.success(f"{indent}RAG verification passed: {len(references)} reference(s) found")
            content_preview = (reply.content or "")[:200]
            self.log.info(f"{indent}  Response preview: {content_preview}...")
            return True

        except Exception as e:
            self.log.warning(f"{indent}RAG verification failed: {e}")
            return False
        finally:
            if chat_session_id:
                try:
                    user_client.delete_chat_sessions([chat_session_id])
                except Exception:
                    pass

    def _get_user_client(self, user_id: str, username: str) -> H2OGPTE:
        self.log.info(f"Creating API key for user: {username}")
        
        key_name = f"temp_migration_{username}_{random_suffix()}"
        
        user_api_key = self.admin_client.create_api_key_for_user(
            user_id=user_id,
            name=key_name,
            collection_id=None,
            expires_in=self.config.api_key_expiry
        )
        
        # Track username and exact key name for cleanup
        with self._key_lock:
            self.users_with_temp_keys.add(username)
            self.created_key_names.add(key_name)
        
        self.log.info(f"API key created: {key_name}")
        
        return H2OGPTE(
            address=self.config.url,
            api_key=user_api_key,
            verify=self._get_ssl_verify_config()
        )
    
    def _cleanup_api_keys(self):
        """Delete all temporary API keys created during migration."""
        if not self.created_key_names:
            self.log.info("No API keys to clean up")
            return
        
        if self.is_self_service:
            # No API keys created in self-service mode
            return
        
        self.log.section("Cleaning Up Generated API Keys")
        self.log.info(f"Cleaning up {len(self.created_key_names)} temporary key(s)...")
        
        key_ids_to_delete = []
        
        try:
            for username in self.users_with_temp_keys:
                try:
                    # Paginated fetch to handle large numbers of keys
                    offset = 0
                    limit = 100
                    
                    while True:
                        batch = self.admin_client.list_all_api_keys(
                            offset=offset,
                            limit=limit,
                            key_filter=username
                        )
                        
                        if not batch:
                            break
                        
                        # Match exact names from our tracking set
                        for key_obj in batch:
                            if key_obj.name in self.created_key_names:
                                key_ids_to_delete.append(key_obj.id)
                                self.log.info(f"  Found temp key: {key_obj.name} (ID: {key_obj.id})")
                        
                        offset += limit
                
                except Exception as e:
                    self.log.warning(f"  Could not list keys for {username}: {e}")
            
            # Delete all found keys in one batch
            if key_ids_to_delete:
                try:
                    self.admin_client.delete_api_keys(key_ids_to_delete)
                    self.log.success(f"Deleted {len(key_ids_to_delete)} temporary API key(s)")
                except Exception as e:
                    self.log.warning(f"Failed to delete API keys in batch: {e}")
            else:
                self.log.info("No temporary API keys found to delete")
                
        except Exception as e:
            self.log.warning(f"Could not clean up temporary API keys: {e}")
            self.log.info(f"Cleanup error details: {str(e)}")
    
    def _print_summary(self):
        summary = self.report.summary

        self.log.header("MIGRATION COMPLETE")
        self.log.info(f"Total Users: {summary['total_users']}")
        self.log.info(f"Total Collections: {summary['total_collections']}")

        if self.config.wait_for_completion:
            self.log.success(f"Completed: {summary['successful_migrations']}")
        else:
            self.log.success(f"Successfully Submitted: {summary['successful_migrations']}")

        if summary['failed_migrations'] > 0:
            self.log.failure(f"Failed: {summary['failed_migrations']}")

        if summary['skipped_migrations'] > 0:
            self.log.warning(f"Skipped: {summary['skipped_migrations']}")

        if summary.get('rag_failures', 0) > 0:
            self.log.warning(f"RAG Verification Failed: {summary['rag_failures']}")

        if summary['duration']:
            self.log.info(f"Duration: {summary['duration']}")

        if not self.config.wait_for_completion and summary['successful_migrations'] > 0:
            self.log.info("")
            if self.config.force_remigrate:
                self.log.info("Run --verify to check job completion. Use --move-connectors/--move-chats to recover connectors and chats (see commands above)")
            else:
                self.log.info("Run --verify to check job completion and migrate connectors" +
                              (", and --migrate-chats to migrate chat sessions" if self.config.migrate_chats else ""))

        self.log.info("")
        self.log.info(f"Detailed log: {self.log.get_log_file()}")
