
import asyncio

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,
            load_org: bool = False,
    ) -> BudgetForm | None:
        if user.role_id == UserRoleEnum.ADMIN:
            return await self.bf_repo.get(
                budget_form_id=budget_form_id,
                load_form_type=load_form_type,
                load_org=load_org,
            )
        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,
            load_org=load_org,
        )

    async def create(
            self, 
            form_type_code: str,
            year: int, 
            org_unit_id: int,
    ) -> BudgetForm | None:
        return await self.bf_repo.create(
            form_type_code=form_type_code,
            year=year,
            org_unit_id=org_unit_id,
        )

    async def create_if_not_exists(
            self, 
            form_type_code: str,
            year: int, 
            org_unit_id: int,
    ) -> BudgetForm | None:
        existing_bf = await self.bf_repo.get(
            form_type_code=form_type_code,
            year=year,
            org_unit=org_unit_id,
        )
        if existing_bf:
            return existing_bf
        return await self.bf_repo.create(
            form_type_code=form_type_code,
            year=year,
            org_unit_id=org_unit_id,
        )
    
    async def bulk_create(
            self, 
            form_type_code: str,
            year: int, 
            org_unit_ids: list[int],
    ) -> BudgetForm | None:
        return await asyncio.gather(
            *[
                self.create(
                    form_type_code=form_type_code,
                    year=year,
                    org_unit_id=org_unit_id,
                )
                for org_unit_id in org_unit_ids
            ] 
        )
    
    async def bulk_create_if_not_exists(
            self, 
            form_type_code: str,
            year: int, 
            org_unit_ids: list[int],
    ) -> BudgetForm | None:
        return await asyncio.gather(
            *[
                self.create_if_not_exists(
                    form_type_code=form_type_code,
                    year=year,
                    org_unit_id=org_unit_id,
                )
                for org_unit_id in org_unit_ids
            ] 
        )
