from sqlalchemy import func, select, text
from sqlalchemy.ext.asyncio import AsyncSession

from src.db.models.form_type import FormTypeEnum
from src.db.models.budget_form import BudgetForm
from sqlalchemy.orm import joinedload



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

    async def get_list(
            self, 
            offset: int | None = None, 
            limit: int | None = None,
            org_unit: int | list[int] | None = None,
            with_count: bool = False,
            load_org: bool = False,
    ) -> list[BudgetForm] | tuple[int, list[BudgetForm]]:
        if with_count:
            query = select(func.count().over().label("total_count"), BudgetForm)
        else:
            query = select(BudgetForm)
        where = []
        if org_unit is not None:
            if isinstance(org_unit, int):
                where.append(BudgetForm.org_unit_id == org_unit)
            else:
                where.append(BudgetForm.org_unit_id.in_(org_unit))

        
        query = query.where(*where)
        if load_org:
            query = query.options(
                joinedload(BudgetForm.org_unit)
            )
        query = query.order_by(BudgetForm.id)

        if offset is not None:
            query = query.offset(offset)
        if limit is not None:
            query = query.limit(limit)

        if with_count:
            result = (await self.db.execute(query)).all()
            if result:
                return result[0][0], [res[1] for res in result]
            else:
                count_query = select(func.count(BudgetForm.id))
                if query.whereclause is not None:
                    count_query = count_query.where(
                        query.whereclause
                    )
                return (await self.db.execute(count_query)).scalar(), []
        else:
            return (await self.db.execute(query)).scalars().all()
    
    async def get(
            self, 
            budget_form_id: int | None = None,
            org_unit: int | list[int] | None = None,
            form_type_code: FormTypeEnum | str | None = None,
            year: int | None = None,
            load_form_type: bool = False,
            load_org: bool = False,
    ) -> BudgetForm | None:
        assert budget_form_id or all((form_type_code, year, org_unit))
        query = select(BudgetForm)
    
        where = []
        if budget_form_id is not None:
            where.append(BudgetForm.id == budget_form_id)

        if org_unit is not None:
            if isinstance(org_unit, int):
                where.append(BudgetForm.org_unit_id == org_unit)
            else:
                where.append(BudgetForm.org_unit_id.in_(org_unit))
        
        if form_type_code is not None:
            if isinstance(form_type_code, FormTypeEnum):
                where.append(BudgetForm.form_type_code == form_type_code.value)
            else:
                where.append(BudgetForm.form_type_code == form_type_code)
        
        if year is not None:
            where.append(BudgetForm.year == year)

        query = query.where(*where)
        if load_org:
            query = query.options(
                joinedload(BudgetForm.org_unit)
            )

        if load_form_type:
            query = query.options(
                joinedload(BudgetForm.form_type)
            )
        return (await self.db.execute(query)).scalars().first()
    
    async def create(
            self, 
            form_type_code: str,
            year: int, 
            org_unit_id: int,
    ) -> BudgetForm | None:
        form_id = (
            await self.db.execute(
                text(
                    """
                    SELECT *
                    FROM v3.add_budget_form(
                        :form_type,
                        :year,
                        :org_unit_id
                    )
                    """
                ),
                {
                    "form_type": form_type_code,
                    "year": year,
                    "org_unit_id": org_unit_id,
                }
            )
        ).scalar_one_or_none()
        if form_id is None:
            return None
        return await self.get(
            budget_form_id=form_id,
        )