import streamlit as st
import requests
import os

# UI Setup
st.set_page_config(page_title="Failure Diagnosis Bot", page_icon="🔧")
st.title("🔧 Failure Diagnosis Bot")
st.subheader("Diagnose mechanical issues based on symptoms.")

# API Config
GROQ_API_KEY = st.secrets.get("GROQ_API_KEY") or os.getenv("GROQ_API_KEY")
GROQ_MODEL = "llama-3.3-70b-versatile"  # Updated model

# Input from user
user_input = st.text_input("Describe the symptom observed in your machine:", placeholder="e.g., Grinding noise in gearbox...")

# Prompt template
def build_prompt(symptom):
    return f"""
You are a mechanical engineering diagnostic chatbot. A user will describe symptoms observed in a mechanical system. Your job is to:

1. Identify possible fault(s)
2. List the probable cause(s)
3. Suggest remedies or corrective actions
4. List the tools needed for diagnosis or repair

Respond using this structure:

**Symptom:**  
{symptom}

**Possible Fault(s):**
- Fault 1
- Fault 2

**Cause(s):**
- Cause 1
- Cause 2

**Remedy/Recommended Action:**
- Step 1
- Step 2

**Tools Needed:**
- Tool 1
- Tool 2
"""

# Groq API call
def get_diagnosis(prompt):
    headers = {
        "Authorization": f"Bearer {GROQ_API_KEY}",
        "Content-Type": "application/json"
    }
    payload = {
        "model": GROQ_MODEL,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
    response = requests.post("https://api.groq.com/openai/v1/chat/completions", headers=headers, json=payload)
    if response.status_code == 200:
        return response.json()['choices'][0]['message']['content']
    else:
        return f"❌ Error: {response.text}"

# Run diagnosis
if user_input and GROQ_API_KEY:
    with st.spinner("Analyzing..."):
        prompt = build_prompt(user_input)
        result = get_diagnosis(prompt)
        st.markdown(result)
elif not GROQ_API_KEY:
    st.warning("Please set your GROQ_API_KEY as a secret or environment variable to use this app.")