#!/usr/bin/env python3
"""
h2oGPTe Embedding Model Migration CLI

A tool for migrating collections from deprecated embedding models to compliant models.
"""

import argparse
import os
import sys

from dotenv import load_dotenv

from h2ogpte_migration.config import MigrationConfig
from h2ogpte_migration.logger import MigrationLogger
from h2ogpte_migration.migrate import CollectionMigrator


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description='Migrate h2oGPTe collections to new embedding models. Preserves permissions, '
                    'lifecycle settings, scheduled connectors, and document metadata.',
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog='''
Examples:
  # Admin: Migrate all users (parallel, jobs run in background)
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --all-users \\
    --use-model-mappings

  # Admin: Verify previously submitted jobs
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --verify

  # Admin: Verify + migrate chat sessions for completed jobs
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --verify \\
    --migrate-chats

  # Admin: Sequential migration with chat migration
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --all-users \\
    --use-model-mappings \\
    --wait-for-completion \\
    --migrate-chats

  # Admin: Migrate specific collections (auto-detects owners)
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --collections col-123,col-456 \\
    --use-model-mappings

  # Admin: Dry-run for specific user
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --admin-key sk-xxx \\
    --users john.doe \\
    --use-model-mappings \\
    --dry-run

  # Self-service: Migrate all your collections
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --user-key sk-xxx \\
    --use-model-mappings

  # Self-service: Migrate specific collections you own
  h2ogpte-migrate \\
    --url https://h2ogpte.dev \\
    --user-key sk-xxx \\
    --collections col-abc,col-def \\
    --target-model "mixedbread-ai/mxbai-embed-large-v1"
        '''
    )
    
    # Required arguments
    parser.add_argument('--url', required=True,
                       help='h2oGPTe instance URL')
    
    # Authentication (mutually exclusive)
    auth_group = parser.add_mutually_exclusive_group(required=True)
    auth_group.add_argument('--admin-key', nargs='?', const='', default=None,
                           help='Admin API key (for admin mode). '
                                'Pass without a value to read from H2OGPTE_ADMIN_KEY env var')
    auth_group.add_argument('--user-key', nargs='?', const='', default=None,
                           help='User API key (for self-service mode). '
                                'Pass without a value to read from H2OGPTE_USER_KEY env var')
    
    # Connection options
    parser.add_argument('--cert', dest='cert_path',
                       help='Path to CA certificate file for SSL verification')
    parser.add_argument('--api-key-expiry', default='30 days',
                       help='Expiry duration for temporary API keys (e.g., "30 days", "7 days"). Default: 30 days')
    
    # Migration mode
    parser.add_argument('--use-model-mappings', action='store_true',
                       help='Use predefined model mappings (ignores --source-model/--target-model)')
    
    # Manual mode arguments (used when --use-model-mappings is not set)
    parser.add_argument('--target-model',
                       help='Target embedding model (required without --use-model-mappings)')
    parser.add_argument('--source-model',
                       help='Source embedding model to migrate from (optional, migrates all if not specified)')
    
    # Scope selection (mutually exclusive, but not required if --verify is used alone)
    scope = parser.add_mutually_exclusive_group(required=False)
    scope.add_argument('--users',
                      help='Comma-separated list of usernames to migrate (e.g., "john.doe,jane.smith")')
    scope.add_argument('--all-users', action='store_true',
                      help='Migrate collections for all users')
    scope.add_argument('--collections',
                      help='Comma-separated list of collection IDs to migrate (auto-detects owners)')
    scope.add_argument('--collections-file',
                      help='Path to a file containing collection IDs (one per line, # comments ignored). '
                           'Works like --collections but reads from a file')
    
    # Migration options
    parser.add_argument('--copy-document', action='store_true', default=False,
                       help='Copy documents instead of referencing (default: False, references preserve chat links after old collection deletion)')
    parser.add_argument('--no-copy-document', dest='copy_document', action='store_false',
                       help='Reference documents instead of copying')
    parser.add_argument('--skip-reparse', action='store_true', default=False,
                       help='Re-embed existing chunks without re-parsing documents. Significantly faster for '
                            'embedding model migrations where only the embedding model changes. '
                            'Requires copy_document=False (default). Cannot be used with --copy-document.')
    parser.add_argument('--migrate-chats', action='store_true',
                       help='Migrate chat sessions from old to new collection. '
                            'In wait mode: migrates immediately after import. '
                            'In verify mode: migrates chats for completed imports.')
    parser.add_argument('--ocr-model', default='auto',
                       help='OCR model to use during document re-parsing (default: auto). '
                            'Use this to override the source collection\'s OCR model, '
                            'e.g., when migrating from a CN model to a non-CN model. '
                            'Examples: auto, off, tesseract')

    # Manual move operations
    parser.add_argument('--move-connectors', action='store_true',
                       help='Move scheduled connectors from --from collection to --to collection')
    parser.add_argument('--move-chats', action='store_true',
                       help='Move chat sessions from --from collection to --to collection')
    parser.add_argument('--from', dest='move_from',
                       help='Source collection ID for --move-connectors/--move-chats')
    parser.add_argument('--to', dest='move_to',
                       help='Target collection ID for --move-connectors/--move-chats')

    # Execution options
    parser.add_argument('--dry-run', action='store_true',
                       help='Preview what would be migrated without executing')

    # Migration tracking options
    parser.add_argument('--db-path', default='migration_tracking.db',
                       help='Path to SQLite database for tracking migrations (default: migration_tracking.db)')
    parser.add_argument('--force-remigrate', action='store_true',
                       help='Re-migrate all collections (ignores DB state, creates new collections)')
    parser.add_argument('--retry-failed', action='store_true',
                       help='Retry previously failed imports (creates new collections for failed ones)')
    parser.add_argument('--wait-for-completion', action='store_true',
                       help='Wait for each import job to complete before moving to the next collection')
    parser.add_argument('--max-concurrent-jobs', type=int, default=1, metavar='N',
                       help='Number of collections to process concurrently with --wait-for-completion '
                            '(default: 1). Each worker runs the full migration cycle (create, import, '
                            'validate, connectors, chats) independently. Has no effect without '
                            '--wait-for-completion.')
    parser.add_argument('--verify', action='store_true',
                       help='Check status of submitted jobs, validate migration (doc counts, '
                            'embedding model, lifecycle settings, document statuses)')
    parser.add_argument('--verify-query',
                       help='RAG verification query — creates a test chat on each verified collection '
                            'and asserts the response has document references. Requires --verify or --wait-for-completion.')

    return parser.parse_args()


