"""
Budget Controls - Cost management for agents and tools.

This module implements budget enforcement:
- Per-agent daily limits (tokens, cost)
- Per-MCP-server limits
- Per-tool limits
- MCP registry allowlist/blocklist

Usage:
    from agent_orchestrator.budget import (
        BudgetEnforcer,
        AgentBudget,
        MCPUsageTracker,
        MCPRegistry,
    )

    # Set up agent budget enforcement
    enforcer = BudgetEnforcer(db)
    result = enforcer.check_budget("claude-code", estimated_input_tokens=10000)
    if result.allowed:
        # Execute request
        enforcer.record_usage("claude-code", input_tokens=9500, output_tokens=2000)

    # Set up MCP tracking
    tracker = MCPUsageTracker(db)
    result = tracker.check_budget("filesystem", "write_file")
    if result.allowed:
        # Execute MCP call
        tracker.record_call("filesystem", "write_file")

    # Set up MCP registry
    registry = MCPRegistry()
    result = registry.check_access("github", "create_pull_request")
    if result.decision == RegistryDecision.REQUIRE_APPROVAL:
        # Request approval
        pass
"""

from .agent_budget import (
    AgentBudget,
    DailyUsage,
    BudgetStatus,
    BudgetCheckResult,
    BudgetEnforcer,
    DEFAULT_BUDGETS,
    TOKEN_COSTS,
    get_budget_enforcer,
    set_budget_enforcer,
    has_budget,
)

from .mcp_budget import (
    MCPServerBudget,
    ToolBudget,
    MCPUsageRecord,
    MCPBudgetStatus,
    MCPBudgetCheckResult,
    MCPUsageTracker,
    DEFAULT_MCP_BUDGETS,
    HIGH_RISK_TOOLS,
    get_mcp_tracker,
    set_mcp_tracker,
    check_mcp_budget,
)

from .registry import (
    MCPRegistry,
    AccessLevel,
    RegistryDecision,
    RegistryCheckResult,
    ServerRegistration,
    ToolRegistration,
    DEFAULT_ALLOWED_SERVERS,
    DEFAULT_BLOCKED_SERVERS,
    DEFAULT_APPROVAL_REQUIRED_TOOLS,
    get_registry,
    set_registry,
    is_server_allowed,
    is_tool_allowed,
)

from .alerts import (
    AlertSeverity,
    BudgetType,
    BudgetConfig,
    BudgetAlert,
    BudgetExhaustedError,
    BudgetAlertSystem,
    get_budget_alert_system,
    set_budget_alert_system,
)

from .optimizer import (
    InefficiencyType,
    UsageRecord,
    Inefficiency,
    CostAnalysis,
    CostProjection,
    CostOptimizer,
    get_cost_optimizer,
    set_cost_optimizer,
)

__all__ = [
    # Agent Budget
    "AgentBudget",
    "DailyUsage",
    "BudgetStatus",
    "BudgetCheckResult",
    "BudgetEnforcer",
    "DEFAULT_BUDGETS",
    "TOKEN_COSTS",
    "get_budget_enforcer",
    "set_budget_enforcer",
    "has_budget",
    # MCP Budget
    "MCPServerBudget",
    "ToolBudget",
    "MCPUsageRecord",
    "MCPBudgetStatus",
    "MCPBudgetCheckResult",
    "MCPUsageTracker",
    "DEFAULT_MCP_BUDGETS",
    "HIGH_RISK_TOOLS",
    "get_mcp_tracker",
    "set_mcp_tracker",
    "check_mcp_budget",
    # MCP Registry
    "MCPRegistry",
    "AccessLevel",
    "RegistryDecision",
    "RegistryCheckResult",
    "ServerRegistration",
    "ToolRegistration",
    "DEFAULT_ALLOWED_SERVERS",
    "DEFAULT_BLOCKED_SERVERS",
    "DEFAULT_APPROVAL_REQUIRED_TOOLS",
    "get_registry",
    "set_registry",
    "is_server_allowed",
    "is_tool_allowed",
    # Budget Alerts
    "AlertSeverity",
    "BudgetType",
    "BudgetConfig",
    "BudgetAlert",
    "BudgetExhaustedError",
    "BudgetAlertSystem",
    "get_budget_alert_system",
    "set_budget_alert_system",
    # Cost Optimizer
    "InefficiencyType",
    "UsageRecord",
    "Inefficiency",
    "CostAnalysis",
    "CostProjection",
    "CostOptimizer",
    "get_cost_optimizer",
    "set_cost_optimizer",
]
