from datetime import date, timedelta

import altair as alt
import pandas as pd
import streamlit as st
from geopy.extra.rate_limiter import RateLimiter
from geopy.geocoders import Nominatim

st.set_page_config(
    page_title="Weather dashboard",
    page_icon=":material/thermostat:",
    layout="wide",
)

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
TIME_RANGES = ["1M", "3M", "6M", "1Y", "YTD", "All"]
VARIABLES = ["Temperature", "Wind speed", "Wind gusts", "Precipitation"]
VAR_COLS = {
    "Temperature": "temperature_2m",
    "Wind speed": "windspeed_10m",
    "Wind gusts": "windgusts_10m",
    "Precipitation": "precipitation",
}
VAR_UNITS = {
    "Temperature": "F",
    "Wind speed": "mph",
    "Wind gusts": "mph",
    "Precipitation": "in",
}


# ---------------------------------------------------------------------------
# Data helpers
# ---------------------------------------------------------------------------
@st.cache_data(show_spinner=False)
def geocode(address: str) -> tuple[float, float]:
    """Return (lat, lon) for *address*, trying Census first then Nominatim."""
    try:
        address2 = address.replace(" ", "+").replace(",", "%2C")
        url = (
            "https://geocoding.geo.census.gov/geocoder/locations/onelineaddress"
            f"?address={address2}&benchmark=2020&format=json"
        )
        df = pd.read_json(url)
        coords = df.iloc[:1, 0][0][0]["coordinates"]
        return coords["y"], coords["x"]
    except Exception:
        geolocator = Nominatim(user_agent="WeatherDashboard")
        geocode_fn = RateLimiter(geolocator.geocode, min_delay_seconds=1)
        location = geocode_fn(address)
        if location is None:
            raise ValueError(f"Could not geocode: {address}")
        return location.latitude, location.longitude


@st.cache_data(show_spinner=False, ttl=900)
def get_weather_data(
    lat: float, lon: float, start_date: str, end_date: str
) -> pd.DataFrame:
    """Fetch hourly weather from Open-Meteo archive API."""
    url = (
        f"https://archive-api.open-meteo.com/v1/archive"
        f"?latitude={lat}&longitude={lon}"
        f"&start_date={start_date}&end_date={end_date}"
        f"&hourly=temperature_2m,precipitation,windspeed_10m,windgusts_10m"
        f"&models=best_match"
        f"&temperature_unit=fahrenheit&windspeed_unit=mph&precipitation_unit=inch"
    )
    raw = pd.read_json(url).reset_index()
    data = pd.DataFrame({c["index"]: c["hourly"] for _, c in raw.iterrows()})
    data["time"] = pd.to_datetime(data["time"])
    data = data.dropna(subset=["temperature_2m"])
    return data


def aggregate_daily(df: pd.DataFrame) -> pd.DataFrame:
    """Compute daily aggregates from hourly data."""
    df = df.copy()
    df["date"] = df["time"].dt.date
    agg = df.groupby("date").agg(
        temperature_2m_min=("temperature_2m", "min"),
        temperature_2m_mean=("temperature_2m", "mean"),
        temperature_2m_max=("temperature_2m", "max"),
        precipitation_sum=("precipitation", "sum"),
        windspeed_10m_min=("windspeed_10m", "min"),
        windspeed_10m_mean=("windspeed_10m", "mean"),
        windspeed_10m_max=("windspeed_10m", "max"),
        windgusts_10m_min=("windgusts_10m", "min"),
        windgusts_10m_mean=("windgusts_10m", "mean"),
        windgusts_10m_max=("windgusts_10m", "max"),
    )
    agg.index = pd.to_datetime(agg.index)
    agg.index.name = "date"
    return agg


def filter_by_time_range(
    df: pd.DataFrame, x_col: str, time_range: str
) -> pd.DataFrame:
    """Filter dataframe by a preset time range."""
    if time_range == "All" or df.empty:
        return df
    df = df.copy()
    df[x_col] = pd.to_datetime(df[x_col])
    max_date = df[x_col].max()
    if time_range == "1M":
        min_date = max_date - timedelta(days=30)
    elif time_range == "3M":
        min_date = max_date - timedelta(days=90)
    elif time_range == "6M":
        min_date = max_date - timedelta(days=180)
    elif time_range == "1Y":
        min_date = max_date - timedelta(days=365)
    elif time_range == "YTD":
        min_date = pd.Timestamp(date(max_date.year, 1, 1))
    else:
        return df
    return df[df[x_col] >= min_date]


@st.cache_data
def to_csv(df: pd.DataFrame) -> bytes:
    return df.to_csv(index=True).encode("utf-8")


