from transformers import ViTImageProcessor, ViTForImageClassification
from PIL import Image
import streamlit as st
import os

# Attempt to import and initialize the Groq client
try:
    from groq import Groq
    groq_key = os.getenv("GROQ_KEY")
    if not groq_key:
        raise ValueError("Environment variable GROQ_KEY not found.")
    client = Groq(api_key=groq_key)
except ImportError as e:
    st.error("Failed to import Groq client or environment variable. Ensure GROQ_KEY is set.")
except Exception as e:
    st.error(f"Error initializing Groq client: {e}")

model_id = 'google/vit-base-patch16-224'

# Streamlit app title
st.title("Product Detection App")

# File uploader for image input
uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"])

if uploaded_file is not None:
    image = Image.open(uploaded_file)
    processor = ViTImageProcessor.from_pretrained(model_id)
    model = ViTForImageClassification.from_pretrained(model_id)
    
    # Process the image and make predictions
    inputs = processor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_idx = logits.argmax(-1).item()
    predicted_class = model.config.id2label[predicted_class_idx]

    # Display the predicted class
    st.write(f"Predicted Class: {predicted_class}")

    # Generate keywords using Groq
    try:
        chat_completion = client.chat.completions.create(
            model="llama3-8b-8192",
            messages=[
                {
                    "role": "system",
                    "content": "You are a seasoned digital marketer with over 15 years of experience specializing in search engine optimization (SEO) and conversion rate optimization (CRO). You have a deep understanding of consumer behavior and know how to craft high-converting keywords that can significantly boost product visibility and sales. Please respond in JSON format."
                },
                {
                    "role": "user",
                    "content": f"Your task is to generate a list of 10 high-converting keywords tailored for selling {predicted_class}. The details of the product for which you need to generate keywords are as follows: - Product Name - Product Category - Target Audience - Key Features/Benefits - Unique Selling Proposition (USP). Keep in mind that the keywords should be relevant, engaging, and designed to attract potential customers effectively. Consider trends, search intent, and competitor analysis to inform your selections."
                }
            ],
             response_format={"type": "json_object"}  # Adjusted based on common practices
        )
        
        # Display the generated keywords
        st.write("Generated Keywords:")
        st.write(chat_completion.choices[0].message.content)
    
    except Exception as e:
        st.error(f"Error generating keywords: {e}")