"""Budget check/update flows."""

from __future__ import annotations

from .shared import Any, Optional, now_cn_naive


class BudgetEnforcementChecksMixin:
    """Pre-call checks and post-call usage updates."""

    async def check_budget_before_ai_call(
        self,
        tenant_id: int,
        estimated_cost_usd: float,
        project_id: Optional[int] = None,
    ) -> dict[str, Any]:
        """AI 调用前预算检查。"""
        budget = await self._get_active_budget(tenant_id=tenant_id, project_id=project_id)
        if not budget:
            return {
                "allowed": True,
                "reason": "no_budget_configured",
                "budget_checked": False,
            }

        now = now_cn_naive()
        await self._ensure_current_period(budget, now)

        current_usage = float(budget.current_usage_usd or 0.0)
        budget_amount = float(budget.budget_amount_usd or 0.0)
        if budget_amount <= 0:
            return {
                "allowed": True,
                "reason": "invalid_budget_amount",
                "budget_checked": True,
            }

        if budget.status == "suspended":
            return {
                "allowed": False,
                "reason": "budget_suspended",
                "budget_checked": True,
                "budget_id": budget.id,
                "current_usage_usd": round(current_usage, 4),
                "budget_amount_usd": round(budget_amount, 4),
            }

        projected_usage = current_usage + max(estimated_cost_usd, 0.0)
        usage_percent = (projected_usage / budget_amount) * 100

        if usage_percent >= float(budget.alert_threshold_percent or 80.0):
            await self.generate_budget_alert(
                budget=budget,
                alert_type="threshold_reached",
                usage_usd=projected_usage,
                usage_percent=usage_percent,
                message=(
                    f"预算使用率达到 {usage_percent:.2f}% "
                    f"(阈值 {float(budget.alert_threshold_percent or 80.0):.2f}%)"
                ),
                action_taken=None,
            )

        if projected_usage > budget_amount and bool(budget.hard_limit_enabled) and budget.status != "overridden":
            await self.enforce_spending_limit(
                budget=budget,
                usage_usd=projected_usage,
                usage_percent=usage_percent,
            )
            return {
                "allowed": False,
                "reason": "budget_exceeded_hard_limit",
                "budget_checked": True,
                "budget_id": budget.id,
                "current_usage_usd": round(current_usage, 4),
                "projected_usage_usd": round(projected_usage, 4),
                "budget_amount_usd": round(budget_amount, 4),
                "usage_percent": round(usage_percent, 2),
            }

        return {
            "allowed": True,
            "reason": "within_budget",
            "budget_checked": True,
            "budget_id": budget.id,
            "current_usage_usd": round(current_usage, 4),
            "projected_usage_usd": round(projected_usage, 4),
            "budget_amount_usd": round(budget_amount, 4),
            "usage_percent": round(usage_percent, 2),
        }

    async def update_budget_usage_after_ai_call(
        self,
        tenant_id: int,
        actual_cost_usd: float,
        project_id: Optional[int] = None,
    ) -> dict[str, Any]:
        """AI 调用成功后累计预算使用量。"""
        budget = await self._get_active_budget(tenant_id=tenant_id, project_id=project_id)
        if not budget:
            return {"updated": False, "reason": "no_budget_configured"}

        now = now_cn_naive()
        await self._ensure_current_period(budget, now)

        cost = max(float(actual_cost_usd or 0.0), 0.0)
        budget.current_usage_usd = float(budget.current_usage_usd or 0.0) + cost

        budget_amount = float(budget.budget_amount_usd or 0.0)
        usage_percent = (float(budget.current_usage_usd) / budget_amount * 100) if budget_amount > 0 else 0.0

        if usage_percent >= float(budget.alert_threshold_percent or 80.0):
            await self.generate_budget_alert(
                budget=budget,
                alert_type="threshold_reached",
                usage_usd=float(budget.current_usage_usd),
                usage_percent=usage_percent,
                message=(
                    f"预算使用率达到 {usage_percent:.2f}% "
                    f"(阈值 {float(budget.alert_threshold_percent or 80.0):.2f}%)"
                ),
                action_taken=None,
            )

        if budget_amount > 0 and float(budget.current_usage_usd) > budget_amount and bool(budget.hard_limit_enabled):
            await self.enforce_spending_limit(
                budget=budget,
                usage_usd=float(budget.current_usage_usd),
                usage_percent=usage_percent,
            )

        await self.db.commit()
        await self.db.refresh(budget)

        return {
            "updated": True,
            "budget_id": budget.id,
            "current_usage_usd": float(budget.current_usage_usd),
            "usage_percent": round(usage_percent, 2),
            "status": budget.status,
        }