#!/usr/bin/env python3
"""
migrate-big-keys.py - Migrate keys that are too big for RIOT

Usage:
    ./migrate-big-keys.py <source> <target> --mem-limit 512MB [OPTIONS]
    ./migrate-big-keys.py <source> <target> --key-file missing-keys.log [OPTIONS]

Requirements:
    - Python 3.7+
    - redis-py (pip install redis)
"""

import argparse
import re
import sys
import time
from pathlib import Path
from typing import List, Tuple, Optional

import redis


# Keys larger than this will use MIGRATE command instead of dump/restore
MIGRATE_THRESHOLD_BYTES = 2 * 1024**3  # 2GB


class Colors:
    RED = "\033[0;31m"
    GREEN = "\033[0;32m"
    YELLOW = "\033[1;33m"
    BLUE = "\033[0;34m"
    NC = "\033[0m"


def log_info(msg: str) -> None:
    print(f"{Colors.GREEN}[INFO]{Colors.NC} {msg}")


def log_warn(msg: str) -> None:
    print(f"{Colors.YELLOW}[WARN]{Colors.NC} {msg}")


def log_error(msg: str) -> None:
    print(f"{Colors.RED}[ERROR]{Colors.NC} {msg}", file=sys.stderr)


def log_step(msg: str) -> None:
    print(f"{Colors.BLUE}[STEP]{Colors.NC} {msg}")


def parse_mem_limit(limit_str: str) -> int:
    match = re.match(r"^(\d+)\s*(B|KB?|MB?|GB?)?$", limit_str.upper())
    if not match:
        raise ValueError(f"Invalid memory limit format: {limit_str}")
    value = int(match.group(1))
    unit = match.group(2) or "B"
    multipliers = {"B": 1, "K": 1024, "KB": 1024, "M": 1024**2, "MB": 1024**2, "G": 1024**3, "GB": 1024**3}
    return value * multipliers[unit]


def format_bytes(num_bytes: int) -> str:
    if num_bytes >= 1024**3:
        return f"{num_bytes / 1024**3:.2f}GB"
    elif num_bytes >= 1024**2:
        return f"{num_bytes / 1024**2:.2f}MB"
    elif num_bytes >= 1024:
        return f"{num_bytes / 1024:.2f}KB"
    return f"{num_bytes}B"


def load_keys_from_file(path: str) -> List[str]:
    """Load unique keys from a file, treating each non-empty line as a complete key."""
    key_file = Path(path)
    try:
        lines = key_file.read_text(encoding="utf-8").splitlines()
    except OSError as e:
        raise ValueError(f"Cannot read key file {path}: {e}") from e

    keys = []
    seen = set()
    for line_number, raw_line in enumerate(lines, start=1):
        line = raw_line.strip()
        if not line:
            continue

        key = line
        if key in seen:
            continue

        seen.add(key)
        keys.append(key)

    if not keys:
        raise ValueError(f"No keys found in key file: {path}")

    return keys


def scan_big_keys(client: redis.Redis, mem_limit_bytes: int, pattern: str, scan_count: int, verbose: bool) -> List[Tuple[bytes, int]]:
    big_keys = []
    total_keys = 0
    for key in client.scan_iter(match=pattern, count=scan_count):
        total_keys += 1
        try:
            mem_usage = client.memory_usage(key)
            if mem_usage and mem_usage > mem_limit_bytes:
                big_keys.append((key, mem_usage))
                if verbose:
                    log_info(f"Big key found: {key.decode('utf-8', errors='replace')} ({format_bytes(mem_usage)})")
        except (redis.RedisError, OSError) as e:
            log_warn(f"Failed to get memory usage for {key}: {e}")
        if verbose and total_keys % 10000 == 0:
            log_info(f"Scanned {total_keys} keys, found {len(big_keys)} big keys so far...")
    log_info(f"Scanned {total_keys} keys, found {len(big_keys)} keys exceeding limit")
    return big_keys


def get_key_type(client: redis.Redis, key: bytes) -> Optional[str]:
    """Get the type of a key."""
    try:
        return client.type(key).decode("utf-8")
    except redis.RedisError:
        return None