# ---------------------------------------------------------------------------
# Sidebar
# ---------------------------------------------------------------------------
with st.sidebar:
    st.header("Settings")

    address = st.text_input(
        "Address",
        value="1000 Main St, Cincinnati, OH 45202",
        placeholder="Enter an address...",
    )
    col_s, col_e = st.columns(2)
    with col_s:
        start_date = st.date_input("Start date", pd.Timestamp(2024, 1, 1))
    with col_e:
        end_date = st.date_input("End date", pd.Timestamp(2025, 11, 10))

    variable = st.selectbox("Variable", VARIABLES, index=2)

    st.caption("Data from Open-Meteo archive API")

# ---------------------------------------------------------------------------
# Geocode + fetch
# ---------------------------------------------------------------------------
try:
    lat, lon = geocode(address)
except Exception:
    st.error(
        "Could not find that address. Please check the spelling and try again.",
        icon=":material/error:",
    )
    st.stop()

start_str = start_date.strftime("%Y-%m-%d")
end_str = end_date.strftime("%Y-%m-%d")

with st.spinner("Fetching weather data..."):
    try:
        hourly = get_weather_data(lat, lon, start_str, end_str)
    except Exception as exc:
        st.error(f"Failed to fetch weather data: {exc}", icon=":material/error:")
        st.stop()

daily = aggregate_daily(hourly)

# ---------------------------------------------------------------------------
# Header
# ---------------------------------------------------------------------------
col_var = VAR_COLS[variable]
unit = VAR_UNITS[variable]

st.markdown("# :material/thermostat: Weather dashboard")
st.caption(f"{address}  ({lat:.4f}, {lon:.4f})")

# ---------------------------------------------------------------------------
# KPI metrics row
# ---------------------------------------------------------------------------
if variable == "Precipitation":
    total_precip = daily["precipitation_sum"].sum()
    avg_daily = daily["precipitation_sum"].mean()
    max_daily = daily["precipitation_sum"].max()
    dry_days = int((daily["precipitation_sum"] < 0.01).sum())

    k1, k2, k3, k4 = st.columns(4)
    k1.metric("Total precipitation", f"{total_precip:.1f} {unit}", border=True)
    k2.metric("Avg daily", f"{avg_daily:.2f} {unit}", border=True)
    k3.metric("Max daily", f"{max_daily:.2f} {unit}", border=True)
    k4.metric("Dry days", f"{dry_days:,}", border=True)
else:
    mean_col = f"{col_var}_mean"
    min_col = f"{col_var}_min"
    max_col = f"{col_var}_max"

    overall_mean = daily[mean_col].mean()
    overall_min = daily[min_col].min()
    overall_max = daily[max_col].max()

    k1, k2, k3, k4 = st.columns(4)
    k1.metric(
        f"Avg {variable.lower()}",
        f"{overall_mean:.1f} {unit}",
        border=True,
    )
    k2.metric(f"Min {variable.lower()}", f"{overall_min:.1f} {unit}", border=True)
    k3.metric(f"Max {variable.lower()}", f"{overall_max:.1f} {unit}", border=True)
    k4.metric("Days of data", f"{len(daily):,}", border=True)

# ---------------------------------------------------------------------------
# Filters row
# ---------------------------------------------------------------------------
with st.popover("Filters", icon=":material/filter_list:"):
    time_range = st.segmented_control("Time range", TIME_RANGES, default="All")
    agg_mode = st.segmented_control(
        "Aggregation", ["Hourly", "Daily"], default="Daily"
    )
    if variable != "Precipitation":
        show_range = st.toggle("Show min/max range", value=True)
    else:
        show_range = False

# ---------------------------------------------------------------------------
# Build filtered data
# ---------------------------------------------------------------------------
if agg_mode == "Hourly":
    chart_df = hourly[["time", col_var]].copy()
    chart_df = filter_by_time_range(chart_df, "time", time_range)
    x_field = "time"
else:
    chart_df = daily.reset_index().copy()
    chart_df = filter_by_time_range(chart_df, "date", time_range)
    x_field = "date"

# ---------------------------------------------------------------------------
# Main charts row
# ---------------------------------------------------------------------------
col1, col2 = st.columns([3, 1])

