# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import boto3
import json
import time
from awslabs.postgres_mcp_server import __user_agent__
from botocore.config import Config
from botocore.exceptions import ClientError
from loguru import logger
from typing import Any, Dict, List, Optional, Tuple


DEFAULT_POSTGRES_PORT = 5432


def internal_create_rds_client(region: str):
    """Create an RDS client with custom user agent configuration."""
    return boto3.client('rds', region_name=region, config=Config(user_agent_extra=__user_agent__))


def internal_get_cluster_valid_endpoints(
    cluster_properties: Dict[str, Any], region: str
) -> List[Tuple[str, int]]:
    """Return the list of valid (host, port) connection endpoints for a cluster.

    A caller-supplied db_endpoint/port is considered legitimate only if it
    matches one of the endpoints returned here. This is used to prevent
    connection strings from pointing at arbitrary hosts (e.g. an
    attacker-controlled FQDN / IP) that could capture DB credentials or
    IAM auth tokens.

    The returned list includes, where present:
      - Writer endpoint (cluster.Endpoint)
      - Reader endpoint (cluster.ReaderEndpoint)
      - Custom endpoints (cluster.CustomEndpoints)
      - Each member instance endpoint (from describe_db_instances)

    Host comparison should be case-insensitive (DNS is case-insensitive);
    this function returns hosts as reported by AWS without normalization.

    Args:
        cluster_properties: Cluster properties previously fetched via
            internal_get_cluster_properties.
        region: AWS region (used to fetch member instance endpoints).

    Returns:
        Non-empty list of (host, port) tuples.

    Raises:
        ValueError: If no valid endpoints could be resolved. A real Aurora
            cluster always has at least a writer endpoint with a port, so
            an empty result indicates malformed cluster properties or a
            transient API state; we treat it as an error rather than
            silently accepting any caller-supplied endpoint.
    """
    endpoints: List[Tuple[str, int]] = []

    cluster_port_raw = cluster_properties.get('Port')
    try:
        cluster_port = int(cluster_port_raw) if cluster_port_raw is not None else 0
    except (TypeError, ValueError):
        cluster_port = DEFAULT_POSTGRES_PORT

    writer_endpoint = cluster_properties.get('Endpoint')
    if writer_endpoint and cluster_port:
        endpoints.append((writer_endpoint, cluster_port))

    reader_endpoint = cluster_properties.get('ReaderEndpoint')
    if reader_endpoint and cluster_port:
        endpoints.append((reader_endpoint, cluster_port))

    for custom_endpoint in cluster_properties.get('CustomEndpoints', []) or []:
        if custom_endpoint and cluster_port:
            endpoints.append((custom_endpoint, cluster_port))

    members = cluster_properties.get('DBClusterMembers', []) or []
    instance_ids = [
        m.get('DBInstanceIdentifier') for m in members if m.get('DBInstanceIdentifier')
    ]
    if instance_ids:
        rds_client = internal_create_rds_client(region=region)
        for instance_id in instance_ids:
            try:
                response = rds_client.describe_db_instances(DBInstanceIdentifier=instance_id)
                for instance in response.get('DBInstances', []):
                    instance_endpoint = instance.get('Endpoint', {})
                    host = instance_endpoint.get('Address')
                    port_raw = instance_endpoint.get('Port')
                    try:
                        port = int(port_raw) if port_raw is not None else 0
                    except (TypeError, ValueError):
                        port = 0
                    if host and port:
                        endpoints.append((host, port))
            except ClientError as e:
                # Don't fail the whole validation if a single instance lookup
                # fails; log and continue. The caller still has to match one of
                # the endpoints we successfully resolved.
                logger.warning(
                    f"Failed to fetch endpoint for member instance '{instance_id}': "
                    f'{e.response["Error"]["Code"]} - {e.response["Error"]["Message"]}'
                )

    if not endpoints:
        # A real Aurora cluster always advertises at least a writer endpoint
        # with a valid port. An empty result here means either the cluster
        # properties dict was malformed (missing/invalid Port, missing
        # Endpoint) or every describe_db_instances call failed. Treat as an
        # error so the caller doesn't build a connection string from caller
        # input.
        raise ValueError(
            f'Cluster has no valid connection endpoints. '
            f'Endpoint={cluster_properties.get("Endpoint")!r}, '
            f'ReaderEndpoint={cluster_properties.get("ReaderEndpoint")!r}, '
            f'Port={cluster_properties.get("Port")!r}'
        )

    return endpoints


