database.py

from sqlalchemy import create_engine, Column, Integer, String, Float, Date, UniqueConstraint, func
from sqlalchemy.orm import declarative_base, sessionmaker
import pandas as pd

engine = create_engine("sqlite:///expenses.db", echo=False)
Session = sessionmaker(bind=engine)
Base = declarative_base()


class Expense(Base):
    __tablename__ = "expenses"

    id = Column(Integer, primary_key=True)
    date = Column(Date, nullable=False)
    description = Column(String, nullable=False)
    amount = Column(Float, nullable=False)
    category = Column(String, nullable=False)

    __table_args__ = (
        UniqueConstraint("date", "description", "amount", name="unique_expense"),
    )


Base.metadata.create_all(engine)


def insert_expenses(df: pd.DataFrame):
    session = Session()
    inserted = 0
    skipped = 0

    for _, row in df.iterrows():
        try:
            expense = Expense(
                date=row["date"],
                description=str(row["description"]),
                amount=float(row["amount"]),
                category=str(row["category"])
            )
            session.add(expense)
            session.commit()
            inserted += 1
        except Exception:
            session.rollback()
            skipped += 1

    session.close()
    return inserted, skipped


def fetch_expenses_by_category(category=None):
    session = Session()

    query = session.query(Expense)
    if category and category != "All":
        query = query.filter(Expense.category == category)

    results = query.all()
    session.close()

    return pd.DataFrame(
        [(r.date, r.description, r.amount, r.category) for r in results],
        columns=["date", "description", "amount", "category"]
    )


def get_total_spent(category=None):
    session = Session()
    query = session.query(func.sum(Expense.amount))

    if category and category != "All":
        query = query.filter(Expense.category == category)

    total = query.scalar()
    session.close()
    return total or 0


def get_category_summary(category=None):
    session = Session()

    query = session.query(
        Expense.category,
        func.sum(Expense.amount).label("total")
    )

    if category and category != "All":
        query = query.filter(Expense.category == category)

    results = query.group_by(Expense.category).all()
    session.close()

    return pd.DataFrame(results, columns=["category", "total"])


def get_monthly_summary(category=None):
    session = Session()

    query = session.query(
        func.strftime("%Y-%m", Expense.date).label("month"),
        func.sum(Expense.amount).label("total")
    )

    if category and category != "All":
        query = query.filter(Expense.category == category)

    results = query.group_by("month").all()
    session.close()

    return pd.DataFrame(results, columns=["month", "total"])

# Budget Model

class Budget(Base):
    __tablename__ = "budgets"

    id = Column(Integer, primary_key=True)
    month = Column(String, nullable=False)
    category = Column(String, nullable=False)  # "Total" or category name
    amount = Column(Float, nullable=False)

    __table_args__ = (
        UniqueConstraint("month", "category", name="unique_budget"),
    )


Base.metadata.create_all(engine)

# Budget Functions

def set_budget(month, category, amount):
    session = Session()

    try:
        budget = Budget(
            month=month,
            category=category,
            amount=amount
        )
        session.add(budget)
        session.commit()
    except:
        session.rollback()
        session.query(Budget).filter_by(
            month=month,
            category=category
        ).update({"amount": amount})
        session.commit()

    session.close()


def get_budget(month, category=None):
    session = Session()

    query = session.query(Budget).filter(Budget.month == month)

    if category:
        query = query.filter(Budget.category == category)

    results = query.all()
    session.close()

    return pd.DataFrame(
        [(r.month, r.category, r.amount) for r in results],
        columns=["month", "category", "amount"]
    )