#!/usr/bin/env python3
"""
CLI tool for acquiring and releasing profile locks.
"""

import argparse
import json
import logging
import os
import sys
import time

try:
    from dotenv import load_dotenv
except ImportError:
    load_dotenv = None

from .profile_manager_tool import ProfileManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def add_profile_allocator_parser(subparsers):
    """Adds the parser for the 'profile-allocator' command."""
    parser = subparsers.add_parser(
        'profile-allocator',
        description='Acquire and release profile locks.',
        formatter_class=argparse.RawTextHelpFormatter,
        help='Acquire and release profile locks.'
    )

    common_parser = argparse.ArgumentParser(add_help=False)
    common_parser.add_argument('--env-file', help='Path to a .env file to load environment variables from.')
    common_parser.add_argument('--redis-host', default=None, help='Redis host. Defaults to MASTER_HOST_IP or REDIS_HOST env var, or localhost.')
    common_parser.add_argument('--redis-port', type=int, default=None, help='Redis port. Defaults to REDIS_PORT env var, or 6379.')
    common_parser.add_argument('--redis-password', default=None, help='Redis password. Defaults to REDIS_PASSWORD env var.')
    common_parser.add_argument('--env', default='dev', help="Environment name for Redis key prefix (e.g., 'stg', 'prod'). Defaults to 'dev'.")
    common_parser.add_argument('--legacy', action='store_true', help="Use legacy key prefix ('profile_mgmt_') without environment.")
    common_parser.add_argument('--key-prefix', default=None, help='Explicit key prefix for Redis. Overrides --env, --legacy and any defaults.')
    common_parser.add_argument('--verbose', action='store_true', help='Enable verbose logging')
    
    allocator_subparsers = parser.add_subparsers(dest='allocator_command', help='Command to execute', required=True)
    
    # Lock command
    lock_parser = allocator_subparsers.add_parser('lock', help='Find and lock an available profile', parents=[common_parser])
    lock_parser.add_argument('--owner', required=True, help='Identifier for the process locking the profile')
    lock_parser.add_argument('--profile-prefix', help='Only lock profiles with this name prefix')
    lock_parser.add_argument('--wait', action='store_true', help='Wait indefinitely for a profile to become available, with exponential backoff.')
    
    # Unlock command
    unlock_parser = allocator_subparsers.add_parser('unlock', help='Unlock a profile', parents=[common_parser])
    unlock_parser.add_argument('name', help='Profile name to unlock')
    unlock_parser.add_argument('--owner', help='Identifier of the owner. If provided, unlock will only succeed if owner matches.')
    
    # Cleanup command
    cleanup_parser = allocator_subparsers.add_parser('cleanup-locks', help='Clean up stale locks', parents=[common_parser])
    cleanup_parser.add_argument('--max-age-seconds', type=int, default=3600,
                                help='Maximum lock age in seconds before it is considered stale (default: 3600)')

    return parser

def main_profile_allocator(args):
    """Main dispatcher for 'profile-allocator' command."""
    if load_dotenv:
        env_file = args.env_file
        if not env_file and args.env and '.env' in args.env and os.path.exists(args.env):
            logger.warning(f"Warning: --env should be an environment name (e.g., 'dev'), not a file path. Treating '{args.env}' as --env-file. The environment name will default to 'dev'.")
            env_file = args.env
            args.env = 'dev'

        was_loaded = load_dotenv(env_file)
        if was_loaded:
            logger.info(f"Loaded environment variables from {env_file or '.env file'}")
        elif args.env_file:
            logger.error(f"The specified --env-file was not found: {args.env_file}")
            return 1

    if args.redis_host is None:
        args.redis_host = os.getenv('MASTER_HOST_IP', os.getenv('REDIS_HOST', 'localhost'))
    if args.redis_port is None:
        args.redis_port = int(os.getenv('REDIS_PORT', 6379))
    if args.redis_password is None:
        args.redis_password = os.getenv('REDIS_PASSWORD')

    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)
    
    if args.key_prefix:
        key_prefix = args.key_prefix
    elif args.legacy:
        key_prefix = 'profile_mgmt_'
    else:
        key_prefix = f"{args.env}_profile_mgmt_"
    
    manager = ProfileManager(
        redis_host=args.redis_host,
        redis_port=args.redis_port,
        redis_password=args.redis_password,
        key_prefix=key_prefix
    )
    
    if args.allocator_command == 'lock':
        if not args.wait:
            profile = manager.lock_profile(args.owner, profile_prefix=args.profile_prefix)
            if profile:
                print(json.dumps(profile, indent=2, default=str))
                return 0
            else:
                print("No available profile could be locked.", file=sys.stderr)
                return 1

        # With --wait, loop with backoff
        lock_attempts = 0
        backoff_seconds = [3, 5, 9, 20, 50, 120, 300]
        while True:
            profile = manager.lock_profile(args.owner, profile_prefix=args.profile_prefix)
            if profile:
                print(json.dumps(profile, indent=2, default=str))
                return 0

            sleep_duration = backoff_seconds[min(lock_attempts, len(backoff_seconds) - 1)]
            logger.info(f"No available profile. Retrying in {sleep_duration}s... (attempt {lock_attempts + 1})")

            try:
                time.sleep(sleep_duration)
            except KeyboardInterrupt:
                logger.warning("Wait for lock interrupted by user.")
                # Use print for stderr as well, since logger might be configured differently by callers
                print("\nWait for lock interrupted by user.", file=sys.stderr)
                return 130 # Standard exit code for Ctrl+C

            lock_attempts += 1
    
    elif args.allocator_command == 'unlock':
        success = manager.unlock_profile(args.name, args.owner)
        return 0 if success else 1
    
    elif args.allocator_command == 'cleanup-locks':
        cleaned_count = manager.cleanup_stale_locks(args.max_age_seconds)
        print(f"Cleaned {cleaned_count} stale lock(s).")
        return 0
        
    return 1 # Should not be reached
