import os
import gradio as gr
from dotenv import load_dotenv
from smolagents import CodeAgent, GoogleSearchTool, OpenAIServerModel

# Load environment variables
load_dotenv()
OPENAI_API_KEY = os.getenv("MISTRAL_API_KEY")
search_tool = GoogleSearchTool(provider="serpapi")

model = OpenAIServerModel(
    model_id="gpt-4",  
    api_key=OPENAI_API_KEY
)

agent = CodeAgent(
    tools=[search_tool],
    model=model,
    add_base_tools=False
)

# === Travel Survey Classification ===
def classify_travel_profile(*answers):
    preferences = {
        "Budget": 1, "Mid-range": 2, "Luxury": 3,
        "Solo": 1, "Couple": 2, "Family": 3, "Group": 3,
        "Weekend (1–3 days)": 1, "Short trip (4–7 days)": 2, "Extended (8+ days)": 3,
        "Nature": 1, "City": 2, "Exotic/Remote": 3,
        "Public Transport": 1, "Taxi/Rideshare": 2, "Private Car": 3
    }
    total = sum(preferences.get(a, 2) for a in answers)
    level = round(total / len(answers))
    category = ["Budget Traveler", "Balanced Traveler", "Luxury Traveler"][level - 1]
    return category, f"🌍 Your travel profile: **{category}**"

# === Budget Estimation ===
def estimate_budget(profile, q1_answer, q2_answer, q3_answer, q4_answer, q5_answer):
    answers = [q1_answer, q2_answer, q3_answer, q4_answer, q5_answer]
    base_cost = {
        "Budget Traveler": 50,
        "Balanced Traveler": 150,
        "Luxury Traveler": 350
    }.get(profile, 100)
    days = {"Weekend (1–3 days)": 2, "Short trip (4–7 days)": 5, "Extended (8+ days)": 10}.get(answers[2], 5)
    group_factor = {"Solo": 1, "Couple": 0.9, "Family": 0.8, "Group": 0.75}.get(answers[1], 1)
    daily_cost = base_cost * group_factor
    total = daily_cost * days
    budget_text = f"""💸 **Estimated Budget**
- **Daily Cost:** ${daily_cost:.2f}
- **Trip Duration:** {days} days
- **Total Estimate:** ${total:.2f}
_For a {profile} going on a {answers[2]} in a group of {answers[1]}, primarily visiting {answers[3]} destinations._
"""
    return budget_text, budget_text


# === Chat Agent ===
def travel_chat(user_message, chat_history, profile, budget_summary):
    if not profile or not budget_summary:
        return chat_history, chat_history + [["⚠️", "Please complete your profile and budget estimate first."]]
    prompt = (
        f"Traveler Profile: {profile}\n"
        f"Budget Summary:\n{budget_summary}\n"
        f"User: {user_message}\n"
        f"Advisor:"
    )
    try:
        reply = agent.run(prompt)
    except Exception as e:
        reply = f"❌ Error: {e}"
    chat_history.append([user_message, reply])
    return chat_history, chat_history

# === Gradio UI ===
with gr.Blocks() as demo:
    travel_profile = gr.State()
    travel_summary = gr.State()
    chat_history = gr.State([])

    gr.Markdown("# ✈️ Trip Budget Advisor")

    with gr.Row():
        q1 = gr.Radio(["Budget", "Mid-range", "Luxury"], label="1. What’s your travel style?")
        q2 = gr.Radio(["Solo", "Couple", "Family", "Group"], label="2. Who are you traveling with?")
        q3 = gr.Radio(["Weekend (1–3 days)", "Short trip (4–7 days)", "Extended (8+ days)"], label="3. Duration?")
        q4 = gr.Radio(["Nature", "City", "Exotic/Remote"], label="4. Destination type?")
        q5 = gr.Radio(["Public Transport", "Taxi/Rideshare", "Private Car"], label="5. Preferred transport?")

    profile_btn = gr.Button("Generate Travel Profile")
    profile_out = gr.Markdown()

    profile_btn.click(
        classify_travel_profile,
        inputs=[q1, q2, q3, q4, q5],
        outputs=[travel_profile, profile_out]
    )

    gr.Markdown("## 🧾 Budget Estimator")
    budget_btn = gr.Button("Estimate Budget")
    budget_out = gr.Markdown()
    budget_btn.click(
        fn=lambda prof, a1, a2, a3, a4, a5: estimate_budget(prof, a1, a2, a3, a4, a5),
        inputs=[travel_profile, q1, q2, q3, q4, q5],
        outputs=[budget_out, travel_summary]
    )


    gr.Markdown("## 💬 Travel Chat Assistant")
    chatbot = gr.Chatbot()
    user_input = gr.Textbox(placeholder="Ask travel questions like best cities under budget...", label="Ask something")
    chat_btn = gr.Button("Send")

    chat_btn.click(
        travel_chat,
        inputs=[user_input, chat_history, travel_profile, travel_summary],
        outputs=[chatbot, chat_history]
    )

demo.launch()