import json
import logging
import base64

import boto3
from botocore.exceptions import ClientError

from ..config import get_app_config
from ..exceptions import ImageError

# Set logging
logger = logging.getLogger(__name__)

# Get configuration
conf = get_app_config()

MODEL_ID = conf['model_id']

# Create Bedrock client using environment variables for AWS credentials
client = boto3.client(
    service_name='bedrock-runtime',
    region_name=conf['region']
)

def generate_image(body):
    """
    Use Amazon Nova Canvas model to generate an image.
    
    Args:
        body (str): Request body
        
    Returns:
        bytes: Image generated by the model
    """
    logger.info(f"Generating image with Amazon Nova Canvas model {MODEL_ID}")

    accept = "application/json"
    content_type = "application/json"

    try:
        response = client.invoke_model(
            body=body, modelId=MODEL_ID, accept=accept, contentType=content_type
        )
        
        response_body = json.loads(response.get("body").read())
        base64_image = response_body.get("images")[0]
        base64_bytes = base64_image.encode('ascii')
        image_bytes = base64.b64decode(base64_bytes)

        finish_reason = response_body.get("error")
        if finish_reason is not None:
            raise ImageError(f"Image generation error. Error: {finish_reason}")

        logger.info(f"Successfully generated image with Amazon Nova Canvas model {MODEL_ID}")
        return image_bytes
        
    except ClientError as err:
        message = err.response["Error"]["Message"]
        logger.error(f"Client error occurred: {message}")
        raise ImageError(f"Client error occurred: {message}") 