import base64
import boto3
import datetime
import json
import random
import sys
from botocore.exceptions import ClientError

BUCKET_NAME = 'xxxxx-xxxxx-bucket'
s3 = boto3.client('s3', region_name="us-east-1")
bedrock = boto3.client(service_name='bedrock-runtime')

class ImageError(Exception):
    "Custom exception for errors returned by the model"

    def __init__(self, message):
        self.message = message

def lambda_handler(event, context):
    try:
        # bedrock の モデルを指定
        model_id = "amazon.titan-text-premier-v1:0"

        age = get_parameter("age", event)
        gender = get_parameter("gender", event)
        donnako = get_parameter("donnako", event)
        okuruhito = get_parameter("okuruhito", event)
        
        # プロンプトの設定
        # bedrock に投げる質問を作成
        prompt = f"""
        以下の条件に合う子に相応しいプレゼントの候補を出してほしい。
        ・年齢：{age}
        ・性別：{gender}
        ・どんな子： {donnako}
        ・贈る人：{okuruhito}

        プレゼントは具体的な商品名だけを書いてください。
        1つ目. 予算100円で買えるもの
        2つ目. 予算3000円で買えるもの
        3つ目. 予算1000円で買えるもの
        4つ目. 手作りのもの
        """
        
        # Bedrock呼び出し
        res_text = generate_text(model_id, create_text_request_body(prompt))
        logger.info(f"{res_text=}")

        # 受け取ったテキストを分割
        presents = res_text.replace("1つ目", "").replace("2つ目", "").replace("3つ目", "").replace("4つ目", "").split("\n")
        logger.info(f"{presents=}")

        for present in presents:
            now = datetime.datetime.now()
            formatted_date = now.strftime('%y%m%d-%H%M%S%f')[:-4]
            file_name = f"{present}-{formatted_date}.png"
            generate_image(present, file_name)

    except ClientError as err:
        message = err.response["Error"]["Message"]
        logger.error("A client error occurred: %s", message)
    except ImageError as err:
        logger.error(err.message)

    return {
        'statusCode': 200,
        'body': json.dumps(prompt)
    }

def create_text_request_body(prompt):
    """titan-textのリクエストするbody作成

    Args:
        prompt : 質問内容

    Returns:
        dict : リクエストボディ
    """    

    return json.dumps({
        "inputText": prompt,
        "textGenerationConfig": {
            "maxTokenCount": 3072,
            "stopSequences": [],
            "temperature": 0,
            "topP": 1
        }
    })

def create_image_request_body(prompt, seed):
    """amazon.titan-image-generator-v2:0にリクエストするbody作成

    Args:
        prompt : 質問内容
        seed : seed値

    Returns:
        dict : リクエストボディ
    """    
    return json.dumps({
        "textToImageParams": {
            "text": prompt
        },
        "taskType": "TEXT_IMAGE",
        "imageGenerationConfig": {
            "cfgScale": 5,
            "seed": seed,
            "quality": "standard",
            "width": 512,
            "height": 512,
            "numberOfImages": 1
        }
    })

def generate_text(model_id, body):
    """テキスト生成

    Args:
        model_id : Bedrockのモデル
        body : リクエストボディ

    Returns:
        str: Bedrockから戻ってきたテキスト
    """    

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

    # bedrockの実行
    response = bedrock.invoke_model(
        body=body, modelId=model_id, accept=accept, contentType=content_type
    )

    # レスポンスデータから戻り値を取得
    response_body = json.loads(response.get("body").read())    
    res = ""
    for result in response_body['results']:
        res += result['outputText']
    
    return res

def generate_image(prompt_data, file_name):
    """画像の生成

    Args:
        prompt_data : 質問内容
        file_name : 画像のファイル名
    """    

    modelId = "amazon.titan-image-generator-v2:0" 
    accept = "application/json"
    contentType = "application/json"

    # 画像のランダム性を与えるためシードを生成
    seed = random.randint(0, 2147483646)

    req_body = create_image_request_body(prompt_data, seed)
    response = bedrock.invoke_model(
        body=req_body, modelId=modelId, accept=accept, contentType=contentType
    )

    model_response = json.loads(response["body"].read())
    base64_image_data = model_response["images"][0]
    image_data = base64.b64decode(base64_image_data)

    # S3バケットへ出力
    s3.put_object(Bucket=BUCKET_NAME, Key=file_name, Body=image_data)