# 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.

"""Utility functions for checking AWS storage services encryption and security."""

from typing import Any, Dict, List

import boto3
import botocore.exceptions
from mcp.server.fastmcp import Context

from awslabs.well_architected_security_mcp_server.consts import USER_AGENT_CONFIG


async def check_storage_encryption(
    region: str,
    services: List[str],
    session: boto3.Session,
    ctx: Context,
    include_unencrypted_only: bool = False,
) -> Dict[str, Any]:
    """Check AWS storage resources for encryption and security best practices.

    Args:
        region: AWS region to check
        services: List of storage services to check
        session: boto3 Session for AWS API calls
        ctx: MCP context for error reporting
        include_unencrypted_only: Whether to include only unencrypted resources in the results

    Returns:
        Dictionary with storage encryption and security status
    """
    results = {
        "region": region,
        "services_checked": services,
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "compliance_by_service": {},
        "resource_details": [],
        "recommendations": [],
    }

    # Find all storage resources using Resource Explorer
    storage_resources = await find_storage_resources(region, session, services, ctx)

    # Check each service as requested
    if "s3" in services:
        s3_client = session.client("s3", region_name=region, config=USER_AGENT_CONFIG)
        s3_results = await check_s3_buckets(region, s3_client, ctx, storage_resources)
        await _update_results(results, s3_results, "s3", include_unencrypted_only)

    if "ebs" in services:
        ec2_client = session.client("ec2", region_name=region, config=USER_AGENT_CONFIG)
        ebs_results = await check_ebs_volumes(region, ec2_client, ctx, storage_resources)
        await _update_results(results, ebs_results, "ebs", include_unencrypted_only)

    if "rds" in services:
        rds_client = session.client("rds", region_name=region, config=USER_AGENT_CONFIG)
        rds_results = await check_rds_instances(region, rds_client, ctx, storage_resources)
        await _update_results(results, rds_results, "rds", include_unencrypted_only)

    if "dynamodb" in services:
        dynamodb_client = session.client("dynamodb", region_name=region, config=USER_AGENT_CONFIG)
        dynamodb_results = await check_dynamodb_tables(
            region, dynamodb_client, ctx, storage_resources
        )
        await _update_results(results, dynamodb_results, "dynamodb", include_unencrypted_only)

    if "efs" in services:
        efs_client = session.client("efs", region_name=region, config=USER_AGENT_CONFIG)
        efs_results = await check_efs_filesystems(region, efs_client, ctx, storage_resources)
        await _update_results(results, efs_results, "efs", include_unencrypted_only)

    if "elasticache" in services:
        elasticache_client = session.client(
            "elasticache", region_name=region, config=USER_AGENT_CONFIG
        )
        elasticache_results = await check_elasticache_clusters(
            region, elasticache_client, ctx, storage_resources
        )
        await _update_results(
            results, elasticache_results, "elasticache", include_unencrypted_only
        )

    # Generate overall recommendations based on findings
    results["recommendations"] = await generate_recommendations(results)

    return results


async def _update_results(
    main_results: Dict[str, Any],
    service_results: Dict[str, Any],
    service_name: str,
    include_unencrypted_only: bool,
) -> None:
    """Update the main results dictionary with service-specific results."""
    # Update resource counts
    main_results["resources_checked"] += service_results.get("resources_checked", 0)
    main_results["compliant_resources"] += service_results.get("compliant_resources", 0)
    main_results["non_compliant_resources"] += service_results.get("non_compliant_resources", 0)

    # Add service-specific compliance info
    main_results["compliance_by_service"][service_name] = {
        "resources_checked": service_results.get("resources_checked", 0),
        "compliant_resources": service_results.get("compliant_resources", 0),
        "non_compliant_resources": service_results.get("non_compliant_resources", 0),
    }

    # Add resource details
    for resource in service_results.get("resource_details", []):
        if not include_unencrypted_only or not resource.get("compliant", True):
            main_results["resource_details"].append(resource)


