# app.py
import streamlit as st
import fitz  # PyMuPDF
import io
import requests
import re
import os
from fpdf import FPDF
from datetime import datetime
from PIL import Image
import base64
import json

# --- Config ---
API_URL = "https://openrouter.ai/api/v1/chat/completions"
MODEL = "mistralai/mistral-7b-instruct"

# Retrieve API key from environment variable (set in Hugging Face secrets)
OPENROUTER_API_KEY = os.getenv("OPENROUTER_API_KEY")

# Set page config
st.set_page_config(
    page_title="🔬 Science Lab Assistant", 
    layout="centered",
    page_icon="🔬",
    initial_sidebar_state="expanded"
)

# Custom CSS for styling
st.markdown("""
    <style>
    /* (Keep all your existing CSS styles here) */
    </style>
    """, unsafe_allow_html=True)

# Header
st.markdown('<p class="header">🔬 Science Lab Assistant</p>', unsafe_allow_html=True)

# Introduction
st.markdown("""
<div style="text-align: center; margin-bottom: 30px;">
    <p style="font-size: 18px;">Your all-in-one science companion! Design experiments, generate reports, 
    and get AI-powered feedback on your lab work.</p>
</div>
""", unsafe_allow_html=True)

# API Key Check
if not OPENROUTER_API_KEY:
    st.error("""
    **API Key Not Configured!**
    
    Please add your OpenRouter API key to Hugging Face Spaces secrets:
    1. Go to your Space settings
    2. Select "Variables and secrets" 
    3. Add a secret named: `OPENROUTER_API_KEY`
    4. Set its value to your actual API key
    5. Redeploy the space
    
    Without this key, the AI features won't work.
    """)
    st.stop()

# Experiment templates
experiments = {
    "Vinegar + Baking Soda": {
        "hypothesis": "Mixing vinegar and baking soda will produce bubbles due to a chemical reaction.",
        "concept": "Acid-base reaction producing carbon dioxide."
    },
    "Floating Egg": {
        "hypothesis": "An egg will float in salt water but sink in plain water.",
        "concept": "Density difference between saltwater and freshwater."
    },
    "Lemon Battery": {
        "hypothesis": "A lemon can produce electricity to power a small LED.",
        "concept": "Chemical energy conversion to electrical energy."
    }
}

# AI Query Function
def query_ai(prompt):
    headers = {
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": MODEL,
        "messages": [
            {"role": "system", "content": "You are a helpful science teacher providing detailed explanations."},
            {"role": "user", "content": prompt}
        ],
        "temperature": 0.7
    }
    try:
        response = requests.post(API_URL, headers=headers, json=payload, timeout=120)
        response.raise_for_status()
        return response.json()['choices'][0]['message']['content']
    except requests.exceptions.HTTPError as err:
        st.error(f"API Error: {err.response.status_code} - {err.response.text}")
        return None
    except Exception as e:
        st.error(f"Error connecting to AI service: {str(e)}")
        return None

# Navigation
app_mode = st.radio("Choose Mode:", ["🧪 Experiment Assistant", "📝 Lab Report Analyzer"], 
                    horizontal=True, label_visibility="collapsed")

# Sidebar
with st.sidebar:
    st.markdown("### 🧪 Experiment Templates")
    st.caption("Quickly start with these pre-defined experiments:")
    
    selected_exp = st.selectbox("Choose an experiment template:", 
                               list(experiments.keys()) + ["Custom Experiment"])
    
    st.markdown("---")
    st.markdown("### 📘 Science Glossary Helper")
    term = st.text_input("Enter a science term (e.g., osmosis, catalyst)")
    if term:
        with st.spinner("Looking up term..."):
            ai_response = query_ai(f"Explain the term '{term}' in simple words for a student.")
        if ai_response:
            st.markdown(f"<div class='concept-box'>{ai_response}</div>", unsafe_allow_html=True)

# --- Experiment Assistant Section ---
if app_mode == "🧪 Experiment Assistant":
    # (Keep all your experiment assistant code here unchanged)
    pass

