"""Token estimates and tier limits used by the executable role planner."""
from __future__ import annotations

import os
from typing import Any

import yaml

ROOT = os.environ.get("CLAUDE_PROJECT_DIR") or os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))))
)
KPI_PATH = os.path.join(ROOT, "org-os", "06-agent-work", "agent-operating-kpi.yaml")

TIER_ROLE_LIMITS = {"light": 2, "standard": 5, "heavy": 12}
TIER_ROLE_TOKENS = {
    "light": {"input": 9000, "output": 2500},
    "standard": {"input": 18000, "output": 4500},
    "heavy": {"input": 32000, "output": 8000},
}
ROLE_TYPE_MULTIPLIER = {
    "coordinator": 0.65,
    "worker": 1.0,
    "recommender": 0.9,
    "reviewer": 0.85,
    "auditor": 0.9,
    "decider": 1.05,
}


def workflow_budget(tier: str, override: int | None = None) -> int:
    if override is not None:
        return max(0, int(override))
    try:
        data = yaml.safe_load(open(KPI_PATH, encoding="utf-8")) or {}
        budgets = data["agent-operating-kpi"]["token-budgets"]["per-wave"]
        return int(budgets[tier])
    except Exception:
        return {"light": 150000, "standard": 500000, "heavy": 2000000}.get(tier, 500000)


def estimate_role(role: dict[str, Any], tier: str, stage: str | None = None) -> dict[str, int]:
    base = TIER_ROLE_TOKENS.get(tier, TIER_ROLE_TOKENS["standard"])
    multiplier = ROLE_TYPE_MULTIPLIER.get(str(role.get("role-type") or "worker"), 1.0)
    if stage in {"verification", "acceptance"}:
        multiplier *= 0.8
    input_tokens = int(base["input"] * multiplier)
    output_tokens = int(base["output"] * multiplier)
    return {"input": input_tokens, "output": output_tokens, "total": input_tokens + output_tokens}


def estimate_plan(selected_roles: list[dict[str, Any]], tier: str, stage: str | None = None) -> dict[str, int]:
    input_tokens = output_tokens = 0
    for role in selected_roles:
        estimate = estimate_role(role, tier, stage)
        input_tokens += estimate["input"]
        output_tokens += estimate["output"]
    synthesis = 0 if len(selected_roles) <= 1 else (3500 if tier == "light" else 7000) * (len(selected_roles) - 1)
    return {
        "input": input_tokens,
        "output": output_tokens,
        "synthesis": synthesis,
        "total": input_tokens + output_tokens + synthesis,
    }


def max_selected_roles(tier: str) -> int:
    return TIER_ROLE_LIMITS.get(tier, TIER_ROLE_LIMITS["standard"])