import gradio as gr
from huggingface_hub import InferenceClient
import json
from typing import Dict, List, Tuple

# Hotel data
HOTELS = [
    {
        "name": "Azure Haven",
        "pricingPerDay": 250,
        "roomTypes": ["Single Room", "Double Room", "Suite"],
        "description": "Azure Haven is a tranquil getaway offering breathtaking ocean views and luxurious amenities. Ideal for a peaceful escape or a romantic retreat."
    },
    {
        "name": "Golden Sands Resort",
        "pricingPerDay": 300,
        "roomTypes": ["Standard Room", "Deluxe Room", "Family Suite"],
        "description": "Nestled on a pristine beach, Golden Sands Resort is perfect for family vacations and adventure seekers. Enjoy watersports, beachside dining, and more."
    },
    {
        "name": "Mountain Bliss Lodge",
        "pricingPerDay": 220,
        "roomTypes": ["Cabin", "Deluxe Cabin", "Mountain Suite"],
        "description": "Mountain Bliss Lodge offers cozy cabins and stunning mountain views. A perfect destination for nature lovers and outdoor enthusiasts."
    },
    {
        "name": "Urban Oasis Hotel",
        "pricingPerDay": 280,
        "roomTypes": ["Single Room", "Double Room", "Executive Suite"],
        "description": "Located in the heart of the city, Urban Oasis Hotel combines modern luxury with convenient access to top attractions, dining, and entertainment."
    },
    {
        "name": "Lakeside Retreat",
        "pricingPerDay": 270,
        "roomTypes": ["Single Room", "Double Room", "Lakeside Suite"],
        "description": "Lakeside Retreat offers serene lakeside accommodations with a focus on relaxation and rejuvenation. Enjoy spa treatments, boating, and picturesque views."
    }
]

# Cost constants
COST_FACTORS = {
    "local_transport_per_day": {
        "Budget Backpacking": 10,
        "Cultural & Sightseeing": 20,
        "Adventure & Outdoor": 25,
        "Food & Culinary": 20,
        "Relaxation & Wellness": 15,
        "Luxury & High-End": 50,
        "Family-Friendly": 30
    },
    "food_cost_per_day": {
        "Budget Backpacking": 30,
        "Cultural & Sightseeing": 50,
        "Adventure & Outdoor": 45,
        "Food & Culinary": 80,
        "Relaxation & Wellness": 60,
        "Luxury & High-End": 120,
        "Family-Friendly": 70
    },
    "activities_cost_per_day": {
        "Budget Backpacking": 20,
        "Cultural & Sightseeing": 40,
        "Adventure & Outdoor": 60,
        "Food & Culinary": 45,
        "Relaxation & Wellness": 80,
        "Luxury & High-End": 150,
        "Family-Friendly": 70
    }
}

client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")

def calculate_flight_cost(origin: str, destination: str, budget: int) -> float:
    """Estimate flight cost based on budget and journey type."""
    # Basic estimation - can be enhanced with actual flight data API
    if budget < 2000:
        return budget * 0.3  # 30% of budget for budget trips
    elif budget < 5000:
        return budget * 0.35  # 35% for mid-range trips
    else:
        return budget * 0.4  # 40% for luxury trips

def calculate_daily_costs(holiday_type: str) -> Dict[str, float]:
    """Calculate daily costs based on holiday type."""
    return {
        "local_transport": COST_FACTORS["local_transport_per_day"][holiday_type],
        "food": COST_FACTORS["food_cost_per_day"][holiday_type],
        "activities": COST_FACTORS["activities_cost_per_day"][holiday_type]
    }

def get_suitable_hotels(budget_per_day: float, holiday_type: str) -> List[Dict]:
    """Get suitable hotels based on daily budget and holiday type."""
    max_hotel_budget = budget_per_day * 0.4  # 40% of daily budget for accommodation
    
    suitable_hotels = []
    for hotel in HOTELS:
        if hotel["pricingPerDay"] <= max_hotel_budget:
            match_score = 0
            # Score hotels based on type match
            if holiday_type == "Relaxation & Wellness" and ("Haven" in hotel["name"] or "Retreat" in hotel["name"]):
                match_score += 2
            elif holiday_type == "Adventure & Outdoor" and ("Mountain" in hotel["name"] or "Sands" in hotel["name"]):
                match_score += 2
            elif holiday_type == "Cultural & Sightseeing" and "Urban" in hotel["name"]:
                match_score += 2
            elif holiday_type == "Family-Friendly" and "Family" in str(hotel["roomTypes"]):
                match_score += 2
            
            suitable_hotels.append({"hotel": hotel, "match_score": match_score})
    
    # Sort by match score and then price
    sorted_hotels = sorted(suitable_hotels, key=lambda x: (-x["match_score"], x["hotel"]["pricingPerDay"]))
    return [hotel["hotel"] for hotel in sorted_hotels]