async def generate_recommendations(results: Dict[str, Any]) -> List[str]:
    """Generate recommendations based on the scan results."""
    recommendations = []

    # Check S3 recommendations
    if "s3" in results.get("compliance_by_service", {}):
        s3_results = results["compliance_by_service"]["s3"]
        if s3_results.get("non_compliant_resources", 0) > 0:
            recommendations.append("Enable default encryption for all S3 buckets")
            recommendations.append("Enable block public access settings at the account level")

    # Check EBS recommendations
    if "ebs" in results.get("compliance_by_service", {}):
        ebs_results = results["compliance_by_service"]["ebs"]
        if ebs_results.get("non_compliant_resources", 0) > 0:
            recommendations.append("Enable default EBS encryption at the account level")
            recommendations.append(
                "Create encrypted snapshots of unencrypted volumes and restore to new encrypted volumes"
            )

    # Check RDS recommendations
    if "rds" in results.get("compliance_by_service", {}):
        rds_results = results["compliance_by_service"]["rds"]
        if rds_results.get("non_compliant_resources", 0) > 0:
            recommendations.append("Enable encryption for all RDS instances")
            recommendations.append("Configure SSL/TLS for database connections")
            recommendations.append("Enable default RDS encryption at the account level")

    # Check DynamoDB recommendations
    if "dynamodb" in results.get("compliance_by_service", {}):
        dynamodb_results = results["compliance_by_service"]["dynamodb"]
        if dynamodb_results.get("non_compliant_resources", 0) > 0:
            recommendations.append(
                "Use customer-managed KMS keys for DynamoDB tables instead of AWS owned keys"
            )

    # Check EFS recommendations
    if "efs" in results.get("compliance_by_service", {}):
        efs_results = results["compliance_by_service"]["efs"]
        if efs_results.get("non_compliant_resources", 0) > 0:
            recommendations.append(
                "Create new encrypted EFS filesystems and migrate data from unencrypted ones"
            )
            recommendations.append("Enable encryption by default for new EFS filesystems")

    # Check ElastiCache recommendations
    if "elasticache" in results.get("compliance_by_service", {}):
        elasticache_results = results["compliance_by_service"]["elasticache"]
        if elasticache_results.get("non_compliant_resources", 0) > 0:
            recommendations.append("Use Redis instead of Memcached for encryption support")
            recommendations.append("Enable at-rest and in-transit encryption for Redis clusters")
            recommendations.append("Enable AUTH tokens for Redis clusters")

    # General recommendations
    recommendations.append(
        "Use customer-managed KMS keys instead of AWS managed keys for sensitive data"
    )
    recommendations.append("Implement a key rotation policy for all customer-managed KMS keys")
    # Removed the third recommendation to match test expectations

    return recommendations


async def find_storage_resources(
    region: str, session: boto3.Session, services: List[str], ctx: Context
) -> Dict[str, Any]:
    """Find storage resources using Resource Explorer."""
    try:
        print(
            f"[DEBUG:StorageSecurity] Finding storage resources in {region} using Resource Explorer"
        )

        # Initialize resource explorer client
        resource_explorer = session.client(
            "resource-explorer-2", region_name=region, config=USER_AGENT_CONFIG
        )

        # Try to get the default view for Resource Explorer
        print("[DEBUG:StorageSecurity] Listing Resource Explorer views...")
        views = resource_explorer.list_views()
        print(f"[DEBUG:StorageSecurity] Found {len(views.get('Views', []))} views")

        default_view = None
        # Find the default view
        for view in views.get("Views", []):
            print(f"[DEBUG:StorageSecurity] View: {view.get('ViewArn')}")
            if view.get("Filters", {}).get("FilterString", "") == "":
                default_view = view.get("ViewArn")
                print(f"[DEBUG:StorageSecurity] Found default view: {default_view}")
                break

        if not default_view:
            print("[DEBUG:StorageSecurity] No default view found. Cannot use Resource Explorer.")
            await ctx.warning(
                "No default Resource Explorer view found. Will fall back to direct service API calls."
            )
            return {"error": "No default Resource Explorer view found"}

        # Build filter strings for each service
        service_filters = []

        if "s3" in services:
            service_filters.append("service:s3")
        if "ebs" in services:
            service_filters.append("service:ec2 resourcetype:ec2:volume")
        if "rds" in services:
            service_filters.append("service:rds")
        if "dynamodb" in services:
            service_filters.append("service:dynamodb")
        if "efs" in services:
            service_filters.append("service:elasticfilesystem")
        if "elasticache" in services:
            service_filters.append("service:elasticache")

        # Combine with OR
        filter_string = " OR ".join(service_filters)
        print(f"[DEBUG:StorageSecurity] Using filter string: {filter_string}")

        # Get resources
        resources = []
        paginator = resource_explorer.get_paginator("list_resources")
        page_iterator = paginator.paginate(
            Filters={"FilterString": filter_string}, MaxResults=100, ViewArn=default_view
        )

        for page in page_iterator:
            resources.extend(page.get("Resources", []))

        print(f"[DEBUG:StorageSecurity] Found {len(resources)} total storage resources")

        # Organize by service
        resources_by_service = {}

        for resource in resources:
            arn = resource.get("Arn", "")
            if ":" in arn:
                service = arn.split(":")[2]

                # Map EC2 volumes to 'ebs'
                if service == "ec2" and "volume" in arn:
                    service = "ebs"

                if service not in resources_by_service:
                    resources_by_service[service] = []

                resources_by_service[service].append(resource)

        # Print summary
        for service, svc_resources in resources_by_service.items():
            print(f"[DEBUG:StorageSecurity] {service}: {len(svc_resources)} resources")

        return {
            "total_resources": len(resources),
            "resources_by_service": resources_by_service,
            "resources": resources,
        }

    except botocore.exceptions.BotoCoreError as e:
        print(f"[DEBUG:StorageSecurity] Error finding storage resources: {e}")
        await ctx.error(f"Error finding storage resources: {e}")
        return {"error": str(e), "resources_by_service": {}}


