import os
from pathlib import Path
from datetime import date
import io

import pandas as pd
import streamlit as st
import altair as alt

# ---------- Constants ----------
DATA_DIR = Path("data")
DATA_FILE = DATA_DIR / "expenses.csv"
CATEGORIES = ["Food", "Transport", "Bills", "Shopping", "Entertainment", "Other"]

# ---------- Helpers ----------
def ensure_data_file():
    """Create data directory and CSV file with headers if missing."""
    DATA_DIR.mkdir(exist_ok=True)
    if not DATA_FILE.exists():
        df = pd.DataFrame(columns=["Date", "Category", "Description", "Amount"])
        df.to_csv(DATA_FILE, index=False)

def load_expenses() -> pd.DataFrame:
    ensure_data_file()
    try:
        df = pd.read_csv(DATA_FILE, parse_dates=["Date"]) if DATA_FILE.exists() else pd.DataFrame(columns=["Date", "Category", "Description", "Amount"]) 
        if "Date" in df.columns:
            df["Date"] = pd.to_datetime(df["Date"]).dt.date
        return df
    except Exception as e:
        st.error(f"Could not load expenses: {e}")
        return pd.DataFrame(columns=["Date", "Category", "Description", "Amount"]) 

def save_expenses(df: pd.DataFrame):
    ensure_data_file()
    df.to_csv(DATA_FILE, index=False)

def add_expense(row: dict):
    df = load_expenses()
    df = pd.concat([df, pd.DataFrame([row])], ignore_index=True)
    save_expenses(df)

def clear_expenses():
    ensure_data_file()
    df = pd.DataFrame(columns=["Date", "Category", "Description", "Amount"])
    save_expenses(df)

def get_summary(df: pd.DataFrame):
    if df.empty:
        return {"total": 0.0, "count": 0}
    total = df["Amount"].sum()
    count = len(df)
    return {"total": total, "count": count}

def csv_download_button(df: pd.DataFrame, filename: str = "expenses.csv"):
    buffer = io.StringIO()
    df.to_csv(buffer, index=False)
    buffer.seek(0)
    st.download_button("Download CSV", buffer, file_name=filename, mime="text/csv")

# ---------- App UI ----------
st.set_page_config(page_title="Expense Tracker", page_icon="💸", layout="centered")
st.title("💸 Simple Expense Tracker")
st.write("Add and track your expenses. Data persists to `data/expenses.csv` on the Space.")

# Load data
df = load_expenses()

# Form to add expense
with st.form("expense_form", clear_on_submit=True):
    col1, col2 = st.columns(2)
    with col1:
        expense_date = st.date_input("Date", value=date.today())
        category = st.selectbox("Category", options=CATEGORIES)
    with col2:
        amount = st.number_input("Amount", min_value=0.0, format="%.2f")
        description = st.text_input("Description")

    submitted = st.form_submit_button("Add expense")

    if submitted:
        if amount <= 0:
            st.warning("Amount must be greater than 0.")
        else:
            row = {
                "Date": expense_date.isoformat(),
                "Category": category,
                "Description": description,
                "Amount": float(amount),
            }
            add_expense(row)
            st.success("Expense added!")
            df = load_expenses()

# Sidebar actions
st.sidebar.header("Actions")
if st.sidebar.button("Refresh"):
    df = load_expenses()

if st.sidebar.button("Clear all expenses"):
    if st.sidebar.checkbox("Confirm clear all (check to confirm)"):
        clear_expenses()
        st.sidebar.success("All expenses cleared")
        df = load_expenses()

# Display table and summary
st.subheader("Expenses")
if df.empty:
    st.info("No expenses yet — add one using the form above.")
else:
    df["Amount"] = pd.to_numeric(df["Amount"], errors="coerce").fillna(0.0)
    if "Date" in df.columns:
        try:
            df["Date"] = pd.to_datetime(df["Date"]).dt.date
        except Exception:
            pass

    st.dataframe(df.sort_values(by="Date", ascending=False))

    summary = get_summary(df)
    st.markdown(f"**Total:** {summary['total']:.2f}   —   **Count:** {summary['count']}")

    # Chart: total by category
    st.subheader("Spending by category")
    cat_df = df.groupby("Category", as_index=False)["Amount"].sum()
    if not cat_df.empty:
        chart = alt.Chart(cat_df).mark_bar().encode(
            x=alt.X("Category:N", sort="-y"),
            y=alt.Y("Amount:Q"),
            tooltip=[alt.Tooltip("Category"), alt.Tooltip("Amount:Q", format=".2f")],
        )
        st.altair_chart(chart, use_container_width=True)

    # Monthly cumulative line
    st.subheader("Monthly spending (cumulative)")
    df_line = df.copy()
    df_line["Date"] = pd.to_datetime(df_line["Date"])
    df_line["YearMonth"] = df_line["Date"].dt.to_period("M").dt.to_timestamp()
    monthly = df_line.groupby("YearMonth")["Amount"].sum().reset_index()
    if not monthly.empty:
        line = alt.Chart(monthly).mark_line(point=True).encode(
            x="YearMonth:T",
            y="Amount:Q",
            tooltip=[alt.Tooltip("YearMonth:T", title="Month"), alt.Tooltip("Amount:Q", format=".2f")],
        )
        st.altair_chart(line, use_container_width=True)

# Download CSV
st.subheader("Export")
if not df.empty:
    csv_download_button(df)

st.caption("Data is stored locally in the app's `data/expenses.csv`. On Hugging Face Spaces this file persists across runs unless you explicitly remove it.")