def internal_get_instance_properties(target_endpoint: str, region: str) -> Dict[str, Any]:
    """Retrieve RDS instance properties from AWS."""
    rds_client = internal_create_rds_client(region=region)
    paginator = rds_client.get_paginator('describe_db_instances')

    # Iterate through all instances
    try:
        for page in paginator.paginate():
            for instance in page['DBInstances']:
                endpoint = instance.get('Endpoint', {}).get('Address')
                if endpoint == target_endpoint:
                    return instance
    except ClientError as e:
        error_code = e.response['Error']['Code']
        logger.exception(
            f'AWS error fetching all instances in region:{region} '
            f'{error_code} - {e.response["Error"]["Message"]}'
        )
        raise
    except Exception as e:
        logger.exception(
            f'Error fetching all instances in region:{region}.  Error: {type(e).__name__}: {e}'
        )
        raise

    not_found_error = f"Instance with endpoint '{target_endpoint}' in region:{region} not found when getting its properties"
    logger.error(not_found_error)
    raise ValueError(not_found_error)


def internal_get_cluster_properties(cluster_identifier: str, region: str) -> Dict[str, Any]:
    """Retrieve RDS cluster properties from AWS.

    Args:
        cluster_identifier: RDS cluster identifier
        region: AWS region (e.g., 'us-east-1')

    Returns:
        Dict[str, Any]: Cluster properties from AWS RDS API

    Raises:
        ValueError: If cluster_identifier or region is empty
        ClientError: If AWS API call fails (cluster not found, access denied, etc.)
        NoCredentialsError: If AWS credentials not configured

    Example:
        >>> props = internal_get_cluster_properties('my-cluster', 'us-east-1')
        >>> print(props['Status'])
    """
    # Input validation
    if not cluster_identifier or not region:
        raise ValueError('cluster_identifier and region are required')

    logger.info(f"Fetching properties for cluster '{cluster_identifier}' in '{region}' ")

    try:
        rds_client = internal_create_rds_client(region)
        response = rds_client.describe_db_clusters(DBClusterIdentifier=cluster_identifier)

        # Safely extract cluster properties
        clusters = response.get('DBClusters', [])
        if not clusters:
            raise ValueError(f"Cluster '{cluster_identifier}' not found in region '{region}'")

        cluster_properties = clusters[0]

        # Log summary only
        logger.info(
            f"Retrieved cluster '{cluster_identifier}': "
            f'Status={cluster_properties.get("Status")}, '
            f'Engine={cluster_properties.get("Engine")}'
        )

        # Full properties at debug level
        logger.debug(
            f'Cluster properties: {json.dumps(cluster_properties, indent=2, default=str)}'
        )

        return cluster_properties

    except ClientError as e:
        error_code = e.response['Error']['Code']
        logger.exception(
            f"AWS error fetching cluster '{cluster_identifier}': "
            f'{error_code} - {e.response["Error"]["Message"]}'
        )
        raise
    except Exception as e:
        logger.exception(f'Error fetching cluster properties: {type(e).__name__}: {e}')
        raise


