import os
import streamlit as st
import easyocr
from openai import OpenAI, OpenAIError  # ✅ Correct import for API errors
from dotenv import load_dotenv
from PIL import Image
import io

# ✅ Load API key from .env
load_dotenv()
api_key = os.getenv("GROQ_API_KEY")

if not api_key:
    st.error("❌ API key not found! Please set `GROQ_API_KEY` in your `.env` file.")
    st.stop()

# ✅ Initialize OpenAI client for Groq API
client = OpenAI(api_key=api_key)

# ✅ Ensure Streamlit config is first
st.set_page_config(page_title="Multimodal AI Assistant", layout="wide")

# ✅ Initialize OCR Reader
reader = easyocr.Reader(["en"])

# ✅ Streamlit App Layout
st.title("📸 Multimodal AI Assistant")
st.write("Upload an image and ask questions based on the extracted text.")

# ✅ File Uploader
uploaded_file = st.file_uploader("Upload an image", type=["png", "jpg", "jpeg"])

if uploaded_file:
    # ✅ Display uploaded image
    image = Image.open(uploaded_file)
    st.image(image, caption="Uploaded Image", use_container_width=True)

    # ✅ Convert image to bytes for EasyOCR
    image_bytes = io.BytesIO(uploaded_file.getvalue()).read()

    # ✅ Extract text using EasyOCR
    with st.spinner("🔍 Extracting text..."):
        extracted_text = reader.readtext(image_bytes, detail=0)

    # ✅ Show extracted text
    extracted_text_str = " ".join(extracted_text) if extracted_text else "No text found"
    st.subheader("📝 Extracted Text:")
    st.write(extracted_text_str)

    # ✅ Question Answering Section
    user_query = st.text_input("Ask a question about the extracted text:")

    if st.button("Get Answer"):
        if not extracted_text_str or extracted_text_str == "No text found":
            st.warning("⚠ No text available for analysis.")
        elif user_query.strip() == "":
            st.warning("⚠ Please enter a question.")
        else:
            with st.spinner("🤖 Thinking..."):
                try:
                    response = client.chat.completions.create(  # ✅ Corrected API call
                        model="llama3-70b-8192",  # ✅ Groq LLaMA 3 model
                        messages=[
                            {"role": "system", "content": "You are an AI assistant analyzing extracted text from images."},
                            {"role": "user", "content": f"Extracted text: {extracted_text_str}\n\nUser question: {user_query}"}
                        ]
                    )
                    answer = response.choices[0].message.content
                    st.subheader("🤖 AI Answer:")
                    st.write(answer)
                except OpenAIError as e:
                    st.error(f"❌ API Error: {e}")