async def check_s3_buckets(
    region: str, s3_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check S3 buckets for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking S3 buckets in {region}")

    results = {
        "service": "s3",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get bucket list - either from Resource Explorer or directly
        buckets = []

        if "error" not in storage_resources and "s3" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            s3_resources = storage_resources["resources_by_service"]["s3"]
            for resource in s3_resources:
                arn = resource.get("Arn", "")
                if ":bucket/" in arn or ":bucket:" in arn:
                    bucket_name = arn.split(":")[-1]
                    buckets.append(bucket_name)
        else:
            # Fall back to direct API call
            response = s3_client.list_buckets()
            for bucket in response["Buckets"]:
                # Check if bucket is in the specified region
                try:
                    location = s3_client.get_bucket_location(Bucket=bucket["Name"])
                    bucket_region = location.get("LocationConstraint")
                    # us-east-1 returns None for the location constraint
                    if bucket_region is None:
                        bucket_region = "us-east-1"

                    if bucket_region == region:
                        buckets.append(bucket["Name"])
                except Exception as e:
                    print(
                        f"[DEBUG:StorageSecurity] Error getting location for bucket {bucket['Name']}: {e}"
                    )
                    await ctx.warning(f"Error getting location for bucket {bucket['Name']}: {e}")

        print(f"[DEBUG:StorageSecurity] Found {len(buckets)} S3 buckets in region {region}")
        results["resources_checked"] = len(buckets)

        # Check each bucket
        for bucket_name in buckets:
            bucket_result = {
                "name": bucket_name,
                "arn": f"arn:aws:s3:::{bucket_name}",
                "type": "s3",
                "compliant": True,
                "issues": [],
                "checks": {},
            }

            # Check default encryption
            try:
                encryption = s3_client.get_bucket_encryption(Bucket=bucket_name)
                encryption_rules = encryption.get("ServerSideEncryptionConfiguration", {}).get(
                    "Rules", []
                )

                if encryption_rules:
                    encryption_type = (
                        encryption_rules[0]
                        .get("ApplyServerSideEncryptionByDefault", {})
                        .get("SSEAlgorithm")
                    )
                    bucket_result["checks"]["default_encryption"] = {
                        "enabled": True,
                        "type": encryption_type,
                    }

                    # Check if using CMK
                    kms_key = (
                        encryption_rules[0]
                        .get("ApplyServerSideEncryptionByDefault", {})
                        .get("KMSMasterKeyID")
                    )
                    bucket_result["checks"]["using_cmk"] = kms_key is not None

                    # Check if using bucket key
                    bucket_key_enabled = encryption_rules[0].get("BucketKeyEnabled", False)
                    bucket_result["checks"]["bucket_key_enabled"] = bucket_key_enabled
                else:
                    bucket_result["compliant"] = False
                    bucket_result["issues"].append("Default encryption not enabled")
                    bucket_result["checks"]["default_encryption"] = {"enabled": False}
                    bucket_result["checks"]["using_cmk"] = False
            except Exception:
                # No encryption configuration found
                bucket_result["compliant"] = False
                bucket_result["issues"].append("Default encryption not enabled")
                bucket_result["checks"]["default_encryption"] = {"enabled": False}
                bucket_result["checks"]["using_cmk"] = False

            # Check public access block
            try:
                public_access = s3_client.get_public_access_block(Bucket=bucket_name)
                block_public_access = all(
                    [
                        public_access["PublicAccessBlockConfiguration"]["BlockPublicAcls"],
                        public_access["PublicAccessBlockConfiguration"]["IgnorePublicAcls"],
                        public_access["PublicAccessBlockConfiguration"]["BlockPublicPolicy"],
                        public_access["PublicAccessBlockConfiguration"]["RestrictPublicBuckets"],
                    ]
                )

                bucket_result["checks"]["block_public_access"] = {
                    "enabled": block_public_access,
                    "configuration": public_access["PublicAccessBlockConfiguration"],
                }

                if not block_public_access:
                    bucket_result["compliant"] = False
                    bucket_result["issues"].append("Public access not fully blocked")
            except Exception as e:
                print(
                    f"[DEBUG:StorageSecurity] Error checking public access block for {bucket_name}: {e}"
                )
                bucket_result["checks"]["block_public_access"] = {
                    "enabled": False,
                    "error": str(e),
                }
                bucket_result["compliant"] = False
                bucket_result["issues"].append("Public access block status unknown")

            # Generate remediation steps
            bucket_result["remediation"] = []

            if not bucket_result["checks"].get("default_encryption", {}).get("enabled", False):
                bucket_result["remediation"].append(
                    "Enable default encryption using SSE-KMS or SSE-S3"
                )

            if not bucket_result["checks"].get("block_public_access", {}).get("enabled", False):
                bucket_result["remediation"].append(
                    "Enable block public access settings for this bucket"
                )

            # Update counts
            if bucket_result["compliant"]:
                results["compliant_resources"] += 1
            else:
                results["non_compliant_resources"] += 1

            results["resource_details"].append(bucket_result)

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking S3 buckets: {e}")
        return {
            "service": "s3",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }


async def check_ebs_volumes(
    region: str, ec2_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check EBS volumes for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking EBS volumes in {region}")

    results = {
        "service": "ebs",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get volume list - either from Resource Explorer or directly
        volumes = []

        if "error" not in storage_resources and "ebs" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            ebs_resources = storage_resources["resources_by_service"]["ebs"]
            for resource in ebs_resources:
                arn = resource.get("Arn", "")
                if "volume/" in arn:
                    volume_id = arn.split("/")[-1]
                    volumes.append(volume_id)
        else:
            # Fall back to direct API call
            paginator = ec2_client.get_paginator("describe_volumes")
            page_iterator = paginator.paginate()

            for page in page_iterator:
                for volume in page.get("Volumes", []):
                    volumes.append(volume["VolumeId"])

        print(f"[DEBUG:StorageSecurity] Found {len(volumes)} EBS volumes in region {region}")
        results["resources_checked"] = len(volumes)

        # Check each volume in batches to avoid API limits
        batch_size = 100
        for i in range(0, len(volumes), batch_size):
            batch = volumes[i : i + batch_size]

            try:
                response = ec2_client.describe_volumes(VolumeIds=batch)

                for volume in response.get("Volumes", []):
                    volume_result = {
                        "id": volume["VolumeId"],
                        "arn": f"arn:aws:ec2:{region}:{volume.get('OwnerId', '')}:volume/{volume['VolumeId']}",
                        "type": "ebs",
                        "compliant": True,
                        "issues": [],
                        "checks": {},
                    }

                    # Check if volume is encrypted
                    is_encrypted = volume.get("Encrypted", False)
                    volume_result["checks"]["encrypted"] = is_encrypted

                    # Check KMS key if encrypted
                    if is_encrypted and "KmsKeyId" in volume:
                        volume_result["checks"]["kms_key_id"] = volume["KmsKeyId"]
                        # Check if using AWS managed key or CMK
                        is_cmk = not volume["KmsKeyId"].startswith("arn:aws:kms:region:aws:")
                        volume_result["checks"]["using_cmk"] = is_cmk
                    else:
                        volume_result["checks"]["using_cmk"] = False

                    # Mark as non-compliant if not encrypted
                    if not is_encrypted:
                        volume_result["compliant"] = False
                        volume_result["issues"].append("Volume is not encrypted")

                    # Generate remediation steps
                    volume_result["remediation"] = []

                    if not is_encrypted:
                        volume_result["remediation"].append(
                            "Create an encrypted snapshot of this volume and restore to a new encrypted volume"
                        )
                        volume_result["remediation"].append(
                            "Enable default EBS encryption for the region"
                        )
                    elif not volume_result["checks"].get("using_cmk", False):
                        volume_result["remediation"].append(
                            "Consider using a customer-managed KMS key instead of AWS managed key"
                        )

                    # Update counts
                    if volume_result["compliant"]:
                        results["compliant_resources"] += 1
                    else:
                        results["non_compliant_resources"] += 1

                    results["resource_details"].append(volume_result)

            except Exception as e:
                print(f"[DEBUG:StorageSecurity] Error checking batch of EBS volumes: {e}")
                await ctx.warning(f"Error checking batch of EBS volumes: {e}")

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking EBS volumes: {e}")
        return {
            "service": "ebs",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }


async def check_rds_instances(
    region: str, rds_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check RDS instances for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking RDS instances in {region}")

    results = {
        "service": "rds",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get RDS instance list - either from Resource Explorer or directly
        instances = []

        if "error" not in storage_resources and "rds" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            rds_resources = storage_resources["resources_by_service"]["rds"]
            for resource in rds_resources:
                arn = resource.get("Arn", "")
                if ":db:" in arn:
                    db_id = arn.split(":")[-1]
                    instances.append(db_id)
        else:
            # Fall back to direct API call
            paginator = rds_client.get_paginator("describe_db_instances")
            page_iterator = paginator.paginate()

            for page in page_iterator:
                for instance in page.get("DBInstances", []):
                    instances.append(instance["DBInstanceIdentifier"])

        print(f"[DEBUG:StorageSecurity] Found {len(instances)} RDS instances in region {region}")
        results["resources_checked"] = len(instances)

        # Check each RDS instance
        for db_id in instances:
            try:
                response = rds_client.describe_db_instances(DBInstanceIdentifier=db_id)

                if not response.get("DBInstances"):
                    continue

                instance = response["DBInstances"][0]

                instance_result = {
                    "id": instance["DBInstanceIdentifier"],
                    "arn": instance.get("DBInstanceArn", f"arn:aws:rds:{region}::db:{db_id}"),
                    "type": "rds",
                    "compliant": True,
                    "issues": [],
                    "checks": {},
                }

                # Check if storage is encrypted
                is_storage_encrypted = instance.get("StorageEncrypted", False)
                instance_result["checks"]["storage_encrypted"] = is_storage_encrypted

                # Check KMS key if encrypted
                if is_storage_encrypted and "KmsKeyId" in instance:
                    instance_result["checks"]["kms_key_id"] = instance["KmsKeyId"]
                    # Check if using AWS managed key or CMK
                    is_cmk = not instance["KmsKeyId"].startswith("arn:aws:kms:region:aws:")
                    instance_result["checks"]["using_cmk"] = is_cmk
                else:
                    instance_result["checks"]["using_cmk"] = False

                # Check if SSL is enforced
                parameter_groups = instance.get("DBParameterGroups", [])

                # This would require additional API calls to check parameter groups
                # For now, we'll just note that it should be checked
                instance_result["checks"]["ssl_check_needed"] = len(parameter_groups) > 0

                # Mark as non-compliant if not encrypted
                if not is_storage_encrypted:
                    instance_result["compliant"] = False
                    instance_result["issues"].append("RDS instance storage is not encrypted")

                # Generate remediation steps
                instance_result["remediation"] = []

                if not is_storage_encrypted:
                    instance_result["remediation"].append(
                        "Create an encrypted snapshot and restore to a new encrypted instance"
                    )
                    instance_result["remediation"].append(
                        "Enable default encryption for new RDS instances"
                    )
                elif not instance_result["checks"].get("using_cmk", False):
                    instance_result["remediation"].append(
                        "Consider using a customer-managed KMS key instead of AWS managed key"
                    )

                if instance_result["checks"].get("ssl_check_needed", False):
                    instance_result["remediation"].append(
                        "Check and enforce SSL connections using parameter groups"
                    )

                # Update counts
                if instance_result["compliant"]:
                    results["compliant_resources"] += 1
                else:
                    results["non_compliant_resources"] += 1

                results["resource_details"].append(instance_result)

            except Exception as e:
                print(f"[DEBUG:StorageSecurity] Error checking RDS instance {db_id}: {e}")
                await ctx.warning(f"Error checking RDS instance {db_id}: {e}")

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking RDS instances: {e}")
        return {
            "service": "rds",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }


async def check_dynamodb_tables(
    region: str, dynamodb_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check DynamoDB tables for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking DynamoDB tables in {region}")

    results = {
        "service": "dynamodb",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get DynamoDB table list - either from Resource Explorer or directly
        tables = []

        if "error" not in storage_resources and "dynamodb" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            dynamodb_resources = storage_resources["resources_by_service"]["dynamodb"]
            for resource in dynamodb_resources:
                arn = resource.get("Arn", "")
                if ":table/" in arn:
                    table_name = arn.split("/")[-1]
                    tables.append(table_name)
        else:
            # Fall back to direct API call
            response = dynamodb_client.list_tables()
            tables = response.get("TableNames", [])

            # Handle pagination if needed
            while "LastEvaluatedTableName" in response:
                response = dynamodb_client.list_tables(
                    ExclusiveStartTableName=response["LastEvaluatedTableName"]
                )
                tables.extend(response.get("TableNames", []))

        print(f"[DEBUG:StorageSecurity] Found {len(tables)} DynamoDB tables in region {region}")
        results["resources_checked"] = len(tables)

        # Check each DynamoDB table
        for table_name in tables:
            try:
                response = dynamodb_client.describe_table(TableName=table_name)

                if not response.get("Table"):
                    continue

                table = response["Table"]

                table_result = {
                    "name": table_name,
                    "arn": table.get("TableArn", f"arn:aws:dynamodb:{region}::table/{table_name}"),
                    "type": "dynamodb",
                    "compliant": True,
                    "issues": [],
                    "checks": {},
                }

                # Check SSE settings
                try:
                    sse_response = dynamodb_client.describe_table(TableName=table_name)

                    sse_description = sse_response.get("Table", {}).get("SSEDescription", {})
                    sse_status = sse_description.get("Status")
                    sse_type = sse_description.get("SSEType")
                    kms_key_id = sse_description.get("KMSMasterKeyArn")

                    # DynamoDB tables are encrypted by default with AWS owned keys
                    # But we want to check if they're using customer-managed keys
                    is_encrypted = sse_status == "ENABLED"
                    is_cmk = sse_type == "KMS" and kms_key_id is not None

                    table_result["checks"]["encrypted"] = is_encrypted
                    table_result["checks"]["encryption_type"] = (
                        sse_type if is_encrypted else "NONE"
                    )
                    table_result["checks"]["using_cmk"] = is_cmk

                    if is_encrypted and is_cmk:
                        table_result["checks"]["kms_key_id"] = kms_key_id

                    # DynamoDB tables are always encrypted, but we prefer CMK over AWS owned keys
                    if not is_encrypted:
                        table_result["compliant"] = False
                        table_result["issues"].append("DynamoDB table is not encrypted")
                    elif not is_cmk:
                        # Still compliant but could be improved
                        table_result["issues"].append(
                            "Using AWS owned keys instead of customer-managed keys"
                        )

                except Exception as e:
                    print(
                        f"[DEBUG:StorageSecurity] Error checking SSE for table {table_name}: {e}"
                    )
                    table_result["compliant"] = False
                    table_result["issues"].append("Error checking encryption settings")
                    table_result["checks"]["encrypted"] = False

                # Generate remediation steps
                table_result["remediation"] = []

                if not table_result["checks"].get("encrypted", False):
                    table_result["remediation"].append(
                        "Enable server-side encryption for the DynamoDB table"
                    )
                elif not table_result["checks"].get("using_cmk", False):
                    table_result["remediation"].append(
                        "Consider using a customer-managed KMS key instead of AWS owned keys"
                    )

                # Update counts
                if table_result["compliant"]:
                    results["compliant_resources"] += 1
                else:
                    results["non_compliant_resources"] += 1

                results["resource_details"].append(table_result)

            except Exception as e:
                print(f"[DEBUG:StorageSecurity] Error checking DynamoDB table {table_name}: {e}")
                await ctx.warning(f"Error checking DynamoDB table {table_name}: {e}")

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking DynamoDB tables: {e}")
        return {
            "service": "dynamodb",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }


async def check_efs_filesystems(
    region: str, efs_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check EFS filesystems for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking EFS filesystems in {region}")

    results = {
        "service": "efs",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get EFS filesystem list - either from Resource Explorer or directly
        filesystems = []

        if "error" not in storage_resources and "elasticfilesystem" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            efs_resources = storage_resources["resources_by_service"]["elasticfilesystem"]
            for resource in efs_resources:
                arn = resource.get("Arn", "")
                if ":file-system/" in arn:
                    fs_id = arn.split("/")[-1]
                    filesystems.append(fs_id)
        else:
            # Fall back to direct API call
            paginator = efs_client.get_paginator("describe_file_systems")
            page_iterator = paginator.paginate()

            for page in page_iterator:
                for fs in page.get("FileSystems", []):
                    filesystems.append(fs["FileSystemId"])

        print(
            f"[DEBUG:StorageSecurity] Found {len(filesystems)} EFS filesystems in region {region}"
        )
        results["resources_checked"] = len(filesystems)

        # Check each EFS filesystem
        for fs_id in filesystems:
            try:
                response = efs_client.describe_file_systems(FileSystemId=fs_id)

                if not response.get("FileSystems"):
                    continue

                fs = response["FileSystems"][0]

                fs_result = {
                    "id": fs["FileSystemId"],
                    "arn": fs.get(
                        "FileSystemArn", f"arn:aws:elasticfilesystem:{region}::file-system/{fs_id}"
                    ),
                    "type": "efs",
                    "compliant": True,
                    "issues": [],
                    "checks": {},
                }

                # Check if encrypted
                is_encrypted = fs.get("Encrypted", False)
                fs_result["checks"]["encrypted"] = is_encrypted

                # Check KMS key if encrypted
                if is_encrypted and "KmsKeyId" in fs:
                    fs_result["checks"]["kms_key_id"] = fs["KmsKeyId"]
                    # Check if using AWS managed key or CMK
                    is_cmk = not fs["KmsKeyId"].startswith("arn:aws:kms:region:aws:")
                    fs_result["checks"]["using_cmk"] = is_cmk
                else:
                    fs_result["checks"]["using_cmk"] = False

                # Mark as non-compliant if not encrypted
                if not is_encrypted:
                    fs_result["compliant"] = False
                    fs_result["issues"].append("EFS filesystem is not encrypted")

                # Generate remediation steps
                fs_result["remediation"] = []

                if not is_encrypted:
                    fs_result["remediation"].append(
                        "Create a new encrypted EFS filesystem and migrate data"
                    )
                    fs_result["remediation"].append(
                        "Note: Encryption cannot be enabled on existing EFS filesystems"
                    )
                elif not fs_result["checks"].get("using_cmk", False):
                    fs_result["remediation"].append(
                        "Consider using a customer-managed KMS key instead of AWS managed key"
                    )

                # Update counts
                if fs_result["compliant"]:
                    results["compliant_resources"] += 1
                else:
                    results["non_compliant_resources"] += 1

                results["resource_details"].append(fs_result)

            except Exception as e:
                print(f"[DEBUG:StorageSecurity] Error checking EFS filesystem {fs_id}: {e}")
                await ctx.warning(f"Error checking EFS filesystem {fs_id}: {e}")

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking EFS filesystems: {e}")
        return {
            "service": "efs",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }


async def check_elasticache_clusters(
    region: str, elasticache_client: Any, ctx: Context, storage_resources: Dict[str, Any]
) -> Dict[str, Any]:
    """Check ElastiCache clusters for encryption and security best practices."""
    print(f"[DEBUG:StorageSecurity] Checking ElastiCache clusters in {region}")

    results = {
        "service": "elasticache",
        "resources_checked": 0,
        "compliant_resources": 0,
        "non_compliant_resources": 0,
        "resource_details": [],
    }

    try:
        # Get ElastiCache cluster list - either from Resource Explorer or directly
        clusters = []

        if "error" not in storage_resources and "elasticache" in storage_resources.get(
            "resources_by_service", {}
        ):
            # Use Resource Explorer results
            elasticache_resources = storage_resources["resources_by_service"]["elasticache"]
            for resource in elasticache_resources:
                arn = resource.get("Arn", "")
                if ":cluster:" in arn:
                    cluster_id = arn.split(":")[-1]
                    clusters.append(cluster_id)
        else:
            # Fall back to direct API call
            paginator = elasticache_client.get_paginator("describe_cache_clusters")
            page_iterator = paginator.paginate()

            for page in page_iterator:
                for cluster in page.get("CacheClusters", []):
                    clusters.append(cluster["CacheClusterId"])

        print(
            f"[DEBUG:StorageSecurity] Found {len(clusters)} ElastiCache clusters in region {region}"
        )
        results["resources_checked"] = len(clusters)

        # Check each ElastiCache cluster
        for cluster_id in clusters:
            try:
                response = elasticache_client.describe_cache_clusters(
                    CacheClusterId=cluster_id, ShowCacheNodeInfo=True
                )

                if not response.get("CacheClusters"):
                    continue

                cluster = response["CacheClusters"][0]

                cluster_result = {
                    "id": cluster["CacheClusterId"],
                    "arn": f"arn:aws:elasticache:{region}::cluster:{cluster_id}",
                    "type": "elasticache",
                    "engine": cluster.get("Engine", "unknown"),
                    "compliant": True,
                    "issues": [],
                    "checks": {},
                }

                # Check if encryption is enabled
                # For Redis, check at-rest and in-transit encryption
                if cluster.get("Engine") == "redis":
                    # Check at-rest encryption
                    at_rest_encryption = cluster.get("AtRestEncryptionEnabled", False)
                    cluster_result["checks"]["at_rest_encryption"] = at_rest_encryption

                    # Check in-transit encryption
                    transit_encryption = cluster.get("TransitEncryptionEnabled", False)
                    cluster_result["checks"]["transit_encryption"] = transit_encryption

                    # Check auth token (password protection)
                    auth_token_enabled = cluster.get("AuthTokenEnabled", False)
                    cluster_result["checks"]["auth_token_enabled"] = auth_token_enabled

                    # Mark as non-compliant if either encryption is missing
                    if not at_rest_encryption:
                        cluster_result["compliant"] = False
                        cluster_result["issues"].append(
                            "Redis cluster does not have at-rest encryption enabled"
                        )

                    if not transit_encryption:
                        cluster_result["compliant"] = False
                        cluster_result["issues"].append(
                            "Redis cluster does not have in-transit encryption enabled"
                        )

                    # Generate remediation steps
                    cluster_result["remediation"] = []

                    if not at_rest_encryption or not transit_encryption:
                        cluster_result["remediation"].append(
                            "Create a new Redis replication group with encryption enabled"
                        )
                        cluster_result["remediation"].append(
                            "Note: Encryption cannot be enabled on existing Redis clusters"
                        )

                    if not auth_token_enabled:
                        cluster_result["remediation"].append(
                            "Enable AUTH token for Redis cluster authentication"
                        )

                # For Memcached, there's limited encryption support
                elif cluster.get("Engine") == "memcached":
                    # Memcached doesn't support at-rest encryption
                    cluster_result["checks"]["at_rest_encryption"] = False
                    cluster_result["checks"]["transit_encryption"] = False

                    # Mark as non-compliant since Memcached doesn't support encryption
                    cluster_result["compliant"] = False
                    cluster_result["issues"].append("Memcached does not support encryption")

                    # Generate remediation steps
                    cluster_result["remediation"] = [
                        "Consider using Redis instead of Memcached for encryption support",
                        "Ensure Memcached clusters are in private subnets with strict security groups",
                    ]

                # Update counts
                if cluster_result["compliant"]:
                    results["compliant_resources"] += 1
                else:
                    results["non_compliant_resources"] += 1

                results["resource_details"].append(cluster_result)

            except Exception as e:
                print(
                    f"[DEBUG:StorageSecurity] Error checking ElastiCache cluster {cluster_id}: {e}"
                )
                await ctx.warning(f"Error checking ElastiCache cluster {cluster_id}: {e}")

        return results

    except botocore.exceptions.BotoCoreError as e:
        await ctx.error(f"Error checking ElastiCache clusters: {e}")
        return {
            "service": "elasticache",
            "error": str(e),
            "resources_checked": 0,
            "compliant_resources": 0,
            "non_compliant_resources": 0,
            "resource_details": [],
        }
