import os
import streamlit as st
from langchain_openai import ChatOpenAI
from langchain.vectorstores import FAISS
import pandas as pd

# Environment setup
os.environ["OPENAI_API_KEY"] = os.getenv("OPENAI_API_KEY")

# Streamlit page setup
st.set_page_config(
    page_title="Airtel AI Agent System",
    page_icon="🤖",
    layout="wide",
    initial_sidebar_state="expanded"
)

# Sidebar setup
with st.sidebar:
    airtel_logo = "https://1000logos.net/wp-content/uploads/2023/06/Airtel-logo.jpg"  # Airtel logo URL
    st.image(airtel_logo, use_column_width=True)
    st.markdown("### Welcome to Airtel AI Help Bot")
    st.markdown("**Instructions:**")
    st.markdown("- Use the chat interface to ask questions about Airtel services.")
    st.markdown("- You can also initiate specific tasks like SIM swap or plan details.")
    st.markdown("Feel free to ask any queries you have!")

# Apply custom CSS for the red-and-white theme
st.markdown(
    """
    <style>
    .stApp {
        background-color: white;
    }
    header {
        background-color: #e60000;
    }
    .stButton>button {
        background-color: #e60000;
        color: white;
    }
    .stTextInput>div>input {
        background-color: #ffffff;
        border: 2px solid #e60000;
        color: black;
    }
    .css-1mgb8k9.e1fqkh3o3 {
        background-color: white;
    }
    .css-18ni7ap.e8zbici0 {
        background-color: white;
    }
    footer {
        background-color: white;
    }
    </style>
    """,
    unsafe_allow_html=True,
)

# Page title
st.title("Airtel AI Agent System Demo")

# Initialize session state to store chat history
if "messages" not in st.session_state:
    st.session_state.messages = []

# User input for chat interface
user_input = st.text_input("Enter your message:")

# Initialize agents' tasks
def conversation_agent(state):
    """Handles general conversations with the user."""
    llm = ChatOpenAI(model_name="gpt-3.5-turbo")
    
    try:
        response = llm.invoke(state["question"])
        # Ensure that response is not None
        if response and response.content:
            parsed_response = response.content  # Directly access content from AIMessage object
        else:
            parsed_response = "I'm sorry, I couldn't generate a response."
    except Exception as e:
        parsed_response = f"An error occurred: {e}"
    
    return {"question": state["question"], "response": parsed_response, "intent": "conversation"}

def website_query_agent(state):
    """Handles queries by retrieving relevant data from Airtel website stored in vector DB."""
    vectorstore = FAISS.load_local("path_to_faiss_index")  # Assuming Airtel data is preloaded in FAISS
    docs = vectorstore.similarity_search(state["question"])
    response = "\n".join([doc.page_content for doc in docs[:3]])  # Return top 3 relevant docs
    return {"question": state["question"], "response": response, "intent": "website_query"}

def fallback_agent(state):
    """Returns a fallback message for unrelated queries."""
    return {"question": state["question"], "response": "Sorry, I don't have an answer for that.", "intent": "fallback"}

def sim_swap_agent(state):
    """Handles SIM swap workflow by triggering a form to collect user details."""
    st.write("---Sim Swap Request---")
    
    # File path to store the data
    file_path = "sim_swap_requests.xlsx"
    
    # Form for collecting user details
    with st.form("sim_swap_form"):
        name = st.text_input("Enter your name:")
        phone_number = st.text_input("Enter your phone number:")
        submit = st.form_submit_button("Submit")

        if submit:
            new_data = {"Name": [name], "Phone Number": [phone_number]}
            df_new = pd.DataFrame(new_data)
            
            # Check if the Excel file already exists
            if os.path.exists(file_path):
                df_existing = pd.read_excel(file_path)
                df_updated = pd.concat([df_existing, df_new], ignore_index=True)
            else:
                df_updated = df_new
            
            # Save the updated DataFrame to Excel
            df_updated.to_excel(file_path, index=False)
            return {"question": state["question"], "response": f"SIM Swap request submitted for {name} with phone number {phone_number}.", "intent": "sim_swap"}
    
    return {"question": state["question"], "response": "Please fill in your details for SIM swap.", "intent": "sim_swap"}

def plan_details_agent(state):
    """Handles plan details workflow by collecting phone number and returning a standard plan."""
    with st.form("plan_details_form"):
        phone_number = st.text_input("Enter your phone number:")
        submit = st.form_submit_button("Submit")

        if submit:
            return {"question": state["question"], "response": f"Plan details for {phone_number}: You are on a standard plan of 999 INR/month.", "intent": "plan_details"}
    
    return {"question": state["question"], "response": "Please enter your phone number to check plan details.", "intent": "plan_details"}

# Decision logic
def decide_agent(state):
    question = state["question"].lower()
    if "sim swap" in question or "swap sim" in question:
        return "sim_swap"
    elif "plan details" in question or "active plan" in question:
        return "plan_details"
    elif "airtel" in question:
        return "conversation"
    elif "website" in question or "airtel website" in question:
        return "website_query"
    else:
        return "fallback"

# Workflow logic
workflow = {
    "conversation": conversation_agent,
    "website_query": website_query_agent,
    "sim_swap": sim_swap_agent,
    "plan_details": plan_details_agent,
    "fallback": fallback_agent,
}

# Function to handle user input and generate responses
def handle_query(query):
    state = {"question": query, "response": "", "intent": ""}
    next_agent = decide_agent(state)
    if next_agent in workflow:
        agent = workflow[next_agent]
        output = agent(state)
        return output["response"]

# Chat Interface
if user_input:
    st.session_state.messages.append({"role": "user", "content": user_input})
    bot_response = handle_query(user_input)
    st.session_state.messages.append({"role": "bot", "content": bot_response})

# Display chat messages in reverse order (latest message on top)
for message in reversed(st.session_state.messages):
    if message["role"] == "user":
        st.chat_message("user").markdown(message["content"])
    else:
        st.chat_message("bot").markdown(message["content"])