"""BudgetMixin — team and task budget management."""
from __future__ import annotations

from typing import Any

from logos_adk.budgeting import create_budget_account


class BudgetMixin:
    """Team and task budget configuration and application."""

    def budget(
        self,
        *,
        team_budget_usd: float | None = None,
        team_budget_tokens: int | None = None,
        task_budget_usd: float | None = None,
        task_budget_tokens: int | None = None,
        threshold_ratio: float = 0.8,
        fail_on_budget_exhausted: bool | None = None,
    ) -> Any:
        """Attach default task/team budgets to the team runtime."""
        self._team_budget_usd = None if team_budget_usd is None else max(0.0, float(team_budget_usd))
        self._team_budget_tokens = None if team_budget_tokens is None else max(0, int(team_budget_tokens))
        self._task_budget_usd = None if task_budget_usd is None else max(0.0, float(task_budget_usd))
        self._task_budget_tokens = None if task_budget_tokens is None else max(0, int(task_budget_tokens))
        self._budget_threshold_ratio = max(0.0, float(threshold_ratio))
        if fail_on_budget_exhausted is not None:
            self._fail_on_budget_exhausted = bool(fail_on_budget_exhausted)
        return self

    def _apply_budget_kwargs(self, kwargs: dict[str, Any]) -> dict[str, Any]:
        effective = dict(kwargs)
        threshold_ratio = max(
            0.0,
            float(effective.get("budget_threshold_ratio", self._budget_threshold_ratio)),
        )
        if "_team_budget" not in effective:
            account = create_budget_account(
                scope="team",
                name=self.name,
                limit_usd=effective.get("team_budget_usd", self._team_budget_usd),
                limit_tokens=effective.get("team_budget_tokens", self._team_budget_tokens),
                threshold_ratio=threshold_ratio,
            )
            if account is not None:
                effective["_team_budget"] = account
        if "task_budget_usd" not in effective and self._task_budget_usd is not None:
            effective["task_budget_usd"] = self._task_budget_usd
        if "task_budget_tokens" not in effective and self._task_budget_tokens is not None:
            effective["task_budget_tokens"] = self._task_budget_tokens
        if "budget_threshold_ratio" not in effective:
            effective["budget_threshold_ratio"] = threshold_ratio
        if "fail_on_budget_exhausted" not in effective and self._fail_on_budget_exhausted:
            effective["fail_on_budget_exhausted"] = True
        return effective
