"""预算强制执行服务

实现：
- AI 调用前预算检查
- 超支硬限制执行
- 预算使用量累计
- 80% 阈值预警与超支告警
"""

from __future__ import annotations

from datetime import datetime, timedelta
from typing import Any, Optional

from sqlalchemy import and_, desc, select
from sqlalchemy.ext.asyncio import AsyncSession

from ..models import Budget, BudgetAlert


class BudgetEnforcementService:
    """预算强制执行服务"""

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

    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 = datetime.utcnow()
        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

        # 阈值预警（80% 默认）
        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 = datetime.utcnow()
        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,
        }

    async def enforce_spending_limit(
        self,
        budget: Budget,
        usage_usd: float,
        usage_percent: float,
    ) -> None:
        """执行超支硬性限制：标记预算为 suspended 并记录告警"""
        budget.status = "suspended"

        await self.generate_budget_alert(
            budget=budget,
            alert_type="exceeded",
            usage_usd=usage_usd,
            usage_percent=usage_percent,
            message=f"预算已超支，当前使用率 {usage_percent:.2f}%",
            action_taken="mark_exceeded",
        )
        await self.generate_budget_alert(
            budget=budget,
            alert_type="suspended",
            usage_usd=usage_usd,
            usage_percent=usage_percent,
            message="预算超支，已暂停新的 AI 调用",
            action_taken="suspend_ai_calls",
        )
        await self.db.commit()

    async def generate_budget_alert(
        self,
        budget: Budget,
        alert_type: str,
        usage_usd: float,
        usage_percent: float,
        message: str,
        action_taken: Optional[str],
    ) -> BudgetAlert:
        """生成预算预警（同类 active 预警存在则复用更新）"""
        result = await self.db.execute(
            select(BudgetAlert).where(
                and_(
                    BudgetAlert.budget_id == budget.id,
                    BudgetAlert.alert_type == alert_type,
                    BudgetAlert.status == "active",
                )
            ).order_by(desc(BudgetAlert.created_at))
        )
        existing = result.scalar_one_or_none()
        if existing:
            existing.usage_usd = usage_usd
            existing.usage_percent = usage_percent
            existing.message = message
            existing.action_taken = action_taken
            await self.db.commit()
            await self.db.refresh(existing)
            return existing

        alert = BudgetAlert(
            tenant_id=budget.tenant_id,
            budget_id=budget.id,
            alert_type=alert_type,
            usage_usd=usage_usd,
            usage_percent=usage_percent,
            message=message,
            is_read=False,
            action_taken=action_taken,
            status="active",
        )
        self.db.add(alert)
        await self.db.commit()
        await self.db.refresh(alert)
        return alert

    async def override_budget_limit(
        self,
        budget_id: int,
        tenant_id: int,
        operator_user_id: int,
        note: str = "",
        duration_hours: int = 24,
    ) -> dict[str, Any]:
        """BOSS/ADMIN 超支解锁"""
        result = await self.db.execute(
            select(Budget).where(
                and_(
                    Budget.id == budget_id,
                    Budget.tenant_id == tenant_id,
                )
            )
        )
        budget = result.scalar_one_or_none()
        if not budget:
            raise ValueError(f"预算 {budget_id} 不存在")

        now = datetime.utcnow()
        expires_at = now + timedelta(hours=max(duration_hours, 1))
        history = dict(budget.override_history or {})
        records = list(history.get("records", []))
        records.append(
            {
                "operator_user_id": operator_user_id,
                "note": note,
                "overridden_at": now.isoformat(),
                "expires_at": expires_at.isoformat(),
            }
        )
        history["records"] = records
        budget.override_history = history
        budget.status = "overridden"

        await self.db.commit()
        await self.db.refresh(budget)
        return {
            "budget_id": budget.id,
            "status": budget.status,
            "expires_at": expires_at.isoformat(),
            "override_count": len(records),
        }

    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
            # 周期切换后恢复 active，允许重新消费
            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

        # default monthly
        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