from __future__ import annotations

from typing import TYPE_CHECKING

from gamma23.types import AgentBudget

if TYPE_CHECKING:
    from gamma23.client import Gamma23


class Agents:
    """Query agent-level information."""

    def __init__(self, client: Gamma23) -> None:
        self._client = client

    def get_budget(self, agent_id: str) -> AgentBudget:
        """Retrieve the budget for an agent.

        Args:
            agent_id: The agent identifier.

        Returns:
            An :class:`AgentBudget` with current spend and remaining balance.
        """
        data = self._client._request("GET", f"/api/v1/agents/{agent_id}/budget")
        return AgentBudget.from_dict(data)