def internal_create_express_cluster(cluster_identifier: str, region: str) -> Dict[str, Any]:
    """Create an Aurora PostgreSQL Express cluster.

    Args:
        cluster_identifier: Unique name for the cluster
        region: AWS region where the cluster is located

    Returns:
        Dict[str, Any]: Cluster properties

    Raises:
        ValueError: If cluster_identifier is invalid
        ClientError: If AWS API call fails
    """
    rds_client = internal_create_rds_client(region)

    # Add default tags
    tags = []
    tags.append({'Key': 'CreatedBy', 'Value': 'MCP'})

    logger.info(f'Creating Aurora Postgres Express cluster: {cluster_identifier}')

    try:
        cluster_create_start_time = time.time()
        rds_client.create_db_cluster(
            DBClusterIdentifier=cluster_identifier,
            Engine='aurora-postgresql',
            WithExpressConfiguration=True,
            Tags=tags,
        )

        logger.info('Waiting for cluster to become available...')
        waiter = rds_client.get_waiter('db_cluster_available')
        waiter.wait(
            DBClusterIdentifier=cluster_identifier, WaiterConfig={'Delay': 5, 'MaxAttempts': 120}
        )

        result = rds_client.describe_db_clusters(DBClusterIdentifier=cluster_identifier)[
            'DBClusters'
        ][0]

        cluster_create_stop_time = time.time()
        elapsed_time = cluster_create_stop_time - cluster_create_start_time
        logger.info(
            f'Express Cluster {cluster_identifier} created successfully and took {elapsed_time:.2f} seconds'
        )
        return result

    except ClientError as e:
        logger.exception(
            f"AWS error creating express cluster '{cluster_identifier}': "
            f'{e.response["Error"]["Code"]} - {e.response["Error"]["Message"]}'
        )
        raise
    except Exception as e:
        logger.exception(
            f"Error creating express cluster '{cluster_identifier}': {type(e).__name__}: {e}"
        )
        raise


