import streamlit as st
import os
from groq import Groq

# Get API Key
api_key = os.getenv("GROQ_API_KEY")

# Initialize Groq client
client = Groq(api_key=api_key) if api_key else None

# Streamlit UI
st.set_page_config(page_title="🧪 ChemEng AI Assistant", page_icon="🧪")
st.title("🧪 Chemical Engineering Equipment Sizing Assistant")

st.markdown("""
Ask any question about chemical process equipment sizing and get AI-powered answers based on domain knowledge.

**Try examples like:**
- `What size of distillation column is used for 5000 kg/hr ethanol separation?`
- `Size a shell-and-tube heat exchanger for 1 MW heat duty.`
- `Suggest the volume of a CSTR with 1-hour residence time at 2000 L/hr.`
""")

# User Input
query = st.text_area("💬 Enter your question:", height=150)

# Model Selector (Only supported by Groq)
model = st.selectbox("🔍 Choose a model:", [
    "llama-3.1-8b-instant",
    "llama-3.3-70b-versatile",
    "openai/gpt-oss-120b"
])

# Process the query
if st.button("🚀 Get Response"):
    if not api_key:
        st.error("🚨 Groq API Key not found. Set the GROQ_API_KEY in your environment.")
    elif not query.strip():
        st.warning("⚠️ Please enter a question.")
    else:
        with st.spinner("Generating response..."):
            try:
                chat_completion = client.chat.completions.create(
                    model=model,
                    messages=[
                        {
                            "role": "system",
                            "content": (
                                "You are a senior chemical engineer. Answer user questions about equipment sizing, "
                                "design rules, and heuristics clearly and concisely. Provide real-world estimates based on "
                                "chemical engineering practices."
                            )
                        },
                        {
                            "role": "user",
                            "content": query
                        }
                    ]
                )
                response = chat_completion.choices[0].message.content
                st.success("✅ Answer:")
                st.markdown(response)

            except Exception as e:
                st.error(f"❌ Error: {str(e)}")