import streamlit as st
import pandas as pd
import matplotlib.pyplot as plt
from datetime import date

from llm_utils import analyze_expenses_with_llm

st.set_page_config(page_title="AI Expense Analyzer", layout="wide")

st.title("💸 AI Expense Analyzer")
st.write("Track expenses and get AI-powered insights using Grok")

# -------------------------
# Session State
# -------------------------
if "expenses" not in st.session_state:
    st.session_state.expenses = pd.DataFrame(
        columns=["Date", "Category", "Amount", "Description"]
    )

# -------------------------
# Expense Form
# -------------------------
st.subheader("➕ Add Expense")

with st.form("expense_form"):
    expense_date = st.date_input("Date", value=date.today())

    category = st.selectbox(
        "Category",
        ["Food", "Transport", "Shopping", "Bills", "Entertainment", "Other"]
    )

    amount = st.number_input("Amount", min_value=0.0, step=1.0)

    description = st.text_input("Description")

    submitted = st.form_submit_button("Add Expense")

    if submitted:
        new_row = {
            "Date": expense_date,
            "Category": category,
            "Amount": amount,
            "Description": description,
        }

        st.session_state.expenses = pd.concat(
            [st.session_state.expenses, pd.DataFrame([new_row])],
            ignore_index=True
        )

        st.success("Expense added successfully!")

# -------------------------
# Expense Table
# -------------------------
st.subheader("📋 Expense History")
st.dataframe(st.session_state.expenses, use_container_width=True)

if st.session_state.expenses.empty:
    st.info("Add expenses to see charts and AI analysis.")
    st.stop()

# -------------------------
# Charts
# -------------------------
st.subheader("📊 Expense Analysis")

col1, col2 = st.columns(2)

with col1:
    category_sum = st.session_state.expenses.groupby("Category")["Amount"].sum()
    fig, ax = plt.subplots()
    category_sum.plot(kind="bar", ax=ax)
    st.pyplot(fig)

with col2:
    fig2, ax2 = plt.subplots()
    ax2.pie(category_sum, labels=category_sum.index, autopct="%1.1f%%")
    ax2.axis("equal")
    st.pyplot(fig2)

# -------------------------
# Trend Chart
# -------------------------
st.write("### Spending Trend")
trend_df = st.session_state.expenses.copy()
trend_df["Date"] = pd.to_datetime(trend_df["Date"])
trend_df = trend_df.groupby("Date")["Amount"].sum()

fig3, ax3 = plt.subplots()
ax3.plot(trend_df.index, trend_df.values, marker="o")
st.pyplot(fig3)

# -------------------------
# AI Analysis
# -------------------------
st.subheader("🤖 AI Expense Insights")

if st.button("Analyze My Expenses"):
    with st.spinner("Analyzing your expenses..."):
        insights = analyze_expenses_with_llm(st.session_state.expenses)
        st.markdown(insights)