def internal_create_serverless_cluster(
    region: str,
    cluster_identifier: str,
    engine_version: str,
    database_name: str,
    master_username: str = 'postgres',
    min_capacity: float = 0.5,
    max_capacity: float = 4,
    enable_cloudwatch_logs: bool = True,
) -> Dict[str, Any]:
    """Create an Aurora PostgreSQL cluster with a single writer instance.

    Credentials are automatically managed by AWS Secrets Manager.

    Args:
        region: region of the cluster
        cluster_identifier: Name of the Aurora cluster
        engine_version: PostgreSQL engine version (e.g., '15.3', '14.7')
        database_name: Name of the default database
        master_username: Master username for the database
        min_capacity: minimum ACU capacity
        max_capacity: maximum ACU capacity
        enable_cloudwatch_logs: Enable CloudWatch logs export

    Returns:
        Dictionary containing cluster information and secret ARN
    """
    if not region:
        raise ValueError('region is required')
    if not cluster_identifier:
        raise ValueError('cluster_identifier is required')
    if not engine_version:
        raise ValueError('engine_version is required')
    if not database_name:
        raise ValueError('database_name is required')

    rds_client = internal_create_rds_client(region=region)

    # Add default tags
    tags = []
    tags.append({'Key': 'CreatedBy', 'Value': 'MCP'})

    # Prepare CloudWatch logs
    enable_cloudwatch_logs_exports = []
    if enable_cloudwatch_logs:
        enable_cloudwatch_logs_exports = ['postgresql']

    try:
        # Create the Aurora cluster
        logger.info(
            f'Creating Aurora PostgreSQL cluster:{cluster_identifier} '
            f'region:{region} engine_version:{engine_version} database_name:{database_name} '
            f'master_username:{master_username}'
        )

        cluster_params = {
            'DBClusterIdentifier': cluster_identifier,
            'Engine': 'aurora-postgresql',
            'EngineVersion': engine_version,
            'MasterUsername': master_username,
            'DatabaseName': database_name,
            'ManageMasterUserPassword': True,  # Enable Secrets Manager integration
            'Tags': tags,
            'DeletionProtection': False,  # Set to True for production
            'CopyTagsToSnapshot': True,
            'EnableHttpEndpoint': True,  # Enable for Data API if needed
            'EnableCloudwatchLogsExports': enable_cloudwatch_logs_exports,
        }

        cluster_params['ServerlessV2ScalingConfiguration'] = {
            'MinCapacity': min_capacity,
            'MaxCapacity': max_capacity,
        }

        # Create the cluster
        cluster_create_start_time = time.time()
        cluster_response = rds_client.create_db_cluster(**cluster_params)

        cluster_info = cluster_response['DBCluster']
        logger.info(
            f'Cluster {cluster_identifier} creation call started successfully. Status: {cluster_info["Status"]}'
        )

        # Wait for cluster to be available
        logger.info('Waiting for cluster to become available...')
        waiter = rds_client.get_waiter('db_cluster_available')
        waiter.wait(
            DBClusterIdentifier=cluster_identifier, WaiterConfig={'Delay': 5, 'MaxAttempts': 120}
        )

        logger.info(f'Cluster {cluster_identifier} is now available')
        cluster_create_stop_time = time.time()
        elapsed_time = cluster_create_stop_time - cluster_create_start_time
        logger.info(f'Cluster creation {cluster_identifier} took {elapsed_time:.2f} seconds')

        # Create the writer instance
        instance_identifier = f'{cluster_identifier}-instance-1'
        logger.info(f'Creating writer instance: {instance_identifier}')

        instance_params = {
            'DBInstanceIdentifier': instance_identifier,
            'DBInstanceClass': 'db.serverless',
            'Engine': 'aurora-postgresql',
            'DBClusterIdentifier': cluster_identifier,
            'PubliclyAccessible': False,  # Set to True if needed
            'Tags': tags,
            'CopyTagsToSnapshot': True,
        }

        instance_create_start_time = time.time()
        rds_client.create_db_instance(**instance_params)

        logger.info(f'Writer instance {instance_identifier} created successfully')

        # Wait for instance to be available
        logger.info(f'Waiting for instance {instance_identifier} to become available...')
        instance_waiter = rds_client.get_waiter('db_instance_available')
        instance_waiter.wait(
            DBInstanceIdentifier=instance_identifier,
            WaiterConfig={
                'Delay': 1,  # check every  seconds
                'MaxAttempts': 1800,  # Try up to 1800 time = 30 mins
            },
        )

        logger.info(f'Instance {instance_identifier} is now available')
        instance_create_stop_time = time.time()
        elapsed_time = instance_create_stop_time - instance_create_start_time
        logger.info(f'Instance creation {instance_identifier} took {elapsed_time:.2f} seconds')

        # Get the final cluster details including the secret ARN
        final_cluster = rds_client.describe_db_clusters(DBClusterIdentifier=cluster_identifier)[
            'DBClusters'
        ][0]

        return final_cluster

    except ClientError as e:
        logger.exception(
            f"AWS error creating serverless cluster '{cluster_identifier}': "
            f'{e.response["Error"]["Code"]} - {e.response["Error"]["Message"]}'
        )
        raise
    except Exception as e:
        logger.exception(
            f"Error creating serverless cluster '{cluster_identifier}': {type(e).__name__}: {e}"
        )
        raise


