from __future__ import annotations

from pydantic import BaseModel, Field


class BudgetPlan(BaseModel):
    """Computed token budget plan for a session."""

    session_budget: int
    agent_allocations: dict[str, int]
    estimated_total: int
    allocation_strategy: str
    warnings: list[str] = Field(default_factory=list)


class BudgetAllocator:
    """Pre-flight budget estimation and proportional allocation."""

    def __init__(self, session_budget: int = 500_000) -> None:
        self.session_budget = session_budget

    def allocate(self, tasks: list[dict]) -> BudgetPlan:
        """Allocate budget by weighted task complexity and expected calls."""
        if not tasks:
            return BudgetPlan(
                session_budget=self.session_budget,
                agent_allocations={},
                estimated_total=0,
                allocation_strategy="proportional_weight",
                warnings=["no tasks provided"],
            )

        per_agent_weight: dict[str, float] = {}
        per_agent_est: dict[str, int] = {}
        estimated_total = 0

        for task in tasks:
            agent_id = str(task.get("agent_id", "unknown"))
            description = str(task.get("description", ""))
            tool_calls = int(task.get("estimated_tool_calls", 1) or 1)
            complexity = float(task.get("complexity", 1.0) or 1.0)

            estimated = self._estimate_task_tokens(description, tool_calls, complexity)
            weight = max(len(description), 1) * max(tool_calls, 1) * max(complexity, 0.1)

            per_agent_est[agent_id] = per_agent_est.get(agent_id, 0) + estimated
            per_agent_weight[agent_id] = per_agent_weight.get(agent_id, 0.0) + weight
            estimated_total += estimated

        total_weight = sum(per_agent_weight.values()) or 1.0
        allocations: dict[str, int] = {}
        for agent_id, weight in per_agent_weight.items():
            allocations[agent_id] = max(100, int(self.session_budget * (weight / total_weight)))

        warnings: list[str] = []
        pct = estimated_total / self.session_budget if self.session_budget else 0.0
        if pct > 0.9:
            warnings.append(f"estimated usage is {pct:.0%} of budget")
        if estimated_total > self.session_budget:
            over = ((estimated_total - self.session_budget) / self.session_budget) * 100
            warnings.append(f"estimated > budget by {over:.1f}%")

        return BudgetPlan(
            session_budget=self.session_budget,
            agent_allocations=allocations,
            estimated_total=estimated_total,
            allocation_strategy="proportional_weight",
            warnings=warnings,
        )

    def _estimate_task_tokens(self, description: str, tool_calls: int, complexity: float) -> int:
        """Heuristic token estimate for a task."""
        turns_estimate = 1 + max(tool_calls, 1)
        prompt_tokens = int(len(description.split()) * 1.3 * turns_estimate)
        return max(32, int(prompt_tokens * max(complexity, 0.1)))
