main.py

from fastapi import FastAPI, HTTPException, Query
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel
import requests
import pandas as pd
from datetime import datetime, date
import google.generativeai as genai
import os
import json
from dotenv import load_dotenv
from sheets_client import add_subscriber_to_sheet
import random

load_dotenv()
app = FastAPI(
    title="Aussie Backpacker Flow API",
    description="Provides flight demand analysis and AI-powered insights for Australian hostels.",
    version="1.0.0"
)

origins = ["*"]
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

FLIGHT_API_KEY = os.getenv("FLIGHT_API_KEY")
GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

if GEMINI_API_KEY:
    try:
        genai.configure(api_key=GEMINI_API_KEY)
    except Exception as e:
        print(f"Could not configure Gemini API: {e}")

AUSTRALIAN_CITY_CODES = {
    "Sydney": "SYD", "Melbourne": "MEL", "Brisbane": "BNE", "Perth": "PER",
    "Adelaide": "ADL", "Canberra": "CBR", "Gold Coast": "OOL", "Cairns": "CNS",
    "Hobart": "HBA", "Darwin": "DRW"
}

class NewsletterPayload(BaseModel):
    email: str
    favDestinations: list[str] = []
    travelOrigin: str = ""
    dob: str = ""

class ChatPayload(BaseModel):
    message: str

def fetch_flight_data(api_key, departure_airport, arrival_airport, date_str):
    url = f"https://api.flightapi.io/onewaytrip/{api_key}/{departure_airport}/{arrival_airport}/{date_str}/1/0/0/Economy/AUD"
    try:
        response = requests.get(url, timeout=30)
        response.raise_for_status()
        return response.json()
    except requests.exceptions.RequestException as e:
        print(f"Error fetching flight data: {e}")
        return None

def parse_and_process_data(data):
    if not data or 'itineraries' not in data or not data.get('itineraries'):
        return pd.DataFrame()
    carriers = {c['id']: c for c in data.get('carriers', [])}
    flight_options = []
    for i, itinerary in enumerate(data.get('itineraries', [])):
        price_info = itinerary.get('pricing_options', [{}])[0].get('price', {})
        price = price_info.get('amount')
        leg_id = itinerary.get('leg_ids', [None])[0]
        leg = next((l for l in data.get('legs', []) if l['id'] == leg_id), None)
        if not all([price, leg]):
            continue
        marketing_carrier_id = leg.get('marketing_carrier_ids', [None])[0]
        carrier_info = carriers.get(marketing_carrier_id, {})
        airline_name = carrier_info.get('name', "Unknown Airline")
        flight_number = "N/A"
        if leg.get('segment_ids'):
            segment_id = leg['segment_ids'][0]
            segment = next((s for s in data.get('segments', []) if s['id'] == segment_id), None)
            if segment:
                carrier_code = carrier_info.get('code', "XX")
                flight_num_part = segment.get('marketing_flight_number', leg['id'][:3])
                flight_number = f"{carrier_code}{flight_num_part}"
        flight_options.append({
            "id": i, "airline": airline_name, "flight": flight_number,
            "departure": pd.to_datetime(leg['departure']).strftime('%H:%M'),
            "arrival": pd.to_datetime(leg['arrival']).strftime('%H:%M'), "price": price
        })
    return pd.DataFrame(flight_options).sort_values(by="price").reset_index(drop=True)

def get_dashboard_ai_analysis(df, origin, destination, date_str):
    if not GEMINI_API_KEY or df.empty:
        return "AI analysis could not be performed due to a configuration issue or lack of data."
    cheapest_flight = df.iloc[0]
    data_summary = f"""Context: Flight search from {origin} to {destination} for {date_str}. Total flights: {len(df)}. Price Range: ${df['price'].min():.2f} to ${df['price'].max():.2f}. Cheapest: {cheapest_flight['airline']} for ${cheapest_flight['price']:.2f}."""
    prompt = f"You are a sharp, concise market analyst for 'Aussie Backpacker Flow'. Based on the following summary, provide actionable insights in markdown bullet points: {data_summary}. Focus on: a one-sentence market snapshot, top budget carriers, demand interpretation, and one specific marketing tip for a hostel manager."
    try:
        model = genai.GenerativeModel('gemini-1.5-flash-latest')
        response = model.generate_content(prompt)
        report_header = f"Based on the data for **{origin} to {destination}** on **{date_str}**, here are the key insights:"
        return f"{report_header}\n\n{response.text}"
    except Exception as e:
        return f"An error occurred during AI analysis: {e}"

