import os
import streamlit as st
from dotenv import load_dotenv
import google.generativeai as gen_ai

# Load environment variables
load_dotenv()

# Configure Streamlit page settings
st.set_page_config(
    page_title="Smart Waste Management System",
    page_icon=":recycle:",  # Favicon emoji
    layout="centered",  # Page layout option
)

# Retrieve the Google API key from the environment
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")

# Check if the API key is loaded
if not GOOGLE_API_KEY:
    st.error("API key not found! Please set the GOOGLE_API_KEY in your .env file.")
    st.stop()

# Configure the Generative AI model
try:
    gen_ai.configure(api_key=GOOGLE_API_KEY)
    model = gen_ai.GenerativeModel("gemini-pro")
except Exception as e:
    st.error(f"Error initializing the Gemini-Pro model: {e}")
    st.stop()

# Function to translate roles between Gemini-Pro and Streamlit terminology
def translate_role_for_streamlit(user_role):
    return "assistant" if user_role == "model" else user_role

# Initialize the chat session if not already present in session state
if "chat_session" not in st.session_state:
    try:
        st.session_state.chat_session = model.start_chat(history=[])
    except Exception as e:
        st.error(f"Error initializing chat session: {e}")
        st.stop()

# Display the app's title
st.title("♻️ Smart Waste Management System")

# Introduction and instructions
st.markdown(
    """
    Welcome to the Smart Waste Management System! This tool helps citizens, municipal workers, 
    recycling companies, and biogas plants collaborate efficiently for effective waste management.

    **Key Features:**
    - **Citizen Role:** Report waste collection issues and track garbage pickup.
    - **Municipal Workers:** Manage schedules and coordinate garbage segregation.
    - **Recycling Companies:** View and respond to requests for plastic waste.
    - **Biogas Plants:** Manage bio-waste for energy production.
    """
)

# Input fields for user role and message
user_role = st.selectbox("Select Your Role:", ["Citizen", "Municipal Worker", "Recycling Company", "Biogas Plant"])

user_prompt = st.chat_input(f"[{user_role}] Enter your query or task...")
if user_prompt:
    # Add the user's message to the chat and display it
    st.chat_message("user").markdown(f"**{user_role}:** {user_prompt}")

    # Customize the prompt based on the user role
    role_specific_prompt = (
        f"You are assisting a {user_role} in a waste management system. The user says: {user_prompt}"
    )

    # Send the prompt to Gemini-Pro and get the response
    try:
        gemini_response = st.session_state.chat_session.send_message(role_specific_prompt)
        # Display Gemini-Pro's response
        with st.chat_message("assistant"):
            st.markdown(gemini_response.text)
    except Exception as e:
        st.error(f"Error processing your message: {e}")

# Add a sidebar for additional information
st.sidebar.title("About")
st.sidebar.markdown(
    """
    The **Smart Waste Management System** aims to:
    - Promote efficient segregation and recycling of waste.
    - Facilitate collaboration between stakeholders.
    - Reduce environmental impact and enhance sustainability.

    **Need Help?** Use the chat to ask questions or provide feedback.
    """
)