# --- Lab Report Analyzer Section ---
else:
    # --- File Upload ---
    st.markdown('<p class="subheader">📤 Upload Your Lab Report</p>', unsafe_allow_html=True)
    uploaded_file = st.file_uploader("Upload PDF only (image support coming soon)", 
                                     type=["pdf"],
                                     label_visibility="collapsed")

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

        if file_ext == "pdf":
            try:
                doc = fitz.open(stream=file_bytes, filetype="pdf")
                for page in doc:
                    lab_text += page.get_text()
                st.success("✅ PDF text extracted successfully!")
            except Exception as e:
                st.error(f"Error reading PDF: {str(e)}")
        else:
            st.warning("Image upload is temporarily disabled. Please upload PDF files only.")
            st.info("Tip: Convert images to PDF using free online tools")

        # Allow text editing
        if lab_text:
            st.markdown('<p class="subheader">✍️ Extracted Text</p>', unsafe_allow_html=True)
            st.caption("Review and edit the extracted text if needed before analysis")
            lab_text = st.text_area("", lab_text, height=300, label_visibility="collapsed")

    # --- AI Evaluation ---
    if lab_text.strip():
        # -- AI Evaluation Prompt --
        full_prompt = f"""You are a science teacher evaluating a student's lab report. Please provide a comprehensive analysis:

        Lab Report:
        {lab_text}

        Evaluation Guidelines:
        1. **Section Check**: Identify which of these sections are present and which are missing:
            - Title
            - Objective
            - Hypothesis
            - Materials
            - Procedure
            - Observations
            - Results
            - Conclusion
            - References
        
        2. **Completeness Score**: 
            - Assign a numerical score from 1-10 based on completeness
            - Justify the score based on missing sections and content quality
        
        3. **Improvement Tips**:
            - For each missing section, explain why it's important
            - Provide specific suggestions for improvement (e.g., "Try writing a more detailed observation section by including quantitative data")
            - Highlight any sections that need more detail or clarity
        
        4. **Structure Response**:
            - Start with: "### Missing Sections:"
            - Then: "### Completeness Score: X/10"
            - Then: "### Improvement Tips:"
            - Finally: "### Detailed Feedback:"
        
        Be concise but thorough in your analysis.
        """

        if st.button("🧪 Analyze Report", use_container_width=True):
            with st.spinner("🔍 Analyzing report with AI. This may take 20-30 seconds..."):
                result = query_ai(full_prompt)
                
            if result:
                st.success("✅ Analysis Complete!")
                st.balloons()
                
                # (Keep all your analysis result display code here unchanged)
                pass
                
        # --- Question Answering Section ---
        st.markdown("---")
        st.markdown('<p class="subheader">❓ Ask About Your Report</p>', unsafe_allow_html=True)
        
        col1, col2 = st.columns([3, 1])
        with col1:
            user_question = st.text_input("Ask a question about your lab report", 
                                         placeholder="e.g., How can I improve my hypothesis?")
        with col2:
            st.markdown("<div style='height: 28px;'></div>", unsafe_allow_html=True)
            ask_button = st.button("🔍 Ask Question", use_container_width=True)
        
        if (ask_button or user_question) and user_question.strip():
            with st.spinner("Thinking..."):
                followup_prompt = f"""Lab Report:
                {lab_text}
                
                Question: {user_question}
                
                Answer the question based on the lab report. If the question can't be answered from the report, 
                suggest what information the student should add to answer it.
                """
                followup_response = query_ai(followup_prompt)
            
            if followup_response:
                st.markdown("### 💬 AI Response")
                st.markdown(f'<div class="tip-box">{followup_response}</div>', unsafe_allow_html=True)
    else:
        # Show sample report if no file uploaded
        st.markdown("---")
        st.markdown('<p class="subheader">📝 Sample Lab Report</p>', unsafe_allow_html=True)
        st.markdown("""
        **Title:** Effect of Temperature on Enzyme Activity  
        **Objective:** To investigate how temperature affects catalase enzyme activity  
        **Hypothesis:** Enzyme activity will increase with temperature up to 37°C, then decrease  
        **Materials:** Test tubes, hydrogen peroxide, liver extract, thermometer  
        **Procedure:**  
        1. Prepare test tubes at 5 different temperatures  
        2. Add equal amounts of hydrogen peroxide and liver extract  
        3. Measure oxygen production  
        **Observations:** More bubbles at 37°C compared to lower or higher temperatures  
        **Conclusion:** Enzyme activity peaks at body temperature  
        """)
        
        st.info("👆 Upload your own lab report to get a personalized analysis!")

# Footer
st.markdown("---")
st.markdown('<div class="footer">🔬 Science Lab Assistant | Made for Students & Educators</div>', unsafe_allow_html=True)