migrate.py

#!/usr/bin/env python3
"""
WebProtégé MongoDB → Keycloak User Migration Script

Usage:
    python migrate.py                    # Run migration with default config
    python migrate.py --dry-run          # Simulate without touching Keycloak
    python migrate.py --reset            # Reset progress and start fresh
    python migrate.py --config my.yaml   # Use a custom config file
"""

import argparse
import os
import re
import time

import yaml
from dotenv import load_dotenv

from migration.keycloak_client import KeycloakClient
from migration.logger_setup import setup_logging
from migration.mongo_reader import MongoReader
from migration.progress_tracker import ProgressTracker
from migration.user_filter import UserFilter, MIGRATE, DISABLE, SKIP
from migration.user_transformer import transform_user


def _resolve_env_vars(value):
    """Replace ${VAR} patterns with environment variable values."""
    if isinstance(value, str):
        return re.sub(
            r"\$\{(\w+)\}",
            lambda m: os.environ.get(m.group(1), m.group(0)),
            value,
        )
    if isinstance(value, dict):
        return {k: _resolve_env_vars(v) for k, v in value.items()}
    if isinstance(value, list):
        return [_resolve_env_vars(item) for item in value]
    return value


def load_config(path: str) -> dict:
    with open(path, "r") as f:
        config = yaml.safe_load(f)
    return _resolve_env_vars(config)


def main():
    parser = argparse.ArgumentParser(
        description="Migrate users from MongoDB to Keycloak"
    )
    parser.add_argument(
        "--config",
        default="config.yaml",
        help="Path to configuration file (default: config.yaml)",
    )
    parser.add_argument(
        "--dry-run",
        action="store_true",
        help="Simulate migration without creating users in Keycloak",
    )
    parser.add_argument(
        "--reset",
        action="store_true",
        help="Reset progress state and start fresh",
    )
    args = parser.parse_args()

    # Load environment variables from .env file (if present)
    load_dotenv()

    # Load config
    config = load_config(args.config)
    log_dir = config.get("progress", {}).get("log_dir", "logs")
    state_file = config.get("progress", {}).get("state_file", "migration_state.json")

    # Set up logging
    logger, skipped_csv, failed_csv, disabled_csv, run_dir = setup_logging(log_dir)

    logger.info("=" * 60)
    logger.info("WebProtégé User Migration")
    logger.info("Mode: %s", "DRY RUN" if args.dry_run else "LIVE")
    logger.info("=" * 60)

    # Progress tracker
    tracker = ProgressTracker(state_file)
    if args.reset:
        tracker.reset()
        logger.info("Progress reset — starting fresh")

    # Connect to MongoDB
    mongo = MongoReader(config["mongodb"])
    mongo.connect()

    total_in_mongo = mongo.get_total_count()
    logger.info("Total users in MongoDB: %d", total_in_mongo)

    # Build duplicate email map and resolve winners by project count
    logger.info("Scanning for duplicate emails...")
    dup_map = mongo.get_all_emails()
    logger.info("Loading project counts for duplicate resolution...")
    project_counts = mongo.get_project_counts()
    for email, usernames in dup_map.items():
        usernames.sort(key=lambda u: (-project_counts.get(u, 0), u))

    # Initialize filter engine
    user_filter = UserFilter(
        config.get("filters", []),
        duplicate_email_map=dup_map,
    )

    # Initialize Keycloak client
    kc = None
    if not args.dry_run:
        kc = KeycloakClient(config["keycloak"])
        logger.info("Keycloak client initialized for %s", config["keycloak"]["base_url"])

    # Migration loop
    resume_id = tracker.last_processed_id
    batch_num = 0
    start_time = time.time()

    try:
        for batch in mongo.iter_batches(resume_after_id=resume_id):
            batch_num += 1
            batch_created = 0
            batch_disabled = 0
            batch_skipped = 0
            batch_failed = 0
            batch_duplicates = 0
            last_id = None

            for doc in batch:
                username = doc.get("_id", "")
                email = doc.get("emailAddress", "")
                last_id = username

                # Apply filters
                action, filter_name, reason = user_filter.apply(doc)
                if action == SKIP:
                    batch_skipped += 1
                    skipped_csv.log(
                        username=username,
                        email=email,
                        filter_name=filter_name,
                        reason=reason,
                    )
                    continue

                # Transform (disabled if filter says so)
                user_enabled = action == MIGRATE
                payload = transform_user(doc, user_enabled=user_enabled)

                if action == DISABLE:
                    disabled_csv.log(
                        username=username,
                        email=email,
                        filter_name=filter_name,
                        reason=reason,
                    )

                if args.dry_run:
                    if action == DISABLE:
                        batch_disabled += 1
                    else:
                        batch_created += 1
                    continue

                # Create in Keycloak
                result = kc.create_user(payload)

                if result["status"] == "created":
                    if action == DISABLE:
                        batch_disabled += 1
                    else:
                        batch_created += 1
                elif result["status"] == "duplicate":
                    batch_duplicates += 1
                    logger.debug(
                        "Duplicate in Keycloak: %s — %s", username, result["detail"]
                    )
                else:
                    batch_failed += 1
                    failed_csv.log(
                        username=username,
                        email=email,
                        error=result["detail"],
                    )
                    logger.warning(
                        "Failed to create user '%s': %s", username, result["detail"]
                    )

            # Update progress after each batch
            if last_id is not None:
                tracker.update_batch(
                    last_id=last_id,
                    created=batch_created,
                    disabled=batch_disabled,
                    skipped=batch_skipped,
                    failed=batch_failed,
                    duplicates=batch_duplicates,
                )

            summary = tracker.summary()
            elapsed = time.time() - start_time
            rate = summary["total_processed"] / elapsed if elapsed > 0 else 0

            logger.info(
                "Batch %d complete | Processed: %d | Created: %d | "
                "Disabled: %d | Skipped: %d | Failed: %d | Dups: %d | Rate: %.0f users/s",
                batch_num,
                summary["total_processed"],
                summary["total_created"],
                summary["total_disabled"],
                summary["total_skipped"],
                summary["total_failed"],
                summary["total_duplicates"],
                rate,
            )

    except KeyboardInterrupt:
        logger.info("Migration interrupted by user — progress saved")
        logger.info("Re-run the script to resume from where it left off")
    finally:
        mongo.close()
        skipped_csv.close()
        failed_csv.close()
        disabled_csv.close()

    # Final summary
    summary = tracker.summary()
    elapsed = time.time() - start_time
    logger.info("=" * 60)
    logger.info("Migration %s", "simulation complete" if args.dry_run else "complete")
    logger.info("Total processed: %d", summary["total_processed"])
    logger.info("  Created:    %d", summary["total_created"])
    logger.info("  Disabled:   %d", summary["total_disabled"])
    logger.info("  Skipped:    %d", summary["total_skipped"])
    logger.info("  Failed:     %d", summary["total_failed"])
    logger.info("  Duplicates: %d", summary["total_duplicates"])
    logger.info("Elapsed time: %.1f seconds", elapsed)
    logger.info("Logs directory: %s", os.path.abspath(run_dir))
    logger.info("=" * 60)

    if not args.dry_run and kc:
        try:
            kc_count = kc.get_user_count()
            logger.info("Keycloak user count (realm '%s'): %d", config["keycloak"]["realm"], kc_count)
        except Exception as e:
            logger.warning("Could not fetch Keycloak user count: %s", e)


if __name__ == "__main__":
    main()