def calculate_total_costs(
    origin: str,
    destination: str,
    budget: int,
    num_days: int,
    holiday_type: str,
    selected_hotel: Dict
) -> Dict[str, float]:
    """Calculate detailed cost breakdown."""
    flight_cost = calculate_flight_cost(origin, destination, budget)
    daily_costs = calculate_daily_costs(holiday_type)
    
    accommodation_cost = selected_hotel["pricingPerDay"] * num_days
    local_transport_cost = daily_costs["local_transport"] * num_days
    food_cost = daily_costs["food"] * num_days
    activities_cost = daily_costs["activities"] * num_days
    
    total_cost = (
        flight_cost +
        accommodation_cost +
        local_transport_cost +
        food_cost +
        activities_cost
    )
    
    return {
        "flight": flight_cost,
        "accommodation": accommodation_cost,
        "local_transport": local_transport_cost,
        "food": food_cost,
        "activities": activities_cost,
        "total": total_cost,
        "per_day": total_cost / num_days
    }

def format_cost_breakdown(costs: Dict[str, float]) -> str:
    """Format cost breakdown for display."""
    return f"""Cost Breakdown:
- Flight: ${costs['flight']:,.2f}
- Accommodation: ${costs['accommodation']:,.2f}
- Local Transportation: ${costs['local_transport']:,.2f}
- Food & Dining: ${costs['food']:,.2f}
- Activities & Entertainment: ${costs['activities']:,.2f}

Total Cost: ${costs['total']:,.2f}
Average Daily Cost: ${costs['per_day']:,.2f}"""

def generate_itinerary(
    destination: str,
    origin: str,
    budget: int,
    num_days: int,
    holiday_type: str,
    system_message: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
) -> str:
    # Calculate initial costs and get hotel recommendations
    recommended_hotels = get_suitable_hotels(budget/num_days, holiday_type)
    if not recommended_hotels:
        return "Error: No suitable hotels found for your budget. Please increase your budget or reduce the number of days."
    
    selected_hotel = recommended_hotels[0]
    costs = calculate_total_costs(origin, destination, budget, num_days, holiday_type, selected_hotel)
    
    if costs["total"] > budget:
        return f"""Warning: The estimated total cost (${costs['total']:,.2f}) exceeds your budget (${budget:,.2f}).
Please consider:
1. Increasing your budget
2. Reducing the number of days
3. Choosing a more budget-friendly holiday type
4. Selecting different destinations

{format_cost_breakdown(costs)}"""
    
    # Construct prompt with accurate cost information
    prompt = f"""Create a detailed {num_days}-day travel itinerary for a trip:
From: {origin}
To: {destination}
Budget Breakdown: {format_cost_breakdown(costs)}
Holiday Type: {holiday_type}

Recommended Accommodation:
{selected_hotel['name']} (${selected_hotel['pricingPerDay']}/night)
{selected_hotel['description']}

Please provide:
1. Day-by-day itinerary with specific activities and timing
2. Daily budget allocation matching the cost breakdown above
3. Specific local transportation recommendations
4. Restaurant suggestions within the food budget (${costs['food']/num_days:.2f}/day)
5. Activity recommendations within the activities budget (${costs['activities']/num_days:.2f}/day)

Format as a clear day-by-day breakdown with budgets for each activity."""

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": prompt}
    ]

    response = ""
    for message in client.chat_completion(
        messages,
        max_tokens=max_tokens,
        stream=True,
        temperature=temperature,
        top_p=top_p,
    ):
        token = message.choices[0].delta.content
        response += token
        yield response

# [Previous Gradio interface code remains the same, just update the button click handler to use the new generate_itinerary function]
# Create the Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🌎 AI Travel Itinerary Generator ✈️")
    gr.Markdown("Enter your travel details for a customized itinerary with accurate cost calculations!")
    
    with gr.Row():
        with gr.Column():
            origin = gr.Textbox(label="Starting Location", placeholder="e.g., New York, USA")
            destination = gr.Textbox(label="Destination", placeholder="e.g., Paris, France")
            budget = gr.Slider(minimum=1000, maximum=15000, value=3000, step=100, label="Total Budget (USD)")
            num_days = gr.Slider(minimum=1, maximum=14, value=3, step=1, label="Number of Days")
            holiday_type = gr.Dropdown(
                choices=list(COST_FACTORS["local_transport_per_day"].keys()),
                label="Type of Holiday",
                value="Cultural & Sightseeing"
            )
    
    with gr.Accordion("Advanced Settings", open=False):
        system_message = gr.Textbox(
            value="You are an experienced travel planner. Provide detailed itineraries that strictly adhere to the given budget constraints and cost breakdown.",
            label="System Message",
            lines=3
        )
        max_tokens = gr.Slider(minimum=1, maximum=2048, value=1024, step=1, label="Max Tokens")
        temperature = gr.Slider(minimum=0.1, maximum=1.0, value=0.7, step=0.1, label="Temperature")
        top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p")

    with gr.Row():
        generate_button = gr.Button("Generate Itinerary", variant="primary")
        output = gr.Textbox(label="Your Custom Itinerary", lines=20, show_copy_button=True)

    generate_button.click(
        fn=generate_itinerary,
        inputs=[
            destination,
            origin,
            budget,
            num_days,
            holiday_type,
            system_message,
            max_tokens,
            temperature,
            top_p,
        ],
        outputs=output,
    )

if __name__ == "__main__":
    demo.launch()