import os
from crewai import Crew
import streamlit as st

from textwrap import dedent
from agents import TravelAgents
from tasks import TravelTasks

from dotenv import load_dotenv

load_dotenv()

api_key = os.getenv("OPENAI_API_KEY")
if api_key:
    os.environ["OPENAI_API_KEY"] = api_key
else:
    raise ValueError("OPENAI_API_KEY environment variable not set")
# This is the main class that you will use to define your custom crew.
# You can define as many agents and tasks as you want in agents.py and tasks.py

class TripCrew:
    def __init__(self, origin_city, destination_city, date_of_departure, date_of_return, departure_flight, return_flight, food_preferences, exercise_preferences):
        self.origin_city = origin_city
        self.destination_city = destination_city
        self.date_of_departure = date_of_departure
        self.date_of_return = date_of_return
        self.departure_flight = departure_flight
        self.return_flight = return_flight
        self.food_preferences = food_preferences
        self.exercise_preferences = exercise_preferences
        
    def run(self):
        # Define your custom agents and tasks in agents.py and tasks.py
        agents = TravelAgents()
        tasks = TravelTasks()

        # Define your custom agents and tasks here
        sleep_specialist = agents.sleep_specialist()
        dietician = agents.dietician()
        fitness_trainer = agents.fitness_trainer()
        expert_itenary_planner = agents.expert_itenary_planner()
        
        # Custom tasks include agent name and variables as input
        sleep_schedule = tasks.sleep_schedule(
            sleep_specialist,
            self.origin_city,
            self.destination_city,
            self.date_of_departure,
            self.date_of_return,
            self.departure_flight,
            self.return_flight
        )

        diet_plan = tasks.diet_plan(
            dietician,
            self.date_of_departure,
            self.date_of_return,
            self.departure_flight,
            self.return_flight,
            self.food_preferences
        )

        fitness_routine = tasks.fitness_routine(
            fitness_trainer,
            self.date_of_departure,
            self.date_of_return,
            self.departure_flight,
            self.return_flight,
            self.exercise_preferences
        )

        plan_travel = tasks.plan_travel(
            expert_itenary_planner,
            sleep_schedule,
            diet_plan,
            fitness_routine
        )

        # Define your custom crew here
        crew = Crew(
            agents=[
                sleep_specialist,
                dietician,
                fitness_trainer,
                expert_itenary_planner
            ],
            tasks=[
                sleep_schedule,
                diet_plan,
                fitness_routine,
                plan_travel
            ],
            verbose=True,
        )

        result = crew.kickoff()
        return result

#This is the main function that you will use to run your custom crew.
if __name__ == "__main__":
    col1, col2 = st.columns(2)
    with col1:
        st.image("logo.jpg", width=250)
    with col2:
        st.markdown("# Trip Planner to Avoid Jet Lag")
        # if st.experimental_user.is_logged_in:
        #     st.write(f"Hello, {st.experimental_user.name}!")
        #     if st.button("Log out"):
        #         st.logout()
        #     st.write("This is a Trip Planner to avoid Jet Lag. Please fill in the details below to generate a plan.")
        # else:
        st.write("This is a Trip Planner to avoid Jet Lag. Please fill in the details below to generate a plan.")
            
    st.text("City names should be in the format: City, Country")
    col1, col2 = st.columns(2)
    with col1:
        origin_city = st.text_input("Origin City")
    with col2:
        destination_city = st.text_input("Destination City")
    
    st.text("Date format: 1 Mar 2025")
    col1, col2 = st.columns(2)
    with col1:
        date_of_departure = st.text_input("Date of Departure")
    with col2:
        date_of_return = st.text_input("Date of Return")
    
    st.text("Flight code should be in the format: Airline Code + Flight Number Eg: AI111")
    col1, col2 = st.columns(2)
    with col1:
        departure_flight = st.text_input("Departure Flight")
    with col2:
        return_flight = st.text_input("Return Flight")
    
    food_preference = st.selectbox(
        "Food Preference (Select one)",
        ("Vegetarian", "Non-Veg", "Vegan"),
    )
    
    exercise_preference = st.selectbox(
        "Exercise Preference (Select one)",
        ("Gym", "Running", "Do not exercise"),
    )
        
    # if not st.experimental_user.is_logged_in:
    #     col1, col2 = st.columns(2)
    #     with col1:
    #         st.write("Login to generate Plan") 
    #     with col2:
    #         if st.button("Log in with Google", type="primary"):
    #             st.login()
    # else:
    if st.button("Generate Plan", type="primary") and origin_city and destination_city and date_of_departure and date_of_return and departure_flight and return_flight and food_preference and exercise_preference:
        with st.spinner('Your personalised plan is getting ready...', show_time=True):
            trip_crew = TripCrew(origin_city, destination_city, date_of_departure, date_of_return, departure_flight, return_flight, food_preference, exercise_preference)
            result = trip_crew.run()
            st.markdown(result)