"""
title: Bedrock Image Description
author: Josh Knapp
version: 0.1.0
description="Provide Direct Bedrock call for image generation"
"""

import subprocess
import json
from pydantic import BaseModel, Field


# Try to import boto3, install if not present
try:
    import boto3
except ImportError:
    print("boto3 package not found. Attempting to install...")
    try:
        subprocess.check_call([sys.executable, "-m", "pip", "install", "boto3"])
        import boto3

        print("boto3 package installed successfully")
    except subprocess.CalledProcessError as e:
        print(f"Failed to install boto3 package: {str(e)}")


class Tools:
    class Valves(BaseModel):
        AWS_ACCESS_KEY: str = Field(
            default="",
            description="AWS Access Key",
        )
        AWS_SECRET_KEY: str = Field(
            default="",
            description="AWS Secret Key",
        )
        AWS_BEDROCK_MODEL: str = Field(
            default="",
            description="AWS Bedrock Model to use"
        )

    def __init__(self):
        self.valves = self.Valves()
        pass

    def analyze_image(self, base64_image: str) -> str:
        """
        Analyze an image using AWS Bedrock's vision model
        Args:
            base64_image (str): Base64 encoded image string
        Returns:
            str: Description of the image
        """
        try:
            # Initialize Bedrock runtime client
            bedrock = boto3.client(
                service_name="bedrock-runtime",
                aws_access_key_id=self.valves.AWS_ACCESS_KEY,
                aws_secret_access_key=self.valves.AWS_SECRET_KEY,
                region_name="us-east-1"  # or your preferred region
            )

            # Prepare the request body
            request_body = {
                "anthropic_version": "bedrock-2023-05-31",
                "max_tokens": 1000,
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "image",
                                "source": {
                                    "type": "base64",
                                    "media_type": "image/jpeg",
                                    "data": base64_image
                                }
                            },
                            {
                                "type": "text",
                                "text": "Please describe this image in detail."
                            }
                        ]
                    }
                ]
            }

            # Invoke the model
            response = bedrock.invoke_model(
                modelId=self.valves.AWS_BEDROCK_MODEL,
                body=json.dumps(request_body)
            )

            # Parse and return the response
            response_body = json.loads(response['body'].read())
            return response_body['messages'][0]['content'][0]['text']

        except Exception as e:
            print(f"Error analyzing image: {str(e)}")
            return f"Error analyzing image: {str(e)}"