import streamlit as st
import pytesseract
from PIL import Image
import fitz  # PyMuPDF
import io
import requests

# --- Config ---
API_KEY = "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"  
API_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL = "mistralai/mistral-7b-instruct"

st.set_page_config(page_title="🧪 AI Science Lab Assistant", layout="centered")
st.title("🧪 AI Science Lab Assistant")

st.markdown("""
This tool helps students evaluate their lab reports by:
- Extracting text from uploaded **PDFs** or **images**
- Automatically **checking for completeness** (Objective, Hypothesis, etc.)
- Giving **improvement tips** 💡
- Letting you **ask questions** about your report 📘
""")

# --- Upload ---
uploaded_file = st.file_uploader("📸 Upload Your Lab Report Image or PDF", type=["jpg", "jpeg", "png", "pdf"])

lab_text = ""
if uploaded_file:
    file_bytes = uploaded_file.read()
    file_ext = uploaded_file.name.split(".")[-1].lower()

    if file_ext == "pdf":
        doc = fitz.open(stream=file_bytes, filetype="pdf")
        for page in doc:
            lab_text += page.get_text()
    else:
        image = Image.open(io.BytesIO(file_bytes))
        lab_text = pytesseract.image_to_string(image)

    st.subheader("📄 Extracted Lab Report Text:")
    st.text_area("", lab_text, height=300)

    # -- AI Evaluation Prompt --
    full_prompt = f"""Here is a science lab report:
{lab_text}

Please evaluate the report based on the following:
- Does it include sections like Objective, Hypothesis, Procedure, Observation, and Conclusion?
- Point out any missing or incomplete sections.
- Give feedback or improvement suggestions like: "Try writing a more detailed observation."
- Grade it roughly on a scale from 1 to 10 for completeness.

Respond clearly:
"""

    def query_ai(prompt):
        headers = {
            "Authorization": f"Bearer {API_KEY}",
            "Content-Type": "application/json"
        }
        payload = {
            "model": MODEL,
            "messages": [
                {"role": "system", "content": "You are a helpful science teacher."},
                {"role": "user", "content": prompt}
            ]
        }
        response = requests.post(API_URL, headers=headers, json=payload)
        return response.json()['choices'][0]['message']['content']

    if st.button("🧠 Evaluate My Lab Report"):
        with st.spinner("Analyzing report with AI..."):
            result = query_ai(full_prompt)
        st.success("✅ Evaluation Complete")
        st.markdown("### 📊 AI Feedback:")
        st.markdown(result)

    # Ask follow-up questions
    st.subheader("🤔 Ask AI About Your Report")
    user_question = st.text_input("Type your question")
    if st.button("🔍 Ask") and user_question:
        with st.spinner("Thinking..."):
            followup_prompt = f"Here is the lab report:
{lab_text}

Now answer this question about it:
{user_question}"
            followup_response = query_ai(followup_prompt)
        st.markdown("### 💬 AI Answer:")
        st.markdown(followup_response)