import PIL.Image
import os
import google.generativeai as genai
from prompt import prompt
import google.ai.generativelanguage as glm
import json

# Configure API key for Google Generative AI
os.environ['GOOGLE_API_KEY'] = "AIzaSyDAG_Xl66vh4ceY81UXe3vrdwP6wIAkpBs"
genai.configure(api_key=os.environ['GOOGLE_API_KEY'])



def get_image_data(image):
    # Initialize the vision model
    try:
        vision_model = genai.GenerativeModel('gemini-1.5-flash')
    except Exception as e:
        print(f"Error initializing vision model: {e}")
        exit()

    # Generate content using the vision model
    try:
        response = vision_model.generate_content([prompt, image],
                                                    generation_config=genai.types.GenerationConfig(
                                                    candidate_count=1,
                                                    stop_sequences=['.'],
                                                    max_output_tokens=200,
                                                    top_p=0.7,
                                                    top_k=4,
                                                    temperature=0.7,
                                                    )  
                                                )

        # Parse the response text to extract category and type
        response_text = response.text.strip().lower()
        category = response_text.split('category:')[1].split('\n')[0].strip() if 'category:' in response_text else 'Unknown'
        type_name = response_text.split('type:')[1].strip() if 'type:' in response_text else 'Unknown'

        # Create the result dictionary
        result = {
            "Category": category.capitalize(),
            "Type": type_name.capitalize()
        }

        # Print the result dictionary
        return(result)

    except Exception as e:
        print(f"Error generating content: {e}")



# # Load the image
# image_path = 'image.png'
# try:
#     image = PIL.Image.open(image_path)
# except Exception as e:
#     print(f"Error loading image: {e}")
#     exit()