import streamlit as st
import datetime
from openai import OpenAI
import os
from prompts import PROMPT_FOR_MANDATE

OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
client = OpenAI(api_key=OPENAI_API_KEY)
st.set_page_config(page_title="AI Trip Planner")

if 'clicked' not in st.session_state:
    st.session_state.clicked = False


def click_button():
    st.session_state.clicked = True


with st.sidebar:
    st.title("Your AI Trip Planner!")
    where_to = st.text_area("Where To?")
    number_of_days = st.text_area("How many days?")
    itinerary_type = st.multiselect(
            "Itinerary Type(s):",
            (
                "Standard",
                "Adventure",
                "Bachelor Party",
                "Bachelor Party",
                "Bicycling",
                "City Experiences",
                "Cultural Immersion",
                "Dog Friendly",
                "Eco Friendly",
                "Educational",
                "Family Friendly",
                "Foodie",
                "Honeymoon",
                "Luxury",
                "Motorcycle Riding",
                "Nature",
                "Relaxation",
                "Rail/Train Trip",
                "Road Trip",
                "Shoestring Budget",
                "Volunteering",
            )
    )
    when_your_trip_start = st.date_input("When your trip start?", datetime.date(2024, 7, 6))
    when_your_trip_end = st.date_input("When your trip end?", datetime.date(2024, 7, 8))
    travel_preference = st.text_area("Any travel preference?")
    st.button("Generate Trip Plan", on_click=click_button)

if st.session_state.clicked:
    if where_to and number_of_days and itinerary_type:
        message_placeholder = st.empty()
        full_response = ""
        prompt = f"""PLACE:```{where_to}``` , NUMBER_OF_DAYS:```{number_of_days}```, TRIP_TYPE: ```{itinerary_type}```""" + \
                 PROMPT_FOR_MANDATE
        for response in client.chat.completions.create(
                model="gpt-4",
                messages=[{'role': 'user', 'content': prompt}],
                stream=True,
        ):
            if response.choices[0].delta.content is not None:
                full_response += response.choices[0].delta.content
                message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)
    elif where_to and number_of_days and itinerary_type or when_your_trip_start or when_your_trip_end or travel_preference:
        message_placeholder = st.empty()
        full_response = ""
        prompt = f"""PLACE:```{where_to}``` , NUMBER_OF_DAYS:```{number_of_days}```, TRIP_TYPE: ```{itinerary_type}```,
TRIP_START: ```{when_your_trip_start}```, TRIP_END:```{when_your_trip_end}```, TRAVEL_PREFERENCE:```{travel_preference}```""" + \
                 PROMPT_FOR_MANDATE
        for response in client.chat.completions.create(
                model="gpt-4",
                messages=[{'role': 'user', 'content': prompt}],
                stream=True,
        ):
            if response.choices[0].delta.content is not None:
                full_response += response.choices[0].delta.content
                message_placeholder.markdown(full_response + "▌")
        message_placeholder.markdown(full_response)