from __future__ import annotations

from app.agents.schemas.task import CollaborationBudget


def build_subtask_budget(
    *,
    execution_mode: str,
    max_parallel_tasks: int,
    max_tool_calls: int = 2,
    max_iterations: int = 2,
    metadata: dict | None = None,
) -> CollaborationBudget:
    return CollaborationBudget(
        mode="collaboration" if execution_mode != "direct" else "direct",
        max_parallel_tasks=max_parallel_tasks,
        remaining_parallel_tasks=max_parallel_tasks,
        max_tool_calls=max_tool_calls,
        remaining_tool_calls=max_tool_calls,
        max_iterations=max_iterations,
        remaining_iterations=max_iterations,
        escalation_threshold=1,
        metadata=metadata or {},
    )