from __future__ import annotations

from sqlalchemy import select

from app.models.budget import BudgetAllocation, BudgetTransaction
from app.schemas.budget import BudgetAllocationRead, BudgetTransactionRead
from app.services.budget_types import SUPPORTED_BUDGET_SUBJECT_CODES
from app.services.pagination import PageResult, paginate_select


class BudgetPaginationMixin:
    def build_allocation_stmt(
        self,
        *,
        fiscal_year: int | None = None,
        period_key: str | None = None,
        department_id: str | None = None,
        department_name: str | None = None,
        cost_center: str | None = None,
    ):
        stmt = select(BudgetAllocation).where(
            BudgetAllocation.subject_code.in_(SUPPORTED_BUDGET_SUBJECT_CODES)
        )
        if fiscal_year is not None:
            stmt = stmt.where(BudgetAllocation.fiscal_year == fiscal_year)
        if period_key:
            stmt = stmt.where(BudgetAllocation.period_key == period_key)
        if department_id:
            stmt = stmt.where(BudgetAllocation.department_id == department_id)
        if department_name:
            stmt = stmt.where(BudgetAllocation.department_name == department_name)
        if cost_center:
            stmt = stmt.where(BudgetAllocation.cost_center == cost_center)
        return stmt.order_by(
            BudgetAllocation.fiscal_year.desc(),
            BudgetAllocation.period_key.asc(),
            BudgetAllocation.department_name.asc(),
            BudgetAllocation.subject_code.asc(),
        )

    def list_allocations_page(
        self,
        *,
        fiscal_year: int | None = None,
        period_key: str | None = None,
        department_id: str | None = None,
        department_name: str | None = None,
        cost_center: str | None = None,
        page: int | None,
        page_size: int | None,
    ) -> PageResult[BudgetAllocationRead]:
        self.ensure_budget_ready()
        result = paginate_select(
            self.db,
            self.build_allocation_stmt(
                fiscal_year=fiscal_year,
                period_key=period_key,
                department_id=department_id,
                department_name=department_name,
                cost_center=cost_center,
            ),
            page=page,
            page_size=page_size,
        )
        return result.map(self.serialize_allocation)

    def list_transactions_page(
        self,
        allocation_id: str,
        *,
        page: int | None,
        page_size: int | None,
    ) -> PageResult[BudgetTransactionRead]:
        self.ensure_budget_ready()
        result = paginate_select(
            self.db,
            select(BudgetTransaction)
            .where(BudgetTransaction.allocation_id == allocation_id)
            .order_by(BudgetTransaction.created_at.desc()),
            page=page,
            page_size=page_size,
        )
        return result.map(BudgetTransactionRead.model_validate)
