import streamlit as st
import requests
from PyPDF2 import PdfReader
from PIL import Image
import os
import certifi

# Load Grok API key from Hugging Face Secrets
GROK_API_KEY = os.environ.get("GROK_API_KEY")

st.set_page_config(page_title="Multimodal Learning Companion", layout="wide")
st.title("📚 Multimodal Learning Companion")
st.markdown(
    "Upload TXT, PDF, or images and get summaries and explanations for students."
)

# File uploader
uploaded_file = st.file_uploader(
    "Upload your learning content",
    type=["txt", "pdf", "png", "jpg", "jpeg"]
)

def call_grok(prompt):
    """Call Grok API and return response text."""
    url = "https://api.grok.ai/v1/chat/completions"
    headers = {
        "Authorization": f"Bearer {GROK_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": "grok-3",
        "messages": [{"role": "user", "content": prompt}]
    }
    try:
        response = requests.post(url, json=payload, headers=headers, verify=certifi.where())
        return response.json()["choices"][0]["message"]["content"]
    except Exception as e:
        st.error(f"Error calling Grok API: {e}")
        return None

if uploaded_file:
    file_name = uploaded_file.name.lower()

    if file_name.endswith(".txt"):
        text = uploaded_file.read().decode("utf-8")
        summary = call_grok(f"Summarize this text for a student:\n{text}")
        if summary:
            st.subheader("📄 Summary")
            st.write(summary)

    elif file_name.endswith(".pdf"):
        reader = PdfReader(uploaded_file)
        text = ""
        for page in reader.pages:
            page_text = page.extract_text()
            if page_text:
                text += page_text + "\n"
        summary = call_grok(f"Summarize this PDF for a student:\n{text}")
        if summary:
            st.subheader("📄 Summary")
            st.write(summary)

    elif file_name.endswith((".png", ".jpg", ".jpeg")):
        img = Image.open(uploaded_file)
        st.image(img, caption="Uploaded Image", use_column_width=True)
        summary = call_grok("Explain this image in simple terms for a student.")
        if summary:
            st.subheader("🖼 Explanation")
            st.write(summary)