from sqlalchemy.ext.asyncio import AsyncSession

from src.repository.user_repository import UserRepository
from src.db.models.app_user import AppUser
from src.db.models.budget_form import BudgetForm
from src.db.models.role import UserRoleEnum
from src.repository.budget_form_repository import BudgetFormRepository



class BudgetFormService:
    def __init__(self, db: AsyncSession):
        self.db = db
        self.bf_repo = BudgetFormRepository(db)
        self.user_repo = UserRepository(db)

    async def get_list(
            self,
            user: AppUser,
            offset: int | None = None,
            limit: int | None = None,
            with_count: bool = False,
            load_org: bool = False,
    ) -> list[BudgetForm] | tuple[int, list[BudgetForm]]:
        if user.role_id == UserRoleEnum.ADMIN:
            return await self.bf_repo.get_list(
                offset=offset,
                limit=limit,
                with_count=with_count,
                load_org=load_org,
            )
        user = await self.user_repo.get(
            user_id=user.id,
            load_orgs=True,
        )
        return await self.bf_repo.get_list(
            offset=offset,
            limit=limit,
            org_unit=[ou.id for ou in user.org_units],
            with_count=with_count,
            load_org=load_org,
        )

    async def get(
            self,
            budget_form_id: int,
            user: AppUser,
            load_form_type: bool = False,
    ) -> list[BudgetForm] | tuple[int, list[BudgetForm]]:
        if user.role_id == UserRoleEnum.ADMIN:
            return await self.bf_repo.get(
                budget_form_id=budget_form_id,
                load_form_type=load_form_type,
            )
        user = await self.user_repo.get(
            user_id=user.id,
            load_orgs=True,
        )
        return await self.bf_repo.get(
            budget_form_id=budget_form_id,
            org_unit=[ou.id for ou in user.org_units],
            load_form_type=load_form_type,
        )