def execute_with_retry(func, max_retries: int = 3, retry_delay: float = 1.0):
    """Execute a function with retry logic for connection errors."""
    last_error = None
    for attempt in range(max_retries):
        try:
            return func()
        except (redis.ConnectionError, redis.TimeoutError, BrokenPipeError, ConnectionResetError) as e:
            last_error = e
            if attempt < max_retries - 1:
                log_warn(f"Connection error (attempt {attempt + 1}/{max_retries}): {e}")
                time.sleep(retry_delay * (attempt + 1))
            else:
                raise
    raise last_error


def migrate_sorted_set_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                                replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a sorted set in chunks using ZSCAN with pagination."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        # Get TTL first
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        # Get total size for progress reporting
        total_members = source.zcard(key)
        if total_members == 0:
            log_warn(f"Sorted set is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating sorted set: {key_str} ({total_members} members)")

        # Check if key exists on target when replace=False
        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        # Use ZSCAN for memory-efficient iteration
        migrated = 0
        cursor = 0
        last_progress_time = time.time()

        # Use smaller pipeline batch size for very large keys to avoid buffer issues
        pipeline_batch_size = min(chunk_size, 1000)

        while True:
            cursor, members = execute_with_retry(
                lambda c=cursor: source.zscan(key, cursor=c, count=chunk_size)
            )

            # Process members in smaller pipeline batches
            for i in range(0, len(members), pipeline_batch_size):
                batch = members[i:i + pipeline_batch_size]
                if batch:
                    def write_batch(b=batch):
                        pipe = target.pipeline(transaction=False)
                        for member, score in b:
                            pipe.zadd(key, {member: score})
                        pipe.execute()
                    execute_with_retry(write_batch)
                    migrated += len(batch)

                    # Progress update every 5 seconds or on completion
                    now = time.time()
                    if verbose and (now - last_progress_time >= 5):
                        log_info(f"  Progress: {migrated}/{total_members} ({100*migrated/total_members:.1f}%)")
                        last_progress_time = now

            if cursor == 0:
                break

        # Restore TTL if it was set
        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {migrated} members migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate sorted set {key_str}: {e}")
        return False


def migrate_hash_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                          replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a hash in chunks using HSCAN."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        total_fields = source.hlen(key)
        if total_fields == 0:
            log_warn(f"Hash is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating hash: {key_str} ({total_fields} fields)")

        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        migrated = 0
        cursor = 0
        last_progress_time = time.time()
        pipeline_batch_size = min(chunk_size, 1000)

        while True:
            try:
                cursor, fields = execute_with_retry(
                    lambda c=cursor: source.hscan(key, cursor=c, count=chunk_size)
                )
            except (redis.RedisError, OSError) as e:
                log_error(f"Failed to read from source (HSCAN) for hash {key_str}: {e}")
                return False

            # Convert to list for batching
            field_items = list(fields.items())
            for i in range(0, len(field_items), pipeline_batch_size):
                batch = dict(field_items[i:i + pipeline_batch_size])
                if batch:
                    try:
                        execute_with_retry(lambda b=batch: target.hset(key, mapping=b))
                    except (redis.RedisError, OSError) as e:
                        log_error(f"Failed to write to target (HSET) for hash {key_str}: {e}")
                        return False
                    migrated += len(batch)

                    now = time.time()
                    if verbose and (now - last_progress_time >= 5):
                        log_info(f"  Progress: {migrated}/{total_fields} ({100*migrated/total_fields:.1f}%)")
                        last_progress_time = now

            if cursor == 0:
                break

        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {migrated} fields migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate hash {key_str}: {e}")
        return False


def migrate_set_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                         replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a set in chunks using SSCAN."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        total_members = source.scard(key)
        if total_members == 0:
            log_warn(f"Set is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating set: {key_str} ({total_members} members)")

        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        migrated = 0
        cursor = 0
        last_progress_time = time.time()
        pipeline_batch_size = min(chunk_size, 1000)

        while True:
            cursor, members = execute_with_retry(
                lambda c=cursor: source.sscan(key, cursor=c, count=chunk_size)
            )

            members_list = list(members)
            for i in range(0, len(members_list), pipeline_batch_size):
                batch = members_list[i:i + pipeline_batch_size]
                if batch:
                    execute_with_retry(lambda b=batch: target.sadd(key, *b))
                    migrated += len(batch)

                    now = time.time()
                    if verbose and (now - last_progress_time >= 5):
                        log_info(f"  Progress: {migrated}/{total_members} ({100*migrated/total_members:.1f}%)")
                        last_progress_time = now

            if cursor == 0:
                break

        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {migrated} members migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate set {key_str}: {e}")
        return False


