from pcluster.aws.common import AWSExceptionHandler, Boto3Client


class ImageBuilderClient(Boto3Client):
    """Imagebuilder Boto3 client."""

    def __init__(self):
        super().__init__("imagebuilder")

    @AWSExceptionHandler.handle_client_exception
    def get_image_resources(self, image_arn):
        """Get image info by ami arn."""
        return self._client.get_image(imageBuildVersionArn=image_arn)

    @AWSExceptionHandler.handle_client_exception
    def get_image_state(self, image_arn):
        """Get image state by ami arn."""
        return self.get_image_resources(image_arn).get("image", {}).get("state", {})

    @AWSExceptionHandler.handle_client_exception
    def get_image_id(self, image_arn):
        """Retrieve image id by image arn."""
        ami_list = self.get_image_resources(image_arn).get("image", {}).get("outputResources", {}).get("amis", None)
        return ami_list[0].get("image") if ami_list else None