with col1:
    with st.container(border=True):
        st.markdown(f"**{variable} over time**")

        if agg_mode == "Hourly":
            chart = (
                alt.Chart(chart_df)
                .mark_line(strokeWidth=1.5)
                .encode(
                    x=alt.X("time:T", title="Date"),
                    y=alt.Y(f"{col_var}:Q", title=f"{variable} ({unit})"),
                    tooltip=[
                        alt.Tooltip("time:T", title="Time"),
                        alt.Tooltip(f"{col_var}:Q", title=variable, format=".1f"),
                    ],
                )
                .properties(height=380)
            )
            st.altair_chart(chart, use_container_width=True)

        elif variable == "Precipitation":
            chart = (
                alt.Chart(chart_df)
                .mark_bar(color="#4B9CD3")
                .encode(
                    x=alt.X("date:T", title="Date"),
                    y=alt.Y("precipitation_sum:Q", title=f"Daily total ({unit})"),
                    tooltip=[
                        alt.Tooltip("date:T", title="Date"),
                        alt.Tooltip(
                            "precipitation_sum:Q",
                            title="Precipitation",
                            format=".2f",
                        ),
                    ],
                )
                .properties(height=380)
            )
            st.altair_chart(chart, use_container_width=True)

        else:
            mean_c = f"{col_var}_mean"
            min_c = f"{col_var}_min"
            max_c = f"{col_var}_max"

            line = (
                alt.Chart(chart_df)
                .mark_line(strokeWidth=2)
                .encode(
                    x=alt.X("date:T", title="Date"),
                    y=alt.Y(f"{mean_c}:Q", title=f"{variable} ({unit})"),
                    tooltip=[
                        alt.Tooltip("date:T", title="Date"),
                        alt.Tooltip(f"{min_c}:Q", title="Min", format=".1f"),
                        alt.Tooltip(f"{mean_c}:Q", title="Mean", format=".1f"),
                        alt.Tooltip(f"{max_c}:Q", title="Max", format=".1f"),
                    ],
                )
            )

            if show_range:
                band = (
                    alt.Chart(chart_df)
                    .mark_area(opacity=0.15)
                    .encode(
                        x=alt.X("date:T"),
                        y=alt.Y(f"{min_c}:Q"),
                        y2=alt.Y2(f"{max_c}:Q"),
                    )
                )
                chart = (band + line).properties(height=380)
            else:
                chart = line.properties(height=380)

            st.altair_chart(chart, use_container_width=True)

with col2:
    with st.container(border=True):
        st.markdown("**Monthly summary**")

        monthly = hourly.copy()
        monthly["month"] = monthly["time"].dt.to_period("M").astype(str)

        if variable == "Precipitation":
            monthly_agg = (
                monthly.groupby("month")["precipitation"].sum().reset_index()
            )
            monthly_agg.columns = ["month", "value"]
        else:
            monthly_agg = monthly.groupby("month")[col_var].mean().reset_index()
            monthly_agg.columns = ["month", "value"]

        bar = (
            alt.Chart(monthly_agg)
            .mark_bar()
            .encode(
                x=alt.X("month:O", title="Month", axis=alt.Axis(labelAngle=-45)),
                y=alt.Y(
                    "value:Q",
                    title=f"{'Total' if variable == 'Precipitation' else 'Avg'} ({unit})",
                ),
                tooltip=[
                    alt.Tooltip("month:O", title="Month"),
                    alt.Tooltip("value:Q", title=variable, format=".1f"),
                ],
            )
            .properties(height=380)
        )
        st.altair_chart(bar, use_container_width=True)

# ---------------------------------------------------------------------------
# Bottom section: distribution + data table
# ---------------------------------------------------------------------------
col_left, col_right = st.columns(2)

with col_left:
    with st.container(border=True):
        st.markdown("**Distribution**")

        if agg_mode == "Hourly":
            hist_data = chart_df[col_var].dropna()
        else:
            if variable == "Precipitation":
                hist_data = chart_df["precipitation_sum"].dropna()
            else:
                hist_data = chart_df[f"{col_var}_mean"].dropna()

        hist_df = pd.DataFrame({"value": hist_data})
        hist = (
            alt.Chart(hist_df)
            .mark_bar()
            .encode(
                x=alt.X(
                    "value:Q",
                    bin=alt.Bin(maxbins=40),
                    title=f"{variable} ({unit})",
                ),
                y=alt.Y("count()", title="Frequency"),
                tooltip=[
                    alt.Tooltip(
                        "value:Q", bin=alt.Bin(maxbins=40), title=variable
                    ),
                    alt.Tooltip("count()", title="Count"),
                ],
            )
            .properties(height=280)
        )
        st.altair_chart(hist, use_container_width=True)

with col_right:
    with st.container(border=True):
        st.markdown("**Raw data**")

        display_df = chart_df.copy()

        st.dataframe(
            display_df,
            height=280,
            hide_index=True,
            column_config={
                "date": st.column_config.DateColumn(
                    "Date", format="MMM DD, YYYY"
                ),
                "time": st.column_config.DatetimeColumn(
                    "Time", format="MMM DD, YYYY HH:mm"
                ),
            },
        )

        csv = to_csv(display_df)
        st.download_button(
            label="Download CSV",
            data=csv,
            file_name=f"weather_{start_str}_to_{end_str}.csv",
            mime="text/csv",
            icon=":material/download:",
        )