def migrate_list_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                          replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a list in chunks using LRANGE with pagination."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        total_elements = source.llen(key)
        if total_elements == 0:
            log_warn(f"List is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating list: {key_str} ({total_elements} elements)")

        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        migrated = 0
        start = 0
        last_progress_time = time.time()
        pipeline_batch_size = min(chunk_size, 1000)

        while start < total_elements:
            end = start + chunk_size - 1
            elements = execute_with_retry(lambda s=start, e=end: source.lrange(key, s, e))

            if elements:
                for i in range(0, len(elements), pipeline_batch_size):
                    batch = elements[i:i + pipeline_batch_size]
                    if batch:
                        execute_with_retry(lambda b=batch: target.rpush(key, *b))
                        migrated += len(batch)

                        now = time.time()
                        if verbose and (now - last_progress_time >= 5):
                            log_info(f"  Progress: {migrated}/{total_elements} ({100*migrated/total_elements:.1f}%)")
                            last_progress_time = now

            start += chunk_size

        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {migrated} elements migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate list {key_str}: {e}")
        return False


def migrate_stream_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                            replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a stream in chunks using XRANGE with pagination."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        total_entries = source.xlen(key)
        if total_entries == 0:
            log_warn(f"Stream is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating stream: {key_str} ({total_entries} entries)")

        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        migrated = 0
        last_id = "-"
        last_progress_time = time.time()
        pipeline_batch_size = min(chunk_size, 1000)

        while True:
            entries = execute_with_retry(lambda lid=last_id: source.xrange(key, min=lid, count=chunk_size))

            if not entries:
                break

            for i in range(0, len(entries), pipeline_batch_size):
                batch = entries[i:i + pipeline_batch_size]
                def write_stream_batch(b=batch):
                    pipe = target.pipeline(transaction=False)
                    for entry_id, fields in b:
                        entry_id_str = entry_id.decode("utf-8") if isinstance(entry_id, bytes) else entry_id
                        decoded_fields = {}
                        for k, v in fields.items():
                            k_str = k.decode("utf-8") if isinstance(k, bytes) else k
                            decoded_fields[k_str] = v
                        pipe.xadd(key, decoded_fields, id=entry_id_str)
                    pipe.execute()
                execute_with_retry(write_stream_batch)
                migrated += len(batch)

                now = time.time()
                if verbose and (now - last_progress_time >= 5):
                    log_info(f"  Progress: {migrated}/{total_entries} ({100*migrated/total_entries:.1f}%)")
                    last_progress_time = now

            # Get the last ID for pagination (exclusive for next iteration)
            last_entry_id = entries[-1][0]
            if isinstance(last_entry_id, bytes):
                last_entry_id = last_entry_id.decode("utf-8")
            # Use exclusive range prefix for next iteration
            last_id = f"({last_entry_id}"

            if len(entries) < chunk_size:
                break

        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {migrated} entries migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate stream {key_str}: {e}")
        return False


def migrate_string_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                            replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a string in chunks using GETRANGE and APPEND."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        ttl_ms = max(0, pttl) if pttl > 0 else 0

        total_length = source.strlen(key)
        if total_length == 0:
            log_warn(f"String is empty: {key_str}")
            return False

        if verbose:
            log_info(f"Migrating string: {key_str} ({format_bytes(total_length)})")

        if replace:
            target.delete(key)
        elif target.exists(key):
            log_error(f"Key already exists on target (use --replace to overwrite): {key_str}")
            return False

        # Use 64MB chunks for strings (chunk_size param is for collection items, not bytes)
        byte_chunk_size = 64 * 1024 * 1024  # 64MB
        migrated = 0
        last_progress_time = time.time()
        first_chunk = True

        while migrated < total_length:
            end = min(migrated + byte_chunk_size - 1, total_length - 1)
            chunk = execute_with_retry(lambda s=migrated, e=end: source.getrange(key, s, e))

            if chunk:
                if first_chunk:
                    # First chunk: use SET to create the key
                    execute_with_retry(lambda c=chunk: target.set(key, c))
                    first_chunk = False
                else:
                    # Subsequent chunks: use APPEND
                    execute_with_retry(lambda c=chunk: target.append(key, c))

                migrated += len(chunk)

                now = time.time()
                if verbose and (now - last_progress_time >= 5):
                    log_info(f"  Progress: {format_bytes(migrated)}/{format_bytes(total_length)} ({100*migrated/total_length:.1f}%)")
                    last_progress_time = now
            else:
                break

        if ttl_ms > 0:
            target.pexpire(key, ttl_ms)

        if verbose:
            log_info(f"  Completed: {format_bytes(migrated)} migrated")
        return True

    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate string {key_str}: {e}")
        return False


