from .base import BaseCheck
from botocore.exceptions import ClientError

class AMIsCheck(BaseCheck):
    """Check for custom AMI usage"""
    
    MANAGED_AMI_TYPES = {
        'AL2_x86_64', 'AL2_x86_64_GPU', 'AL2_ARM_64', 
        'AL2023_x86_64_STANDARD', 'AL2023_ARM_64_STANDARD', 
        'BOTTLEROCKET_ARM_64', 'BOTTLEROCKET_x86_64'
    }
    
    def _execute_check(self):
        custom_ami_resources = []
        
        # Check managed node groups for launch templates (which may have custom AMIs)
        node_groups = self.eks_client.list_nodegroups(clusterName=self.cluster_name)
        for ng_name in node_groups.get('nodegroups', []):
            try:
                ng_details = self.eks_client.describe_nodegroup(
                    clusterName=self.cluster_name,
                    nodegroupName=ng_name
                )['nodegroup']
                
                # Check if using launch template with custom AMI
                if ng_details.get('launchTemplate'):
                    lt_id = ng_details['launchTemplate']['id']
                    lt_version = ng_details['launchTemplate'].get('version', '$Latest')
                    
                    try:
                        # Get launch template details to check AMI
                        lt_response = self.ec2_client.describe_launch_template_versions(
                            LaunchTemplateId=lt_id,
                            Versions=[lt_version]
                        )
                        
                        if lt_response['LaunchTemplateVersions']:
                            lt_data = lt_response['LaunchTemplateVersions'][0]['LaunchTemplateData']
                            if 'ImageId' in lt_data:
                                # Has custom AMI specified in launch template
                                ami_id = lt_data['ImageId']
                                custom_ami_resources.append(f"Node group {ng_name}: custom AMI {ami_id} in launch template {lt_id}")
                    except ClientError:
                        # If we can't check the launch template, assume it might be custom
                        custom_ami_resources.append(f"Node group {ng_name}: unable to verify launch template {lt_id}")
                
                # Check AMI type - if it's not a standard AWS managed type
                ami_type = ng_details.get('amiType', '')
                if ami_type and ami_type not in self.MANAGED_AMI_TYPES:
                    custom_ami_resources.append(f"Node group {ng_name}: unsupported AMI type {ami_type}")
            except ClientError:
                continue
        
        if custom_ami_resources:
            resource_summary = ', '.join(custom_ami_resources[:3])
            if len(custom_ami_resources) > 3:
                resource_summary += '...'
            
            return self._create_result(
                'FAIL',
                f'Found {len(custom_ami_resources)} custom AMI resources',
                [
                    'Custom AMIs not supported in Auto Mode',
                    'Switch to AWS managed AMIs',
                    f'Resources: {resource_summary}'
                ]
            )
        else:
            return self._create_result(
                'PASS',
                'Using AWS managed AMIs'
            )