import streamlit as st
from huggingface_hub import InferenceClient
import os

# ---------------- PAGE CONFIG ----------------
st.set_page_config(page_title="Smart Waste Segregation Advisor")

st.title("♻️ AI-Powered Smart Waste Segregation Advisor")
st.subheader("SDG 12: Responsible Consumption and Production")

st.markdown("""
This tool helps users identify waste categories and provides guidance
for responsible disposal using AI-based reasoning powered by IBM Granite.
""")

# ---------------- USER INPUTS ----------------
uploaded_image = st.file_uploader(
    "Upload an image of the waste item (optional)",
    type=["jpg", "png", "jpeg"]
)

user_text = st.text_input("Optional: Describe the waste item")

st.divider()

# ---------------- ANALYZE BUTTON ----------------
if st.button("Analyze Waste"):

    # Simulated image interpretation (Option 1 design choice)
    if uploaded_image:
        image_description = "An image showing a waste item uploaded by the user"
    else:
        image_description = "No image provided"

    # Initialize IBM Granite via Hugging Face Inference API
    client = InferenceClient(
        model="ibm-granite/granite-4.0-mini",
        token=os.getenv("HF_TOKEN")
    )

    prompt = f"""
You are an AI-powered Smart Waste Segregation Advisor.

Tasks:
1. Identify the waste item.
2. Classify it as one of the following:
   - Wet Waste
   - Dry Waste
   - E-Waste
3. Explain briefly why it belongs to this category.
4. Suggest the correct and safe disposal method.
5. Explain the environmental impact if disposed of incorrectly.

User Input:
Image description: {image_description}
User text: {user_text}
"""

    with st.spinner("Analyzing using IBM Granite AI..."):
        response = client.chat_completion(
            messages=[
                {
                    "role": "system",
                    "content": "You are an AI assistant focused on sustainability and responsible waste management."
                },
                {
                    "role": "user",
                    "content": prompt
                }
            ],
            max_tokens=250,
            temperature=0.7
        )

    ai_output = response.choices[0].message["content"]

    st.markdown("### 🧠 AI Analysis Result")
    st.write(ai_output)