import streamlit as st
import openai
from dotenv import load_dotenv
import os

# Load environment variables
load_dotenv()

# Configure OpenAI API
client = openai.OpenAI(api_key=os.getenv("OPENAI_API_KEY"))

# System prompt template
SYSTEM_PROMPT = """You are an AI assistant for the University of Engineering and Technology (UET). 
Your role is to provide accurate and helpful information about {university}.
You should:
1. Provide information about academic programs, admission requirements, and campus facilities
2. Answer questions about university policies, procedures, and important dates
3. Guide students about registration, courses, and academic matters
4. Share information about research opportunities and faculty
5. Direct users to official university resources when needed

Always maintain a professional and helpful tone. If you're unsure about any information, 
acknowledge the limitation and suggest contacting the university directly."""

def get_chatbot_response(messages):
    """Get response from OpenAI API"""
    try:
        response = client.chat.completions.create(
            model="gpt-3.5-turbo", 
            messages=messages,
            temperature=0.7,
            max_tokens=500
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

def main():
    st.title("UET Chatbot")
    st.write("Welcome to the UET Chatbot! Ask me anything about UET Peshawar or UET Mardan.")

    # Dropdown to select university
    university = st.selectbox("Select University", ["UET Peshawar", "UET Mardan"])

    # Initialize chat history
    if "messages" not in st.session_state:
        st.session_state.messages = [{"role": "system", "content": SYSTEM_PROMPT.format(university=university)}]

    # Update system prompt if university changes
    if st.session_state.messages[0]["content"] != SYSTEM_PROMPT.format(university=university):
        st.session_state.messages = [{"role": "system", "content": SYSTEM_PROMPT.format(university=university)}]

    # Display chat history
    for message in st.session_state.messages:
        if message["role"] != "system":
            with st.chat_message(message["role"]):
                st.write(message["content"])

    # Chat input
    if prompt := st.chat_input("What would you like to know?"):
        # Add user message to chat history
        st.session_state.messages.append({"role": "user", "content": prompt})
        
        # Display user message
        with st.chat_message("user"):
            st.write(prompt)

        # Get chatbot response
        with st.chat_message("assistant"):
            response = get_chatbot_response(st.session_state.messages)
            st.write(response)
            
        # Add assistant response to chat history
        st.session_state.messages.append({"role": "assistant", "content": response})

if __name__ == "__main__":
    main()