# Imports & Config
import streamlit as st
from openai import OpenAI
from llama_index.core import VectorStoreIndex, download_loader, SimpleDirectoryReader
from llama_index.readers.web import BeautifulSoupWebReader
from bs4 import BeautifulSoup
import requests
import openai

# Title
st.title("✈️ Airline Compensation Bot 💬")

# Introduction
st.markdown("**Flight delayed or cancelled yet again?**")
st.markdown("Don’t waste time digging through complicated airline policies! **Just tell us your airline, delay/cancellation reason, and any other relevant details** and our chatbot will instantly tell you what compensation you're entitled to - by law and by airline policy.")
st.divider()

#OpenAI API Key
with st.sidebar:
    openai_api_key = st.text_input("Enter your OpenAI API Key", type="password")
    "[Get an OpenAI API key](https://platform.openai.com/account/api-keys)"
    "[View the source code](https://github.com/streamlit/llm-examples/blob/main/Chatbot.py)"
    st.divider()
    st.markdown(
            """
            **Under the hood:**

            Large language models like ChatGPT excel at general queries but can falter on specialized topics. This tool uses LlamaIndex for Retrieval Augmented Generation (RAG), pulling data directly from trusted sources like the US Department of Transportation (DOT) and airline customer service documents. This ensures a more reliable chatbot that provides accurate information tailored to your specific needs, rather than relying on broad internet data.

            """
        )

# Check if an API key is provided
if not openai_api_key:
    st.warning("Please enter your OpenAI API to use this app 🔑")
    st.stop()  # Stop the app until an API key is provided

# Initialize OpenAI API Client with user-provided API key
openai.api_key = openai_api_key

# Load data
loader = BeautifulSoupWebReader()
documents = loader.load_data(urls=[
"https://content.spirit.com/Shared/en-us/Documents/Contract_of_Carriage.pdf",
"https://secure.dot.gov/air-travel-complaint",
"https://www.ecfr.gov/current/title-14/chapter-II/subchapter-A/part-259",
"https://www.ecfr.gov/current/title-14/chapter-II/subchapter-A/part-259#259.5",
"https://www.federalregister.gov/documents/2024/04/26/2024-07177/refunds-and-other-consumer-protections",
"https://www.federalregister.gov/documents/2024/08/12/2024-17602/refunds-and-other-consumer-protections-2024-faa-reauthorization",
"https://www.flyfrontier.com/legal/customer-service-plan?mobile=true",
"https://www.hawaiianairlines.com/about-us/customer-service-plan",
"https://www.jetblue.com/customer-assurance/customer-service-plan",
"https://www.jetblue.com/legal/customer-service-plan",
"https://www.reginfo.gov/public/do/eAgendaViewRule?pubId=201910&RIN=2105-AE57",
"https://www.reginfo.gov/public/do/eAgendaViewRule?pubId=202110&RIN=2105-AF04",
"https://www.southwest.com/assets/pdfs/corporate-commitments/customer-service-plan.pdf?clk=7396032",
"https://www.southwest.com/swa-resources/pdfs/corporate-commitments/contract-of-carriage.pdf?clk=CSP_Form",
"https://www.transportation.gov/airconsumer",
"https://www.transportation.gov/airconsumer/air-travel-tips#FAQ",
"https://www.transportation.gov/airconsumer/air-travelers-tell-it-judge",
"https://www.transportation.gov/airconsumer/airline-consumer-contacts",
"https://www.transportation.gov/airconsumer/fly-rights",
"https://www.transportation.gov/individuals/aviation-consumer-protection/bumping-oversales",
"https://www.transportation.gov/individuals/aviation-consumer-protection/refunds",
"https://www.transportation.gov/individuals/aviation-consumer-protection/tarmac-delays",
"https://www.transportation.gov/lost-delayed-or-damaged-baggage",
"https://www.transportation.gov/resources/individuals/aviation-consumer-protection/airline-cancellation-delay-dashboard-html",
    ])

# RAG
index = VectorStoreIndex.from_documents(documents)

def response_generator(query):
    try:
        # Create an index from the documents
        query_engine = index.as_query_engine()
        response = query_engine.query(query)

    except Exception as e:
        # Log or handle the exception
        response = f"An error occurred: {e}"

    return response


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

# Display chat messages from history on app rerun
for message in st.session_state.messages:
    with st.chat_message(message["role"]):
        st.markdown(message["content"])

# Accept user input
if prompt := st.chat_input("How can I help you?"):
    # Add user message to chat history
    st.session_state.messages.append({"role": "user", "content": prompt})
    
    # Display user message in chat message container
    with st.chat_message("user"):
        st.markdown(prompt)
    
    # Generate and display assistant response
    # st.write('Before Response generator')
    response = response_generator(prompt)
  
    # Display assistant response in chat message container
    with st.chat_message("assistant"):
        st.markdown(response)
    
    # Add assistant response to chat history
    st.session_state.messages.append({"role": "assistant", "content": response})