import os
import streamlit as st
import folium
from streamlit_folium import folium_static

# Set the page configuration
st.set_page_config(page_title="Google Maps 3D", layout="wide")

# Set your Google Maps API key
api_key = os.getenv("GM_TOKEN")

# Define state codes
state_codes = ["AL", "AK", "AZ", "AR", "CA", "CO", "CT", "DE", "FL", "GA", "HI", "ID", "IL", "IN", "IA", "KS", "KY", "LA", "ME", "MD", "MA", "MI", "MN", "MS", "MO", "MT", "NE", "NV", "NH", "NJ", "NM", "NY", "NC", "ND", "OH", "OK", "OR", "PA", "RI", "SC", "SD", "TN", "TX", "UT", "VT", "VA", "WA", "WV", "WI", "WY"]

# Define a dictionary of city coordinates
city_coords = {
    "Mound, MN": [44.9327, -93.6663],
    "San Francisco, CA": [37.7749, -122.4194],
    "New York, NY": [40.7128, -74.0060],
    "Chicago, IL": [41.8781, -87.6298],
    # Add more cities and their coordinates as needed
}

# Initialize session state
if "city" not in st.session_state:
    st.session_state["city"] = "Mound, MN"
if "state" not in st.session_state:
    st.session_state["state"] = "MN"

# Create Streamlit UI elements
st.title("🌍 Google Maps 3D")
city = st.text_input("Enter a city 🏙️", value=st.session_state["city"])
state = st.selectbox("Select a state 🎌", options=state_codes, index=state_codes.index(st.session_state["state"]))
search_button = st.button("Search 🔍")

# Update session state
st.session_state["city"] = city
st.session_state["state"] = state

# Get the latitude and longitude of the selected city and state
if search_button:
    city_state = f"{city}, {state}"
    if city_state in city_coords:
        map_center = city_coords[city_state]
        map_zoom = 12
    else:
        st.warning(f"Could not find the city '{city}' in the state '{state}'.")
        map_center = [37.7749, -122.4194]  # Example coordinates for San Francisco
        map_zoom = 12
else:
    map_center = city_coords["Mound, MN"]
    map_zoom = 12

# Create a Folium map instance
m = folium.Map(location=map_center, zoom_start=map_zoom)

# Add the Google Maps 3D tile layer
google_3d_layer = folium.TileLayer(
    tiles=f"https://tile.googleapis.com/v1/3dtiles/root.json?key={api_key}",
    attr="Google",
    name="Google Maps 3D",
    overlay=True,
    control=True,
)
google_3d_layer.add_to(m)

# Display the Folium map using streamlit-folium
folium_static(m, width=800, height=600)