def main():
    """Main entry point."""
    try:
        # Load .env file from current directory (if present)
        load_dotenv()

        # Parse arguments
        args = parse_args()

        # Resolve API keys: flag value takes precedence, empty flag reads from env var
        admin_key = args.admin_key
        user_key = args.user_key

        if admin_key == '':
            admin_key = os.environ.get('H2OGPTE_ADMIN_KEY')
            if not admin_key:
                print("Error: --admin-key passed without a value and H2OGPTE_ADMIN_KEY env var is not set",
                      file=sys.stderr)
                sys.exit(2)
        if user_key == '':
            user_key = os.environ.get('H2OGPTE_USER_KEY')
            if not user_key:
                print("Error: --user-key passed without a value and H2OGPTE_USER_KEY env var is not set",
                      file=sys.stderr)
                sys.exit(2)

        # Parse users into list
        users_list = None
        if args.users:
            users_list = [u.strip() for u in args.users.split(',')]

        # Parse collections into list (from --collections or --collections-file)
        collections_list = None
        if args.collections:
            collections_list = [c.strip() for c in args.collections.split(',') if c.strip()]
        elif args.collections_file:
            try:
                with open(args.collections_file, 'r') as f:
                    collections_list = [
                        line.strip() for line in f
                        if line.strip() and not line.strip().startswith('#')
                    ]
            except FileNotFoundError:
                print(f"Error: Collections file not found: {args.collections_file}", file=sys.stderr)
                sys.exit(2)
            if not collections_list:
                print(f"Error: No collection IDs found in {args.collections_file}", file=sys.stderr)
                sys.exit(2)

        # Create configuration
        config = MigrationConfig(
            url=args.url,
            admin_api_key=admin_key,
            user_api_key=user_key,
            use_model_mappings=args.use_model_mappings,
            target_model=args.target_model,
            cert_path=args.cert_path,
            api_key_expiry=args.api_key_expiry,
            source_model=args.source_model,
            users=users_list,
            collection_ids=collections_list,
            all_users=args.all_users,
            dry_run=args.dry_run,
            copy_document=args.copy_document,
            skip_reparse=args.skip_reparse,
            migrate_chats=args.migrate_chats,
            ocr_model=args.ocr_model,
            move_connectors=args.move_connectors,
            move_chats=args.move_chats,
            move_from=args.move_from,
            move_to=args.move_to,
            verify_query=args.verify_query,
            db_path=args.db_path,
            force_remigrate=args.force_remigrate,
            retry_failed=args.retry_failed,
            wait_for_completion=args.wait_for_completion,
            verify=args.verify,
            max_concurrent_jobs=args.max_concurrent_jobs
        )
        
        # Validate configuration
        config.validate()
        
        # Setup logging (all logs now shown by default)
        logger = MigrationLogger()
        
        # Run migration
        migrator = CollectionMigrator(config, logger)
        report = migrator.run()
        
        # Exit with appropriate code
        if getattr(report, 'move_failed', False):
            sys.exit(1)

        verify_counters = getattr(report, 'verify_counters', None)
        if verify_counters:
            # Verify mode: exit 1 if any jobs failed or canceled
            if verify_counters['failed'] + verify_counters['canceled'] > 0:
                sys.exit(1)
        elif report.summary['failed_migrations'] > 0:
            sys.exit(1)

        sys.exit(0)
            
    except ValueError as e:
        print(f"{e}", file=sys.stderr)
        sys.exit(2)
    except KeyboardInterrupt:
        print("\nMigration interrupted by user", file=sys.stderr)
        sys.exit(130)
    except Exception as e:
        print(f"Fatal Error: {e}", file=sys.stderr)
        sys.exit(1)


if __name__ == '__main__':
    main()
