import streamlit as st
from dotenv import load_dotenv
import os
from groq import Groq
import base64
from PIL import Image
import io

# Set page config
st.set_page_config(
    page_title="Image Analysis with AI",
    page_icon="🖼️",
    layout="centered"
)

# Add title and description
st.title("AI Image Analysis")
st.markdown("Upload an image or provide an image URL to get AI-generated commentary.")

# Load environment variables
load_dotenv()

# Initialize Groq client with the correct configuration
api_key = os.getenv("GROQ_API_KEY")
if not api_key:
    st.error("Please set your GROQ_API_KEY in the .env file")
    st.stop()

client = Groq(
    api_key=api_key
)

# Create input field for image URL
image_url = st.text_input("Enter Image URL", "https://static.seekingalpha.com/uploads/2016/1/19/saupload_fredgraph.jpg")

# Add a file uploader
uploaded_file = st.file_uploader("Or upload an image", type=["jpg", "jpeg", "png"])

def get_image_url():
    if uploaded_file is not None:
        # Convert uploaded file to base64
        image = Image.open(uploaded_file)
        buffered = io.BytesIO()
        image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        return f"data:image/png;base64,{img_str}"
    return image_url

# Add a button to trigger analysis
if st.button("Analyze Image"):
    try:
        # Show loading spinner
        with st.spinner("Analyzing image..."):
            current_image_url = get_image_url()
            
            # Create the completion request
            completion = client.chat.completions.create(
                model="meta-llama/llama-4-scout-17b-16e-instruct",  # Updated model name
                messages=[
                    {
                        "role": "user",
                        "content": [
                            {
                                "type": "text",
                                "text": "Write a detailed commentary on the trend observed in the image?"
                            },
                            {
                                "type": "image_url",
                                "image_url": {
                                    "url": current_image_url
                                }
                            }
                        ]
                    }
                ],
                temperature=1,
                max_tokens=300,  # Updated parameter name
                top_p=1,
                stream=False
            )

            # Display the image
            if uploaded_file is not None:
                st.image(uploaded_file, caption="Analyzed Image", use_column_width=True)
            else:
                st.image(image_url, caption="Analyzed Image", use_column_width=True)
            
            # Display the analysis
            st.subheader("AI Analysis")
            st.write(completion.choices[0].message.content)

    except Exception as e:
        st.error(f"An error occurred: {str(e)}")

# Add footer
st.markdown("---")
st.markdown("Built with Streamlit and Groq AI")