app.py

import streamlit as st
import pandas as pd
import requests
import folium
from streamlit_folium import st_folium
from astral import LocationInfo
from astral.sun import sun
from datetime import datetime
import os
from dotenv import load_dotenv
from groq import Groq # Moved this import to the top

# --- Configuration and API Key Loading ---
load_dotenv()
OPENWEATHER_API_KEY = os.getenv("OPENWEATHER_API_KEY")
GROQ_API_KEY = os.getenv("GROQ_API_KEY")

# --- Basic API Key Validation ---
if not OPENWEATHER_API_KEY:
    st.error("OpenWeatherMap API key not found. Please set OPENWEATHER_API_KEY in your .env file.")
    st.stop() # Stop execution if critical key is missing

if not GROQ_API_KEY:
    st.error("Groq API key not found. Please set GROQ_API_KEY in your .env file.")
    st.stop() # Stop execution if critical key is missing

# Initialize Groq client
groq_client = Groq(api_key=GROQ_API_KEY)


st.set_page_config(page_title="AI Flight Assistant", layout="wide")

# --- Custom CSS for better aesthetics ---
st.markdown("""
    <style>
    .main {background-color: #f8f9fa;}
    h1, h2, h3 {color: #003366;}
    .st-emotion-cache-nahz7x { /* Adjust sidebar width, check this class name if not working */
        width: 300px;
    }
    .st-emotion-cache-1cyp85f { /* Adjust main content padding, check this class name if not working */
        padding-left: 2rem;
        padding-right: 2rem;
    }
    .stButton>button {
        background-color: #007bff;
        color: white;
        border-radius: 0.5rem;
        padding: 0.5rem 1rem;
    }
    .stButton>button:hover {
        background-color: #0056b3;
        color: white;
    }
    </style>
""", unsafe_allow_html=True)

st.title("✈️ AI Flight Planning Assistant")
st.markdown("Plan VFR flights smartly with real-time weather, time-of-day checks, and interactive mapping.")

# --- Initialize session state variables ---
# This ensures these variables exist even on the first run
if 'flight_planned' not in st.session_state:
    st.session_state.flight_planned = False
if 'dep_weather' not in st.session_state:
    st.session_state.dep_weather = None
if 'arr_weather' not in st.session_state:
    st.session_state.arr_weather = None
if 'dep_sun' not in st.session_state:
    st.session_state.dep_sun = None
if 'arr_sun' not in st.session_state:
    st.session_state.arr_sun = None
# Initialize input values in session state
if 'dep_lat' not in st.session_state: st.session_state.dep_lat = 12.9538
if 'dep_lon' not in st.session_state: st.session_state.dep_lon = 77.4903
if 'arr_lat' not in st.session_state: st.session_state.arr_lat = 13.0827
if 'arr_lon' not in st.session_state: st.session_state.arr_lon = 80.2707
if 'dep_airport' not in st.session_state: st.session_state.dep_airport = "VOBL"
if 'arr_airport' not in st.session_state: st.session_state.arr_airport = "VOMM"
if 'flight_date' not in st.session_state: st.session_state.flight_date = datetime.today()

# Initialize chat history in session state
if "messages" not in st.session_state:
    st.session_state.messages = []


# --- Functions for fetching data ---
def get_weather(lat, lon):
    url = f"https://api.openweathermap.org/data/2.5/weather?lat={lat}&lon={lon}&appid={OPENWEATHER_API_KEY}&units=metric"
    try:
        res = requests.get(url)
        res.raise_for_status() # Raises HTTPError for bad responses (4xx or 5xx)
        data = res.json()
        return {
            "temp": data['main']['temp'],
            "weather": data['weather'][0]['description'],
            "wind_speed": data['wind']['speed'],
            "city_name": data.get('name', 'N/A')
        }
    except requests.exceptions.RequestException as e:
        st.error(f"Error fetching weather data: {e}. Check API key or coordinates.")
        return None
    except KeyError as e:
        st.error(f"Unexpected weather data format from API: {e}.")
        return None

def get_sun_times(lat, lon, date):
    try:
        city = LocationInfo("Custom", "World", "UTC", lat, lon)
        s = sun(city.observer, date=date)
        return s
    except Exception as e:
        st.error(f"Error calculating sun times: {e}.")
        return None

# --- Sidebar for Flight Inputs ---
st.sidebar.header("✈️ Flight Parameters")

