"""Budget repository"""

from typing import List, Optional
from sqlalchemy.orm import Session
from app.db.models import Budget
from app.db.repositories.base import BaseRepository


class BudgetRepository(BaseRepository[Budget]):
    """Budget data access operations"""

    def __init__(self, session: Session):
        super().__init__(session, Budget)

    def get_family_budgets(self, family_id: int) -> List[Budget]:
        """Get all active budgets for family"""
        return (
            self.session.query(Budget)
            .filter(Budget.family_id == family_id, Budget.is_active == True)
            .all()
        )

    def get_category_budget(self, family_id: int, category_id: int) -> Optional[Budget]:
        """Get budget for specific category"""
        return (
            self.session.query(Budget)
            .filter(
                Budget.family_id == family_id,
                Budget.category_id == category_id,
                Budget.is_active == True
            )
            .first()
        )

    def get_general_budget(self, family_id: int) -> Optional[Budget]:
        """Get general budget (no category)"""
        return (
            self.session.query(Budget)
            .filter(
                Budget.family_id == family_id,
                Budget.category_id == None,
                Budget.is_active == True
            )
            .first()
        )

    def update_spent_amount(self, budget_id: int, amount: float) -> Optional[Budget]:
        """Update spent amount for budget"""
        budget = self.get_by_id(budget_id)
        if budget:
            budget.spent_amount += amount
            self.session.commit()
            self.session.refresh(budget)
        return budget