def setup_aurora_iam_policy_for_current_user(
    db_user: str, cluster_resource_id: str, cluster_region: str
) -> Optional[str]:
    """Create or update IAM policy for Aurora access.

    Maintains one policy per user, adding new clusters as they're created.

    ⚠️  If running as assumed role, this will attempt to attach the policy to
        the BASE ROLE (not the session). This requires iam:AttachRolePolicy permission.

    Args:
        db_user: PostgreSQL username (must have rds_iam role granted in database)
        cluster_resource_id: The DBI resource ID (e.g., 'cluster-ABCD123XYZ')
        cluster_region: AWS region where the Aurora cluster is located

    Returns:
        Policy ARN if successful, None otherwise

    Raises:
        ValueError: If running as federated user, root, or invalid identity
        boto3 exceptions: For AWS API errors (except AccessDenied on attach)
    """
    # Validate inputs
    if not db_user or not isinstance(db_user, str):
        raise ValueError('db_user must be a non-empty string')
    if not cluster_resource_id or not isinstance(cluster_resource_id, str):
        raise ValueError('cluster_resource_id must be a non-empty string')
    if not cluster_region or not isinstance(cluster_region, str):
        raise ValueError('cluster_region must be a non-empty string')

    # Initialize clients
    sts = boto3.client('sts', config=Config(user_agent_extra=__user_agent__))
    iam = boto3.client('iam', config=Config(user_agent_extra=__user_agent__))

    # 1. Get current IAM identity
    try:
        identity = sts.get_caller_identity()
        account_id = identity['Account']
        arn = identity['Arn']
        user_id = identity['UserId']

        logger.info('Current Identity:')
        logger.info(f'  ARN: {arn}')
        logger.info(f'  Account: {account_id}')
        logger.info(f'  UserID: {user_id}')

    except Exception as e:
        logger.exception(f'❌ Error getting caller identity: {e}')
        raise

    # ============================================================================
    # 🔵 MODIFIED: Extract base role from assumed role session
    # ============================================================================
    # 2. Extract username/role from ARN and determine identity type
    current_user = None
    current_role = None
    identity_type = None

    if ':user/' in arn:
        # Standard IAM user: arn:aws:iam::123456789012:user/username
        current_user = arn.split(':user/')[-1].split('/')[-1]
        identity_type = 'user'
        logger.info('  Type: IAM User')
        logger.info(f'  Username: {current_user}')

    elif ':assumed-role/' in arn:
        # 🔵 MODIFIED: Extract BASE ROLE name from assumed role session
        # Assumed role ARN: arn:aws:sts::123456789012:assumed-role/RoleName/session-name
        # We want to extract "RoleName" (the base role)
        parts = arn.split(':assumed-role/')[-1].split('/')
        current_role = parts[0]  # This is the BASE ROLE name
        session_name = parts[1] if len(parts) > 1 else 'unknown'

        identity_type = 'role'
        logger.info('  Type: Assumed Role Session')
        logger.info(f'  Base Role: {current_role}')
        logger.info(f'  Session Name: {session_name}')
        logger.info(f'  → Will attach policy to base role: {current_role}')
        logger.warning(
            f"⚠️  Policy will be attached to role '{current_role}'\n"
            f'   This will grant Aurora access to ALL users/services that assume this role.'
        )

    elif ':federated-user/' in arn:
        logger.error('  Type: Federated User')
        raise ValueError(
            'Cannot attach policies to federated users.\n'
            'Please use the parent IAM user or role instead.'
        )

    elif ':root' in arn:
        logger.error('  Type: Root User')
        raise ValueError(
            'Cannot (and should not) attach policies to root user.\n'
            'Please use an IAM user instead.'
        )

    else:
        raise ValueError(f'Unexpected ARN format: {arn}')

    # 3. Prepare new resource ARN
    policy_name = f'AuroraIAMAuth-{db_user}'
    policy_arn = f'arn:aws:iam::{account_id}:policy/{policy_name}'

    new_resource_arn = (
        f'arn:aws:rds-db:{cluster_region}:{account_id}:dbuser:{cluster_resource_id}/{db_user}'
    )

    logger.info('\nPolicy Configuration:')
    logger.info(f'  Policy Name: {policy_name}')
    logger.info(f'  New Resource: {new_resource_arn}')
    logger.info(f'  Cluster Region: {cluster_region}')
    logger.info(f'  Cluster Resource ID: {cluster_resource_id}')

    # 4. Create or update policy

    try:
        # Try to get existing policy
        existing_policy = iam.get_policy(PolicyArn=policy_arn)
        logger.info(f'\n✓ Policy already exists: {policy_name}')

        # Get current policy document
        policy_version = iam.get_policy_version(
            PolicyArn=policy_arn, VersionId=existing_policy['Policy']['DefaultVersionId']
        )

        current_doc = policy_version['PolicyVersion']['Document']
        current_resources = current_doc['Statement'][0]['Resource']

        # Normalize to list (could be string or list)
        if isinstance(current_resources, str):
            current_resources = [current_resources]

        logger.info(f'  Current resources in policy: {len(current_resources)}')
        for idx, res in enumerate(current_resources, 1):
            logger.info(f'    {idx}. {res}')

        # Check if new resource already exists
        if new_resource_arn in current_resources:
            logger.info('\n✓ Cluster already included in policy - no update needed')
        else:
            # Add new resource to the list
            current_resources.append(new_resource_arn)
            logger.info('\n→ Adding new cluster to policy...')

            # Create updated policy document
            updated_doc = {
                'Version': '2012-10-17',
                'Statement': [
                    {'Effect': 'Allow', 'Action': 'rds-db:connect', 'Resource': current_resources}
                ],
            }

            # Handle AWS policy version limits (max 5 versions per policy)
            versions = iam.list_policy_versions(PolicyArn=policy_arn)['Versions']
            logger.info(f'  Current policy versions: {len(versions)}/5')

            if len(versions) >= 5:
                # Find oldest non-default version to delete
                non_default_versions = [v for v in versions if not v['IsDefaultVersion']]
                if non_default_versions:
                    oldest_version = sorted(non_default_versions, key=lambda v: v['CreateDate'])[0]
                    logger.info(
                        f'  Deleting oldest version: {oldest_version["VersionId"]} (created {oldest_version["CreateDate"]})'
                    )
                    iam.delete_policy_version(
                        PolicyArn=policy_arn, VersionId=oldest_version['VersionId']
                    )

            # Create new policy version
            new_version = iam.create_policy_version(
                PolicyArn=policy_arn,
                PolicyDocument=json.dumps(updated_doc, indent=2),
                SetAsDefault=True,
            )

            logger.info('✓ Successfully updated policy')
            logger.info(f'  New version: {new_version["PolicyVersion"]["VersionId"]}')
            logger.info(f'  Total resources now: {len(current_resources)}')

    except iam.exceptions.NoSuchEntityException:
        # Policy doesn't exist - create new one
        logger.info("\nPolicy doesn't exist, creating new policy...")

        policy_document = {
            'Version': '2012-10-17',
            'Statement': [
                {'Effect': 'Allow', 'Action': 'rds-db:connect', 'Resource': [new_resource_arn]}
            ],
        }

        try:
            policy_response = iam.create_policy(
                PolicyName=policy_name,
                PolicyDocument=json.dumps(policy_document, indent=2),
                Description=f'IAM authentication for Aurora PostgreSQL user {db_user} across all clusters',
            )
            policy_arn = policy_response['Policy']['Arn']
            logger.info(f'✓ Successfully created new policy: {policy_name}')
            logger.info(f'  Policy ARN: {policy_arn}')

        except iam.exceptions.EntityAlreadyExistsException:
            logger.info('✓ Policy was just created by another process')

        except Exception as e:
            logger.exception(f'\n❌ Error creating policy: {e}')
            raise

    except Exception as e:
        logger.exception(f'\n❌ Error checking/updating policy: {e}')
        raise

    # ============================================================================
    # 🔵 MODIFIED: Attach to base role with better error handling
    # ============================================================================
    # 5. Attach policy to current user OR base role
    try:
        if identity_type == 'user':
            # IAM User - attach directly
            attached_policies = iam.list_attached_user_policies(UserName=current_user)
            already_attached = any(
                p['PolicyArn'] == policy_arn for p in attached_policies['AttachedPolicies']
            )

            if already_attached:
                logger.info(f'\n✓ Policy already attached to user: {current_user}')
            else:
                iam.attach_user_policy(UserName=current_user, PolicyArn=policy_arn)
                logger.info(f'\n✓ Successfully attached policy to user: {current_user}')

            # Display summary
            logger.info(f'\nAttached policies for user {current_user}:')
            attached_policies = iam.list_attached_user_policies(UserName=current_user)
            for policy in attached_policies['AttachedPolicies']:
                marker = '  → ' if policy['PolicyArn'] == policy_arn else '    '
                logger.info(f'{marker}{policy["PolicyName"]}')

        elif identity_type == 'role':
            # 🔵 MODIFIED: Attach to BASE ROLE (not session)
            logger.info(f'\n→ Attempting to attach policy to base role: {current_role}')

            try:
                # Check if already attached to the base role
                attached_policies = iam.list_attached_role_policies(RoleName=current_role)
                already_attached = any(
                    p['PolicyArn'] == policy_arn for p in attached_policies['AttachedPolicies']
                )

                if already_attached:
                    logger.info(f'\n✓ Policy already attached to role: {current_role}')
                else:
                    # Attach to the BASE ROLE
                    iam.attach_role_policy(RoleName=current_role, PolicyArn=policy_arn)
                    logger.info(f'\n✓ Successfully attached policy to role: {current_role}')
                    logger.warning(
                        f"⚠️  All users/services assuming role '{current_role}' now have Aurora access"
                    )

                # Display summary
                logger.info(f'\nAttached policies for role {current_role}:')
                attached_policies = iam.list_attached_role_policies(RoleName=current_role)
                for policy in attached_policies['AttachedPolicies']:
                    marker = '  → ' if policy['PolicyArn'] == policy_arn else '    '
                    logger.info(f'{marker}{policy["PolicyName"]}')

            except iam.exceptions.AccessDeniedException:
                # 🔵 MODIFIED: Graceful handling of permission denied
                logger.exception(
                    f"\n❌ Access Denied: Cannot attach policy to role '{current_role}'"
                )
                logger.error("   Your session does not have 'iam:AttachRolePolicy' permission")
                logger.info(f'\n✓ Policy created successfully: {policy_arn}')
                logger.info('   But could not be attached automatically.')
                logger.info('\n📋 MANUAL STEPS REQUIRED:')
                logger.info('\n Option 1: Have an administrator attach the policy to the role')
                logger.info('   aws iam attach-role-policy \\')
                logger.info(f'     --role-name {current_role} \\')
                logger.info(f'     --policy-arn {policy_arn}')
                logger.info('\n Option 2: Attach to your individual IAM user (if you have one)')
                logger.info('   aws iam attach-user-policy \\')
                logger.info('     --user-name YOUR_IAM_USERNAME \\')
                logger.info(f'     --policy-arn {policy_arn}')
                logger.info('\n Option 3: Grant the role permission to attach policies')
                logger.info(
                    f"   (Admin needs to add iam:AttachRolePolicy to role '{current_role}')"
                )

                # Return policy ARN even though not attached
                return policy_arn

            except iam.exceptions.NoSuchEntityException:
                logger.exception(f"\n❌ Role '{current_role}' not found")
                raise

        return policy_arn

    except iam.exceptions.NoSuchEntityException:
        entity_name = current_user if identity_type == 'user' else current_role
        entity_type = 'User' if identity_type == 'user' else 'Role'
        logger.exception(f"\n❌ Error: {entity_type} '{entity_name}' not found")
        raise

    except iam.exceptions.LimitExceededException:
        entity_name = current_user if identity_type == 'user' else current_role
        entity_type = 'user' if identity_type == 'user' else 'role'
        logger.exception(
            f"\n❌ Error: Managed policy limit exceeded for {entity_type} '{entity_name}'"
        )
        raise

    except Exception as e:
        logger.exception(f'\n❌ Error attaching policy: {e}')
        raise


