"""Budget storage and period helpers."""

from __future__ import annotations

from .shared import Budget, Optional, and_, datetime, desc, select


class BudgetEnforcementStorageMixin:
    """Budget lookup and period-window management helpers."""

    async def _get_active_budget(self, tenant_id: int, project_id: Optional[int]) -> Optional[Budget]:
        """优先读取项目预算，否则回退租户预算。"""
        if project_id is not None:
            result = await self.db.execute(
                select(Budget)
                .where(
                    and_(
                        Budget.tenant_id == tenant_id,
                        Budget.budget_type == "project",
                        Budget.project_id == project_id,
                    )
                )
                .order_by(desc(Budget.created_at))
            )
            project_budget = result.scalar_one_or_none()
            if project_budget:
                return project_budget

        result = await self.db.execute(
            select(Budget)
            .where(
                and_(
                    Budget.tenant_id == tenant_id,
                    Budget.budget_type == "tenant",
                )
            )
            .order_by(desc(Budget.created_at))
        )
        return result.scalar_one_or_none()

    async def _ensure_current_period(self, budget: Budget, now: datetime) -> None:
        """若预算周期已过，自动切换到新周期并清零使用量。"""
        start = budget.current_period_start
        end = budget.current_period_end
        if start is None or end is None or now >= end:
            next_start, next_end = self._compute_period_window(now, budget.period)
            budget.current_period_start = next_start
            budget.current_period_end = next_end
            budget.current_usage_usd = 0.0
            budget.status = "active"
            await self.db.commit()

    def _compute_period_window(self, now: datetime, period: str) -> tuple[datetime, datetime]:
        period = (period or "monthly").lower()
        if period == "yearly":
            start = datetime(now.year, 1, 1)
            end = datetime(now.year + 1, 1, 1)
            return start, end
        if period == "quarterly":
            quarter = (now.month - 1) // 3
            start_month = quarter * 3 + 1
            start = datetime(now.year, start_month, 1)
            if start_month == 10:
                end = datetime(now.year + 1, 1, 1)
            else:
                end = datetime(now.year, start_month + 3, 1)
            return start, end

        start = datetime(now.year, now.month, 1)
        if now.month == 12:
            end = datetime(now.year + 1, 1, 1)
        else:
            end = datetime(now.year, now.month + 1, 1)
        return start, end