import gradio as gr
import folium
import geopandas as gpd
import requests
import unidecode

gdf_polygons = gpd.read_parquet("gdf_polygons_deploy.parquet")

def get_polygon(gdf_polygons,x):
    polygon = gdf_polygons.query(f"(polygon_id == {x['polygon']['id']}) and (polygon_type == '{x['polygon']['type']}')")
    polygon = polygon.iloc[0].geometry
    return polygon

def standardize_user_input(user_input):
    # lowercase
    user_input = user_input.lower()

    # strip
    user_input = user_input.strip()

    # remove accents
    user_input = unidecode.unidecode(user_input)

    return user_input

def get_autocomplete(polygon_id_general: str = None, user_input: str = None):
    url = 'https://data.dev.dd360.mx/autocomplete/v1/data'
    data = {"last_search": polygon_id_general, "user_input": standardize_user_input(user_input) if user_input else None}

    try:
        response = requests.post(url, json=data, timeout=1)  
        response_json = response.json()  
        for item in response_json:
            item["polygon_id_general"] = str(item["polygon"]["type"]) + "_" + str(item["polygon"]["id"])
            item["polygon"] = get_polygon(gdf_polygons,item)
            item["latitude"] = item["centroid"]["latitude"]
            item["longitude"] = item["centroid"]["longitude"]
        return response_json  
    except requests.exceptions.Timeout:
        return []  
    except Exception:
        return []  

# Function to query DynamoDB based on user input
def query_dynamodb(user_input, last_selection):
    if user_input:
        return get_autocomplete(user_input = user_input)

    if last_selection and last_selection[0]:
        return get_autocomplete(polygon_id_general = last_selection[0]["polygon_id_general"])

    return []

# Function to update the dropdown based on user input
def update_dropdown(query, last_selection):
    suggestions = query_dynamodb(query, last_selection)
    selected_value = suggestions[0]["address_recommendation"] if suggestions else ""
    return gr.update(choices=[s["address_recommendation"] for s in suggestions], value=selected_value), suggestions

def show_map(selected_address, suggestions, last_selection):
    selected_coords = next((s for s in suggestions if s["address_recommendation"] == selected_address), None)

    if selected_coords:
        m = folium.Map(location=[selected_coords["latitude"], selected_coords["longitude"]], zoom_start=15)

        folium.Marker(
            location=[selected_coords["latitude"], selected_coords["longitude"]],
            popup=selected_coords["address_recommendation"],
            icon=folium.Icon(color='blue')
        ).add_to(m)

        # Add the polygon to the map
        if 'polygon' in selected_coords:
            folium.GeoJson(
                selected_coords['polygon'],
                style_function=lambda x: {'color': 'orange', 'fillOpacity': 0.3}
            ).add_to(m)

        map_html = m._repr_html_()

        # Update last_selection
        last_selection[0] = selected_coords  # Update the last_selection with the current selection
        return map_html, last_selection  # Return the HTML and the updated last_selection
    else:
        return None, last_selection  # Return None if no selection

with gr.Blocks() as demo:
    gr.Markdown("<h2 style='text-align: center; font-size: 40px;'>Address Autocomplete</h2>")  # Larger Title
    gr.Markdown(
        "<ul>"
        "<li>Type an address to see matching suggestions.</li>"
        "<li>Selecting one will show its polygon.</li>"
        "<li>Previous searches will also suggest addresses.</li>"
        "<li>The service supports a maximum input of 15 characters.</li>"
        "</ul>"
    ) 

    query_input = gr.Textbox(label="Type your address")
    address_dropdown = gr.Dropdown(label="Address Suggestions", choices=[], value="", interactive=True)
    map_output = gr.HTML(label="Map")  
    suggestions = gr.State([])  
    last_selection = gr.State([None])  

    query_input.change(fn=update_dropdown, inputs=[query_input, last_selection], outputs=[address_dropdown, suggestions])
    address_dropdown.change(fn=show_map, inputs=[address_dropdown, suggestions, last_selection], outputs=[map_output, last_selection])

demo.launch()