with st.sidebar:
    dep_airport = st.text_input("Departure Airport ICAO Code", value=st.session_state.dep_airport, key="dep_airport_input")
    arr_airport = st.text_input("Arrival Airport ICAO Code", value=st.session_state.arr_airport, key="arr_airport_input")
    flight_date = st.date_input("Flight Date", value=st.session_state.flight_date, key="flight_date_input")

    dep_lat = st.number_input("Departure Latitude", value=st.session_state.dep_lat, key="dep_lat_input")
    dep_lon = st.number_input("Departure Longitude", value=st.session_state.dep_lon, key="dep_lon_input")
    arr_lat = st.number_input("Arrival Latitude", value=st.session_state.arr_lat, key="arr_lat_input")
    arr_lon = st.number_input("Arrival Longitude", value=st.session_state.arr_lon, key="arr_lon_input")

    # Update session state when inputs change (this runs on every rerun)
    st.session_state.dep_airport = dep_airport
    st.session_state.arr_airport = arr_airport
    st.session_state.flight_date = flight_date
    st.session_state.dep_lat = dep_lat
    st.session_state.dep_lon = dep_lon
    st.session_state.arr_lat = arr_lat
    st.session_state.arr_lon = arr_lon


# --- Plan Flight Button Logic ---
if st.sidebar.button("Plan Flight"):
    with st.spinner("Fetching weather and sunlight data..."):
        st.session_state.dep_weather = get_weather(st.session_state.dep_lat, st.session_state.dep_lon)
        st.session_state.arr_weather = get_weather(st.session_state.arr_lat, st.session_state.arr_lon)
        st.session_state.dep_sun = get_sun_times(st.session_state.dep_lat, st.session_state.dep_lon, st.session_state.flight_date)
        st.session_state.arr_sun = get_sun_times(st.session_state.arr_lat, st.session_state.arr_lon, st.session_state.flight_date)
        st.session_state.flight_planned = True # Set status to true after fetching
    st.rerun() # Trigger a rerun to display content in tabs immediately

# --- Main Content Area with Tabs ---
tab1, tab2 = st.tabs(["📊 Flight Plan Overview", "🤖 AI Flight Assistant"])

with tab1: # Flight Plan Overview Tab
    if st.session_state.flight_planned:
        st.subheader("Current Flight Plan Details")

        display_col1, display_col2 = st.columns(2)

        with display_col1:
            st.markdown("#### 🌤️ Departure Weather")
            if st.session_state.dep_weather:
                st.metric(label="Location", value=st.session_state.dep_weather.get('city_name', 'N/A'))
                st.metric(label="Temperature", value=f"{st.session_state.dep_weather['temp']}°C")
                st.metric(label="Condition", value=st.session_state.dep_weather['weather'])
                st.metric(label="Wind Speed", value=f"{st.session_state.dep_weather['wind_speed']} m/s")
            else:
                st.info("Departure weather data not available.")

            st.markdown("#### ☀️ Departure Sunlight")
            if st.session_state.dep_sun:
                st.write(f"🌅 Sunrise: {st.session_state.dep_sun['sunrise'].time()}")
                st.write(f"🌇 Sunset: {st.session_state.dep_sun['sunset'].time()}")
            else:
                st.info("Departure sun times not available.")

        with display_col2:
            st.markdown("#### 🌤️ Arrival Weather")
            if st.session_state.arr_weather:
                st.metric(label="Location", value=st.session_state.arr_weather.get('city_name', 'N/A'))
                st.metric(label="Temperature", value=f"{st.session_state.arr_weather['temp']}°C")
                st.metric(label="Condition", value=st.session_state.arr_weather['weather'])
                st.metric(label="Wind Speed", value=f"{st.session_state.arr_weather['wind_speed']} m/s")
            else:
                st.info("Arrival weather data not available.")

            st.markdown("#### ☀️ Arrival Sunlight")
            if st.session_state.arr_sun:
                st.write(f"🌅 Sunrise: {st.session_state.arr_sun['sunrise'].time()}")
                st.write(f"🌇 Sunset: {st.session_state.arr_sun['sunset'].time()}")
            else:
                st.info("Arrival sun times not available.")

        # --- Map ---
        st.subheader("🗺️ Flight Path")
        if all(isinstance(coord, (int, float)) for coord in [st.session_state.dep_lat, st.session_state.dep_lon, st.session_state.arr_lat, st.session_state.arr_lon]):
            m = folium.Map(location=[(st.session_state.dep_lat + st.session_state.arr_lat)/2, (st.session_state.dep_lon + st.session_state.arr_lon)/2], zoom_start=6)
            folium.Marker([st.session_state.dep_lat, st.session_state.dep_lon], tooltip=f"Departure ({st.session_state.dep_airport})").add_to(m)
            folium.Marker([st.session_state.arr_lat, st.session_state.arr_lon], tooltip=f"Arrival ({st.session_state.arr_airport})").add_to(m)
            folium.PolyLine([[st.session_state.dep_lat, st.session_state.dep_lon], [st.session_state.arr_lat, st.session_state.arr_lon]], color="blue").add_to(m)
            st_folium(m, width=800, height=400)
        else:
            st.warning("Invalid latitude/longitude provided for map display.")
    else:
        st.info("Enter flight parameters in the sidebar and click 'Plan Flight' to see details.")


