"""Runtime task/team budget accounting for economic autonomy."""

from __future__ import annotations

from dataclasses import asdict, dataclass
from typing import Any, MutableMapping, TypedDict

from logos_adk.types import Usage


class BudgetAccount(TypedDict):
    """JSON-safe mutable runtime budget account."""

    scope: str
    name: str | None
    limit_usd: float | None
    limit_tokens: int | None
    spent_usd: float
    spent_tokens: int
    threshold_ratio: float


@dataclass
class BudgetStatus:
    """Computed budget status for a runtime account."""

    scope: str
    name: str | None
    spent_usd: float
    remaining_usd: float | None
    limit_usd: float | None
    spent_tokens: int
    remaining_tokens: int | None
    limit_tokens: int | None
    ratio: float
    near_limit: bool
    exhausted: bool
    threshold_ratio: float

    def to_dict(self) -> dict[str, Any]:
        return asdict(self)


def create_budget_account(
    *,
    scope: str,
    name: str | None = None,
    limit_usd: float | None = None,
    limit_tokens: int | None = None,
    threshold_ratio: float = 0.8,
    spent_usd: float = 0.0,
    spent_tokens: int = 0,
) -> BudgetAccount | None:
    """Create a JSON-safe runtime budget account, or ``None`` if no limits are set."""
    normalized_limit_usd = None if limit_usd is None else max(0.0, float(limit_usd))
    normalized_limit_tokens = None if limit_tokens is None else max(0, int(limit_tokens))
    if normalized_limit_usd is None and normalized_limit_tokens is None:
        return None
    return BudgetAccount(
        scope=str(scope or "task"),
        name=name,
        limit_usd=normalized_limit_usd,
        limit_tokens=normalized_limit_tokens,
        spent_usd=max(0.0, float(spent_usd)),
        spent_tokens=max(0, int(spent_tokens)),
        threshold_ratio=max(0.0, float(threshold_ratio)),
    )


def ensure_budget_account(
    value: Any,
    *,
    scope: str,
    name: str | None = None,
    limit_usd: float | None = None,
    limit_tokens: int | None = None,
    threshold_ratio: float = 0.8,
) -> BudgetAccount | None:
    """Normalize an incoming runtime budget payload to a JSON-safe budget account."""
    if value is None:
        return create_budget_account(
            scope=scope,
            name=name,
            limit_usd=limit_usd,
            limit_tokens=limit_tokens,
            threshold_ratio=threshold_ratio,
        )
    if not isinstance(value, MutableMapping):
        raise TypeError("Budget account must be a mutable mapping")
    account = create_budget_account(
        scope=str(value.get("scope", scope)),
        name=value.get("name", name),
        limit_usd=value.get("limit_usd", limit_usd),
        limit_tokens=value.get("limit_tokens", limit_tokens),
        threshold_ratio=value.get("threshold_ratio", threshold_ratio),
        spent_usd=value.get("spent_usd", 0.0),
        spent_tokens=value.get("spent_tokens", 0),
    )
    if account is None:
        return None
    value.clear()
    value.update(account)
    return value


def record_budget_usage(
    account: MutableMapping[str, Any] | None,
    *,
    usage: Usage,
    amount_usd: float,
) -> BudgetStatus | None:
    """Record runtime usage into the mutable budget account."""
    if account is None:
        return None
    account["spent_usd"] = round(max(0.0, float(account.get("spent_usd", 0.0))) + max(0.0, float(amount_usd)), 8)
    account["spent_tokens"] = max(0, int(account.get("spent_tokens", 0))) + max(
        0,
        int(usage.input_tokens) + int(usage.output_tokens),
    )
    return budget_status(account)


def budget_status(account: MutableMapping[str, Any] | None) -> BudgetStatus | None:
    """Compute the current status of a runtime budget account."""
    if account is None:
        return None
    spent_usd = max(0.0, float(account.get("spent_usd", 0.0)))
    spent_tokens = max(0, int(account.get("spent_tokens", 0)))
    limit_usd = account.get("limit_usd")
    limit_tokens = account.get("limit_tokens")
    threshold_ratio = max(0.0, float(account.get("threshold_ratio", 0.8)))

    ratio_components: list[float] = []
    exhausted = False
    remaining_usd: float | None = None
    remaining_tokens: int | None = None

    if limit_usd is not None:
        limit_usd = max(0.0, float(limit_usd))
        remaining_usd = max(0.0, round(limit_usd - spent_usd, 8))
        ratio_components.append(1.0 if limit_usd == 0.0 else spent_usd / limit_usd)
        exhausted = exhausted or spent_usd >= limit_usd
    if limit_tokens is not None:
        limit_tokens = max(0, int(limit_tokens))
        remaining_tokens = max(0, limit_tokens - spent_tokens)
        ratio_components.append(1.0 if limit_tokens == 0 else spent_tokens / limit_tokens)
        exhausted = exhausted or spent_tokens >= limit_tokens

    ratio = max(ratio_components, default=0.0)
    return BudgetStatus(
        scope=str(account.get("scope", "task")),
        name=account.get("name"),
        spent_usd=round(spent_usd, 8),
        remaining_usd=remaining_usd,
        limit_usd=limit_usd,
        spent_tokens=spent_tokens,
        remaining_tokens=remaining_tokens,
        limit_tokens=limit_tokens,
        ratio=round(ratio, 6),
        near_limit=ratio >= threshold_ratio,
        exhausted=exhausted,
        threshold_ratio=threshold_ratio,
    )


__all__ = [
    "BudgetAccount",
    "BudgetStatus",
    "budget_status",
    "create_budget_account",
    "ensure_budget_account",
    "record_budget_usage",
]
