import os
import streamlit as st
from langchain.chains import ConversationChain
from langchain_openai import ChatOpenAI
from langchain.memory import ConversationBufferMemory
from streamlit_extras.add_vertical_space import add_vertical_space
from streamlit_extras.chat_elements import message

# Set OpenAI API Key
os.environ["OPENAI_API_KEY"] = ""

# Initialize the chatbot
@st.cache_resource
def init_chatbot():
    memory = ConversationBufferMemory()
    chatbot = ConversationChain(
        llm=ChatOpenAI(model="gpt-4o-mini"),
        memory=memory,
        verbose=False
    )
    return chatbot

chatbot = init_chatbot()

# Custom Styling
st.markdown("""
    <style>
        body {
            background-color: #f5f5f5;
        }
        .chat-container {
            background-color: #ffffff;
            padding: 20px;
            border-radius: 10px;
            box-shadow: 0px 4px 6px rgba(0,0,0,0.1);
        }
    </style>
""", unsafe_allow_html=True)

# Sidebar - Settings
st.sidebar.title("⚙️ Settings")
model_choice = st.sidebar.radio("Select Model", ("gpt-4o-mini", "gpt-4", "gpt-3.5-turbo"))

# Update model based on user selection
if model_choice:
    chatbot.llm = ChatOpenAI(model=model_choice)

# Title and Description
st.title("💬 LangChain AI Chatbot")
st.write("### Hi, I'm a chatbot built with Langchain powered by GPT. How can I assist you today?")

# Chat history
if "chat_history" not in st.session_state:
    st.session_state.chat_history = []

# User Input
user_input = st.text_input("You:", placeholder="Ask me anything...")

# Process input
if user_input:
    with st.spinner("Thinking..."):
        response = chatbot.run(user_input)
        st.session_state.chat_history.append(("user", user_input))
        st.session_state.chat_history.append(("bot", response))

# Display chat history
st.write("### 🗨️ Conversation")
chat_container = st.container()

with chat_container:
    for role, text in st.session_state.chat_history:
        if role == "user":
            message(text, is_user=True, avatar_style="thumbs")
        else:
            message(text, is_user=False, avatar_style="bottts")

# Add some spacing
add_vertical_space(2)

# Collapsible Chat History
with st.expander("📜 Chat History"):
    for role, text in st.session_state.chat_history:
        st.write(f"**{role.capitalize()}**: {text}")

# Footer
st.markdown("---")
st.markdown("Developed with ❤️ using Streamlit & LangChain")