def generate_dynamic_trend_data(origin, destination):
    price_tiers = {
        "short": (90, 180),
        "medium": (150, 300),
        "long": (250, 450)
    }
    long_haul_cities = ["Perth"]
    tier = "long" if origin in long_haul_cities or destination in long_haul_cities else "short"
    if origin in ["Sydney", "Melbourne", "Brisbane"] and destination in ["Cairns", "Adelaide", "Hobart"]:
        tier = "medium"

    min_base, max_base = price_tiers[tier]
    
    data = {}
    routes_to_generate = {f"{origin[:3].upper()}-{destination[:3].upper()}": (min_base, max_base)}
    other_routes = list(price_tiers.keys())
    random.shuffle(other_routes)
    for i in range(2):
        tier_key = other_routes[i]
        mock_origin, mock_dest = ("SYD","MEL") if tier_key == "short" else ("BNE","CNS") if tier_key == "medium" else ("PER","SYD")
        routes_to_generate[f"{mock_origin}-{mock_dest}"] = price_tiers[tier_key]

    for route, (min_b, max_b) in routes_to_generate.items():
        route_data = []
        for i in range(1, 31):
            day_factor = (i % 7 - 3) * 5
            min_price = round(min_b + day_factor + random.uniform(-10, 10))
            max_price = round(max_b + day_factor + random.uniform(-10, 20))
            avg_price = round((min_price + max_price) / 2)
            route_data.append({"day": i, "minPrice": min_price, "avgPrice": avg_price, "maxPrice": max_price})
        data[route] = route_data
    return data

def generate_mock_heatmap_data():
    airports = list(AUSTRALIAN_CITY_CODES.keys())
    routes = []
    for _ in range(10):
        from_city, to_city = random.sample(airports, 2)
        price = random.randint(70, 400)
        status = 'green' if price < 150 else 'yellow' if price < 280 else 'red'
        routes.append({"from": from_city, "to": to_city, "status": status, "price": f"${price}"})
    return routes

def generate_mock_topmovers_data():
    routes = ["SYD → MEL", "BNE → CNS", "MEL → ADL", "SYD → CBR", "SYD → OOL"]
    cities = ["Gold Coast (OOL)", "Cairns (CNS)", "Sydney (SYD)"]
    price_drops = [{"route": r, "airline": random.choice(["Jetstar", "Rex", "Virgin"]), "drop": f"{random.randint(15, 35)}%", "price": f"${random.randint(70, 120)}"} for r in random.sample(routes, 5)]
    high_demand = [{"city": c, "flights": random.randint(30, 120), "change": f"+{random.randint(8, 25)}%"} for c in cities]
    return {"priceDrops": price_drops, "highDemand": high_demand}

@app.get("/", response_class=HTMLResponse)
def read_root():
    html_content = """
    <html>
        <head>
            <title>Aussie Backpacker Flow API</title>
            <style>
                body { font-family: -apple-system, BlinkMacSystemFont, "Segoe UI", Roboto, sans-serif; display: flex; justify-content: center; align-items: center; height: 100vh; margin: 0; background-color: #0d1117; color: #c9d1d9; }
                .container { max-width: 700px; text-align: center; background-color: #161b22; border: 1px solid #30363d; border-radius: 12px; padding: 50px; }
                h1 { color: #58a6ff; font-size: 2.5rem; margin-bottom: 1rem; }
                code { background-color: #30363d; padding: 4px 8px; border-radius: 6px; font-family: "SF Mono", "Consolas", monospace; }
                a { color: #58a6ff; text-decoration: none; }
                .status { display: inline-block; background-color: #238636; color: white; padding: 8px 18px; border-radius: 20px; font-weight: bold; margin-bottom: 1.5rem; }
            </style>
        </head>
        <body>
            <div class="container">
                <p class="status">API is Online</p>
                <h1>Aussie Backpacker Flow API</h1>
                <p>This is the backend server for the flight demand analysis tool.</p>
                <p>To see the interactive API documentation, visit <a href="/docs">/docs</a>.</p>
            </div>
        </body>
    </html>
    """
    return HTMLResponse(content=html_content)

