# test file for gradio

# import openai
import os
from openai import OpenAI
import gradio as gr

client = OpenAI()

def travel_agent(input_str):
    completion = client.chat.completions.create(
        model="gpt-4",
        messages=[
            {"role": "system",
             "content": "You will be provided with constraints, and your task is to provide a travel itinerary which satisfies these constraints. Suggest restaurants and hotels as appropriate."},
            {"role": "user", "content": input_str}
        ]
    )
    return (completion.choices[0].message.content)


def get_itinerary(destination, duration, events, cost, individuals, constraints, food, mobility, transport):
    # create prompt
    input_str = ('I want to go to ' + destination + ' for ' + str(duration) + ' days.' +
                 '. I want the itinerary to include the following activities:' + ','.join(events) +
                 '. On average, I want to spend ' + str(cost) + ' dollars per day for a group of ' + str(individuals) +
                 '. Some additional constraints are: ' + constraints + '.')

    personalize = {
        "Preferred food": food,
        "Mobility constraint": mobility,
        "Preferred transport": transport
    }
    if any(s != '' for s in personalize.values()):
        for key, value in personalize.items():
            if value != '':
                input_str = input_str + key + ': ' + value + '. '

    # print(input_str)
    # call openAI with that string
    itin = travel_agent(input_str)

    # create a fun image
    image_str = "Create a 2 by 2 collage of hyper-realistic images from " + destination + " based on places visited in this itinerary: " + itin
    response = client.images.generate(
        model="dall-e-3",
        prompt=image_str,
        size="1024x1024",
        quality="standard",
        n=1,
    )

    image_url = response.data[0].url
    return itin, image_url


with gr.Blocks(theme=gr.themes.Default()) as demo:
    gr.Markdown(
        """
        # The NS Travel Agency
        We are the best AI travel agency in the world! Tell us about your travel plans and we will create an itinerary for you!
        """)
    with gr.Row():
        with gr.Column():
            destination = gr.Textbox(label="Where would you like to go?")
        with gr.Column():
            with gr.Row():
                duration = gr.Slider(label="How many days do you plan to stay?", value=5, minimum=1, maximum=30, step=1)
                individuals = gr.Slider(label="How many individuals will be traveling?", value=1, minimum=1, maximum=10,
                                        step=1)
                cost = gr.Slider(label="What is your daily budget in dollars?", value=100, minimum=50, maximum=1000,
                                 step=50)
    with gr.Row():
        with gr.Column():
            constraints = gr.Textbox(label="Are there any constraints or events you wish to add?")
        with gr.Column():
            events = gr.CheckboxGroup(["Hiking", "Nature", "Sports", "Museums", "Shopping"],
                                      label="What activities do you enjoy (select all)")
    with gr.Accordion("Personalize further!", open=False):
        with gr.Row():
            food = gr.Textbox(label="What kinda food do you like?")
        with gr.Row():
            mobility = gr.Textbox(label="Do you have any mobility or accessibility needs?")
        with gr.Row():
            transport = gr.Textbox(label="What is your preferred mode of transportation?")

    with gr.Row():
        btn = gr.Button("Create Itinerary", variant='primary')
    with gr.Row():
        with gr.Column():
            itin = gr.Text(label="Itinerary", autoscroll=False)
        with gr.Column():
            collage = gr.Image()

    btn.click(get_itinerary, inputs=[destination, duration, events, cost, individuals, constraints, food, mobility, transport],
              outputs=[itin, collage])

if __name__ == "__main__":
    demo.launch(auth=("AI", os.getenv('PASSWD')))