with tab2: # AI Flight Assistant Tab
    st.subheader("🤖 AI Flight Chat Assistant")
    st.markdown("Ask me questions related to your current flight plan, weather, or general aviation topics.")

    # --- Display Chat History ---
    for message in st.session_state.messages:
        with st.chat_message(message["role"]):
            st.markdown(message["content"])

    # --- Chat Input ---
    user_query = st.chat_input("Ask a question about your flight plan...")

    if user_query:
        st.session_state.messages.append({"role": "user", "content": user_query})
        with st.chat_message("user"):
            st.markdown(user_query)

        # --- RAG for Chatbot: Provide flight context to Llama 3 ---
        # Ensure that context is built only if a flight has been planned
        flight_context = ""
        if st.session_state.flight_planned:
            flight_context = f"""
            Current Flight Plan Details:
            Departure Airport ICAO: {st.session_state.dep_airport}
            Arrival Airport ICAO: {st.session_state.arr_airport}
            Flight Date: {st.session_state.flight_date}
            Departure Latitude: {st.session_state.dep_lat}, Longitude: {st.session_state.dep_lon}
            Arrival Latitude: {st.session_state.arr_lat}, Longitude: {st.session_state.arr_lon}
            """
            if st.session_state.dep_weather:
                flight_context += f"\nDeparture Weather: At {st.session_state.dep_weather.get('city_name', 'N/A')}: Temperature {st.session_state.dep_weather['temp']}°C, Condition {st.session_state.dep_weather['weather']}, Wind Speed {st.session_state.dep_weather['wind_speed']} m/s."
            if st.session_state.arr_weather:
                flight_context += f"\nArrival Weather: At {st.session_state.arr_weather.get('city_name', 'N/A')}: Temperature {st.session_state.arr_weather['temp']}°C, Condition {st.session_state.arr_weather['weather']}, Wind Speed {st.session_state.arr_weather['wind_speed']} m/s."
            if st.session_state.dep_sun:
                flight_context += f"\nDeparture Sunrise: {st.session_state.dep_sun['sunrise'].time()}, Sunset: {st.session_state.dep_sun['sunset'].time()}."
            if st.session_state.arr_sun:
                flight_context += f"\nArrival Sunrise: {st.session_state.arr_sun['sunrise'].time()}, Sunset: {st.session_state.arr_sun['sunset'].time()}."
        else:
            flight_context = "No specific flight plan data available yet. Please plan a flight first."


        messages_for_llm = [
            {"role": "system", "content": f"You are a helpful AI flight assistant. Here's the flight context: {flight_context}\nAnswer questions based on this context and general aviation knowledge. If a question is outside the provided context, state that you don't have that specific information but can answer general aviation questions. Keep answers concise."}
        ]

        # Add recent chat history to LLM messages (excluding system message for every turn)
        for msg in st.session_state.messages:
            if msg["role"] != "system": # Avoid sending system message multiple times
                messages_for_llm.append({"role": msg["role"], "content": msg["content"]})


        with st.chat_message("assistant"):
            with st.spinner("Thinking..."):
                try:
                    chat_completion = groq_client.chat.completions.create(
                        messages=messages_for_llm,
                        model="llama3-8b-8192",
                        temperature=0.7,
                        max_tokens=500,
                        stream=True
                    )
                    full_response = ""
                    response_placeholder = st.empty()
                    for chunk in chat_completion:
                        if chunk.choices[0].delta.content:
                            full_response += chunk.choices[0].delta.content
                            response_placeholder.markdown(full_response + "▌") # Add blinking cursor
                    response_placeholder.markdown(full_response) # Final response without cursor
                    st.session_state.messages.append({"role": "assistant", "content": full_response})
                except Exception as e:
                    st.error(f"Error communicating with AI: {e}")
                    st.session_state.messages.append({"role": "assistant", "content": f"Sorry, I encountered an error: {e}"})