@app.post("/api/newsletter")
def handle_newsletter(payload: NewsletterPayload):
    success, message = add_subscriber_to_sheet(
        email=payload.email,
        fav_destinations=payload.favDestinations,
        origin=payload.travelOrigin,
        dob=payload.dob
    )
    if not success:
        raise HTTPException(status_code=400, detail=message)
    return {"message": message}

@app.get("/api/dashboard-data")
def get_dashboard_data(origin: str, destination: str, date: str):
    origin_code = AUSTRALIAN_CITY_CODES.get(origin)
    dest_code = AUSTRALIAN_CITY_CODES.get(destination)
    if not origin_code or not dest_code:
        raise HTTPException(status_code=400, detail="Invalid city name provided.")
    raw_data = fetch_flight_data(FLIGHT_API_KEY, origin_code, dest_code, date)
    flight_df = parse_and_process_data(raw_data)
    if flight_df.empty:
        raise HTTPException(status_code=404, detail=f"No live flight data could be found for {origin} to {destination} on {date}. Please try another route or date.")
    
    cheapest_flight_row = flight_df.loc[flight_df['price'].idxmin()]
    
    dashboard_payload = {
        "insightCards": {
            "cheapestFlight": {"price": cheapest_flight_row['price'], "airline": cheapest_flight_row['airline'], "flightNumber": cheapest_flight_row['flight']},
            "busiestAirline": {"name": flight_df['airline'].mode()[0], "flightCount": len(flight_df)},
            "bestDeal": {"name": cheapest_flight_row['airline'], "savings": "Top Value"}
        },
        "aiAnalystReport": get_dashboard_ai_analysis(flight_df, origin, destination, date),
        "flightPriceChart": flight_df.groupby('airline')['price'].min().reset_index().rename(columns={'airline': 'name'}).to_dict('records'),
        "flightDataTable": flight_df.to_dict('records'),
        "detailedTrendChart": generate_dynamic_trend_data(origin, destination),
        "airfareHeatmap": generate_mock_heatmap_data(),
        "topMovers": generate_mock_topmovers_data(),
    }
    return dashboard_payload

@app.post("/api/chat")
def handle_chat(payload: ChatPayload):
    user_message = payload.message
    if not GEMINI_API_KEY:
        return {"reply": "Chatbot is disabled. Backend needs a Gemini API key."}
    try:
        model = genai.GenerativeModel('gemini-1.5-flash-latest')
        intent_prompt = f"From the user's message, extract origin city, destination city, and a date (today is {date.today().strftime('%Y-%m-%d')}). Respond ONLY with a valid JSON object. Keys: 'origin', 'destination', 'date'. Missing values should be null. Message: '{user_message}'"
        response = model.generate_content(intent_prompt)
        json_str = response.text.strip().replace("```json", "").replace("```", "")
        params = json.loads(json_str)
        if not all(params.get(k) for k in ['origin', 'destination', 'date']):
            return {"reply": "I can help with that! To give you the best info, I need the origin city, destination city, and the date you're interested in."}
        origin_code = next((code for name, code in AUSTRALIAN_CITY_CODES.items() if name.lower() in params['origin'].lower()), None)
        dest_code = next((code for name, code in AUSTRALIAN_CITY_CODES.items() if name.lower() in params['destination'].lower()), None)
        if not origin_code or not dest_code:
            return {"reply": "Sorry, I couldn't recognise those city names. Please use major Australian cities."}
        raw_data = fetch_flight_data(FLIGHT_API_KEY, origin_code, dest_code, params['date'])
        flight_df = parse_and_process_data(raw_data)
        if flight_df.empty:
            return {"reply": f"I couldn't find any flights from {params['origin']} to {params['destination']} on {params['date']}."}
        summary_prompt = f"You are a helpful travel assistant. Based on this flight data, write a short, conversational summary of the cheapest option, and maybe one other good one. Data: {flight_df.head(3).to_markdown()}"
        summary_response = model.generate_content(summary_prompt)
        return {"reply": summary_response.text}
    except Exception as e:
        return {"reply": f"I had a little trouble with that. My apologies. Error: {e}"}