def migrate_key_chunked(source: redis.Redis, target: redis.Redis, key: bytes,
                         mem_usage: int, replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a key using chunked approach based on its type."""
    key_str = key.decode("utf-8", errors="replace")
    key_type = get_key_type(source, key)

    if key_type is None:
        log_warn(f"Key no longer exists: {key_str}")
        return False

    if verbose:
        log_info(f"Using chunked migration for {key_type} key: {key_str} ({format_bytes(mem_usage)})")

    if key_type == "zset":
        return migrate_sorted_set_chunked(source, target, key, replace, verbose, chunk_size)
    elif key_type == "hash":
        return migrate_hash_chunked(source, target, key, replace, verbose, chunk_size)
    elif key_type == "set":
        return migrate_set_chunked(source, target, key, replace, verbose, chunk_size)
    elif key_type == "list":
        return migrate_list_chunked(source, target, key, replace, verbose, chunk_size)
    elif key_type == "stream":
        return migrate_stream_chunked(source, target, key, replace, verbose, chunk_size)
    elif key_type == "string":
        return migrate_string_chunked(source, target, key, replace, verbose, chunk_size)
    else:
        log_warn(f"Unsupported key type '{key_type}' for chunked migration: {key_str}")
        log_warn(f"  Attempting standard dump/restore (may fail for very large keys)")
        return migrate_key_dump_restore(source, target, key, replace, verbose)


def migrate_key_dump_restore(source: redis.Redis, target: redis.Redis, key: bytes, replace: bool, verbose: bool) -> bool:
    """Migrate a key using dump/restore approach (for smaller keys)."""
    key_str = key.decode("utf-8", errors="replace")
    try:
        pttl = source.pttl(key)
        if pttl == -2:
            log_warn(f"Key no longer exists: {key_str}")
            return False
        ttl_ms = max(0, pttl)
        dump_data = source.dump(key)
        if dump_data is None:
            log_warn(f"Key no longer exists: {key_str}")
            return False
        if verbose:
            log_info(f"Migrating: {key_str} (dump: {format_bytes(len(dump_data))}, TTL: {ttl_ms}ms)")
        target.restore(key, ttl_ms, dump_data, replace=replace)
        if verbose:
            log_info(f"  Restored successfully")
        return True
    except (redis.RedisError, OSError) as e:
        log_error(f"Failed to migrate key {key_str}: {e}")
        return False


def migrate_key(source: redis.Redis, target: redis.Redis, key: bytes,
                mem_usage: int, replace: bool, verbose: bool, chunk_size: int = 10000) -> bool:
    """Migrate a key, choosing the appropriate method based on size."""
    if mem_usage > MIGRATE_THRESHOLD_BYTES:
        return migrate_key_chunked(source, target, key, mem_usage, replace, verbose, chunk_size)
    return migrate_key_dump_restore(source, target, key, replace, verbose)


def collect_explicit_keys(source: redis.Redis, keys: List[str]) -> List[Tuple[bytes, int]]:
    """Resolve explicit key names to byte keys and memory usage."""
    resolved_keys = []
    missing_keys = 0

    for key in keys:
        key_bytes = key.encode("utf-8")
        mem_usage = source.memory_usage(key_bytes)
        if mem_usage is None:
            log_warn(f"Key does not exist: {key}")
            missing_keys += 1
            continue
        resolved_keys.append((key_bytes, mem_usage))

    if missing_keys:
        log_warn(f"Skipped {missing_keys} missing keys")

    return resolved_keys


def main() -> int:
    parser = argparse.ArgumentParser(description="Migrate keys exceeding memory limit.")
    parser.add_argument("source", help="Source Redis URI")
    parser.add_argument("target", help="Target Redis URI")
    parser.add_argument("--mem-limit", help="Memory limit (e.g., 512MB, 1GB). Required unless --key or --key-file is used.")
    parser.add_argument("--scan-count", type=int, default=100)
    parser.add_argument("--key-pattern", default="*")
    parser.add_argument("--key", help="Migrate a specific key (skips scanning)")
    parser.add_argument("--key-file", help="Migrate keys listed in a file; each non-empty line is treated as a complete key")
    parser.add_argument("--replace", action="store_true")
    parser.add_argument("--dry-run", action="store_true")
    parser.add_argument("--verbose", action="store_true")
    parser.add_argument("--chunk-size", type=int, default=10000,
                        help="Batch size for chunked migration (default: 10000)")
    args = parser.parse_args()

    explicit_key_options = sum(1 for option in (args.key, args.key_file) if option)
    if explicit_key_options > 1:
        log_error("Use only one of --key or --key-file")
        return 1

    # mem-limit is not required when an explicit key selection is specified
    if not args.mem_limit and not args.key and not args.key_file:
        log_error("--mem-limit is required unless --key or --key-file is specified")
        return 1

    mem_limit_bytes = 0
    if args.mem_limit:
        try:
            mem_limit_bytes = parse_mem_limit(args.mem_limit)
        except ValueError as e:
            log_error(str(e))
            return 1
        log_info(f"Memory limit: {args.mem_limit} ({mem_limit_bytes} bytes)")

    # Add redis:// prefix if missing
    source_uri = args.source if "://" in args.source else f"redis://{args.source}"
    target_uri = args.target if "://" in args.target else f"redis://{args.target}"

    log_step("Testing Redis connections...")
    # Use long socket timeouts for large key operations
    socket_timeout = 3600  # 1 hour
    socket_connect_timeout = 30
    try:
        source = redis.Redis.from_url(
            source_uri,
            decode_responses=False,
            socket_timeout=socket_timeout,
            socket_connect_timeout=socket_connect_timeout,
            socket_keepalive=True,
            health_check_interval=30
        )
        source.ping()
    except (redis.RedisError, OSError) as e:
        log_error(f"Cannot connect to source: {e}")
        return 1
    try:
        target = redis.Redis.from_url(
            target_uri,
            decode_responses=False,
            socket_timeout=socket_timeout,
            socket_connect_timeout=socket_connect_timeout,
            socket_keepalive=True,
            health_check_interval=30
        )
        target.ping()
    except (redis.RedisError, OSError) as e:
        log_error(f"Cannot connect to target: {e}")
        return 1
    log_info("Both Redis connections successful")

    # If explicit keys are provided, migrate just those keys
    if args.key:
        big_keys = collect_explicit_keys(source, [args.key])
        if not big_keys:
            log_error(f"Key does not exist: {args.key}")
            return 1
        log_info(f"Key: {args.key} ({format_bytes(big_keys[0][1])})")
    elif args.key_file:
        try:
            requested_keys = load_keys_from_file(args.key_file)
        except ValueError as e:
            log_error(str(e))
            return 1

        log_info(f"Loaded {len(requested_keys)} unique keys from {args.key_file}")
        big_keys = collect_explicit_keys(source, requested_keys)
        if not big_keys:
            log_error("No listed keys exist on source")
            return 1
        log_info(f"Resolved {len(big_keys)} keys from key file")
    else:
        log_step(f"Scanning for keys exceeding {args.mem_limit}...")
        big_keys = scan_big_keys(source, mem_limit_bytes, args.key_pattern, args.scan_count, args.verbose)
        if not big_keys:
            log_info("No keys to migrate")
            return 0

    if args.dry_run:
        log_info(f"[DRY RUN] Would migrate {len(big_keys)} keys:")
        for key, mem in big_keys:
            key_str = key.decode("utf-8", errors="replace")
            key_type = get_key_type(source, key) or "unknown"
            method = "chunked" if mem > MIGRATE_THRESHOLD_BYTES else "dump/restore"
            print(f"  {key_str} ({format_bytes(mem)}) [{key_type}] -> {method}")
        return 0

    log_step(f"Migrating {len(big_keys)} big keys...")
    success, errors = 0, 0
    for key, mem_usage in big_keys:
        if migrate_key(source, target, key, mem_usage, args.replace, args.verbose, args.chunk_size):
            success += 1
        else:
            errors += 1
    print()
    log_info(f"Migration complete! Successful: {success}, Failed: {errors}")
    return 1 if errors > 0 else 0


if __name__ == "__main__":
    sys.exit(main())