import streamlit as st
import pandas as pd
import plotly.express as px
import pydeck as pdk
from salesforce_integration import fetch_poles
from modules.visuals import display_dashboard, display_charts

# Title
st.title("VIEP Smart Poles Dashboard")

# Fetch the raw data from Salesforce
df = fetch_poles()

# --- Sidebar Filters ---
st.sidebar.header("Filter Data")
selected_alert_levels = st.sidebar.multiselect(
    "Alert Level", ["Red", "Yellow", "Green"], default=["Red", "Yellow", "Green"]
)
selected_camera_status = st.sidebar.selectbox("Camera Status", ["All", "Online", "Offline"])
site_options = ["All"] + df["Site__c"].dropna().unique().tolist()
selected_site = st.sidebar.selectbox("Site", site_options, index=0)

# --- Filtering Logic ---
filtered_df = df[df["Alert_Level__c"].isin(selected_alert_levels)]
if selected_camera_status != "All":
    filtered_df = filtered_df[filtered_df["Camera_Status__c"] == selected_camera_status]
if selected_site != "All":
    filtered_df = filtered_df[filtered_df["Site__c"] == selected_site]

# --- Display System Summary ---
display_dashboard(filtered_df)

# --- Pole Table ---
st.subheader("Pole Table")
st.dataframe(filtered_df, use_container_width=True)

# --- Energy Generation Chart ---
st.subheader("⚙ Energy Generation (Solar vs Wind)")
st.plotly_chart(px.bar(
    filtered_df, x="Name", y=["Solar_Generation__c", "Wind_Generation__c"], barmode="group"
))

# --- Alert Level Breakdown Chart ---
display_charts(filtered_df)

# Define function to generate heatmap based on site
def generate_heatmap_for_site(site_name, df):
    site_df = df[df['Site__c'] == site_name]

    # Ensure that Alert_Level__c is treated as a string (for color mapping)
    site_df['Alert_Level__c'] = site_df['Alert_Level__c'].astype(str)

    # Define color mapping for alert levels
    color_map = {
        "Green": [0, 255, 0],
        "Yellow": [255, 255, 0],
        "Red": [255, 0, 0]
    }

    # Create a color column based on Alert_Level__c
    site_df["color"] = site_df["Alert_Level__c"].map(color_map)

    # Create a Pydeck map for the site
    layer = pdk.Layer(
        "ScatterplotLayer",
        data=site_df,
        get_position='[Location_Longitude__c, Location_Latitude__c]',  # Ensure the column names are correct
        get_color="color",
        get_radius=80,  # You can adjust the radius if needed
        pickable=True,
        auto_highlight=True
    )

    # Handle potential missing coordinates for map center
    latitude = site_df["Location_Latitude__c"].mean() if not site_df["Location_Latitude__c"].isnull().all() else 0
    longitude = site_df["Location_Longitude__c"].mean() if not site_df["Location_Longitude__c"].isnull().all() else 0

    # View state for the map
    view_state = pdk.ViewState(
        latitude=latitude,
        longitude=longitude,
        zoom=10,
        pitch=40
    )

    tooltip = {
        "html": """
            <b>Pole Name:</b> {Name}<br>
            <b>Site:</b> {Site__c}<br>
            <b>Alert Level:</b> {Alert_Level__c}<br>
            <b>RFID Tag:</b> {RFID_Tag__c}<br>
            <b>Tilt:</b> {Tilt__c}<br>
            <b>Vibration:</b> {Vibration__c}
        """,
        "style": {
            "backgroundColor": "steelblue",
            "color": "white"
        }
    }

    # Return the heatmap
    return pdk.Deck(
        map_style="mapbox://styles/mapbox/dark-v10",
        initial_view_state=view_state,
        layers=[layer],
        tooltip=tooltip
    )

# Divide into four columns (Hyderabad, Kurnool, Ballari, Gadwal)
col1, col2 = st.columns(2)
with col1:
    st.subheader("Hyderabad")
    st.pydeck_chart(generate_heatmap_for_site("Hyderabad", df))
    st.subheader("Kurnool")
    st.pydeck_chart(generate_heatmap_for_site("Kurnool", df))

with col2:
    st.subheader("Ballari")
    st.pydeck_chart(generate_heatmap_for_site("Ballari", df))
    st.subheader("Gadwal")
    st.pydeck_chart(generate_heatmap_for_site("Gadwal", df))