def internal_delete_cluster(region: str, cluster_id: str) -> None:
    """Delete an existing Amazon RDS (Aurora) cluster.

    Args:
        region: AWS region where the cluster is located
        cluster_id: The DB cluster identifier

    Raises:
        ClientError: If deletion fails or the cluster is not found
        Exception: If cluster was not created by MCP tool (safety check)
    """
    rds = internal_create_rds_client(region)

    # Check cluster exists and verify it was created by MCP
    try:
        resp = rds.describe_db_clusters(DBClusterIdentifier=cluster_id)
        cluster = resp['DBClusters'][0]
        status = cluster['Status']
        arn = cluster['DBClusterArn']

        tag_resp = rds.list_tags_for_resource(ResourceName=arn)
        tags = tag_resp.get('TagList', [])
        created_by_mcp = False
        if tags:
            for tag in tags:
                if tag['Key'] == 'CreatedBy' and tag['Value'] == 'MCP':
                    created_by_mcp = True
        logger.info(f"Found cluster '{cluster_id}' (status={status})")

        if not created_by_mcp:
            logger.error('Can only delete clusters created by MCP tool')
            raise Exception('Can only delete clusters created by MCP tool')

    except ClientError as e:
        if e.response['Error']['Code'] == 'DBClusterNotFoundFault':
            logger.info(f"Cluster '{cluster_id}' does not exist")
            return
        logger.exception(f"Error checking cluster '{cluster_id}' for deletion")
        raise

    # Delete all DB instances in the cluster
    members = cluster.get('DBClusterMembers', [])
    instance_ids = [m['DBInstanceIdentifier'] for m in members]

    if instance_ids:
        logger.info(f'Deleting {len(instance_ids)} instance(s): {", ".join(instance_ids)}')

    for inst_id in instance_ids:
        try:
            rds.delete_db_instance(
                DBInstanceIdentifier=inst_id,
                SkipFinalSnapshot=True,
                DeleteAutomatedBackups=True,
            )
            logger.info(f"Deletion of instance '{inst_id}' initiated")
        except ClientError as e:
            code = e.response['Error']['Code']
            if code in ('DBInstanceNotFound', 'DBInstanceNotFoundFault'):
                logger.info(f"Instance '{inst_id}' already deleted")
            else:
                logger.exception(f"Error deleting instance '{inst_id}': {e}")
                raise

    # Wait for all instances to be fully deleted (max ~20 minutes)
    max_instance_attempts = 240
    if instance_ids:
        logger.info('Waiting for instances to be deleted...')
    remaining = set(instance_ids)
    instance_attempts = 0

    while remaining:
        instance_attempts += 1
        if instance_attempts > max_instance_attempts:
            raise TimeoutError(
                f'Timed out waiting for instance(s) to be deleted: {", ".join(remaining)}'
            )
        done = []
        for inst_id in list(remaining):
            try:
                rds.describe_db_instances(DBInstanceIdentifier=inst_id)
                # still exists
            except ClientError as e:
                if e.response['Error']['Code'] in (
                    'DBInstanceNotFound',
                    'DBInstanceNotFoundFault',
                ):
                    logger.info(f"Instance '{inst_id}' deleted")
                    done.append(inst_id)
                else:
                    raise
        for d in done:
            remaining.discard(d)
        if remaining:
            time.sleep(5)

    # Delete cluster
    try:
        rds.delete_db_cluster(
            DBClusterIdentifier=cluster_id,
            SkipFinalSnapshot=True,
        )
        logger.info(f"Deletion of cluster '{cluster_id}' initiated")
    except ClientError as e:
        logger.exception(f"Error deleting cluster '{cluster_id}': {e}")
        raise

    # Poll for cluster deletion (max ~20 minutes)
    max_cluster_attempts = 240
    logger.info('Waiting for cluster to be deleted...')
    cluster_attempts = 0
    while True:
        cluster_attempts += 1
        if cluster_attempts > max_cluster_attempts:
            raise TimeoutError(f"Timed out waiting for cluster '{cluster_id}' to be deleted")
        try:
            rds.describe_db_clusters(DBClusterIdentifier=cluster_id)
            time.sleep(5)
        except ClientError as e:
            if e.response['Error']['Code'] == 'DBClusterNotFoundFault':
                logger.info(f"Cluster '{cluster_id}' deleted successfully")
                break
            else:
                raise
