from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from src.db.models.budget_line import BudgetLine


class BudgetLineRepository:
    def __init__(self, db: AsyncSession):
        self.db = db

    async def get(self, budget_line_id: int) -> BudgetLine | None:
        query = select(BudgetLine).where(BudgetLine.id == budget_line_id).limit(1)
        return (await self.db.execute(query)).scalar_one_or_none()

    async def get_list(self, budget_line_ids: list[int]) -> list[BudgetLine]:
        query = select(BudgetLine).where(BudgetLine.id.in_(budget_line_ids))
        return (await self.db.execute(query)).scalars().all()

    async def get_ids(
            self,
            budget_form_id: int,
    ) -> list[int]:
        
        query = select(BudgetLine.id).where(BudgetLine.budget_form_id == budget_form_id)
        return (await self.db.execute(query)).scalars().all()
        
