"""The rupee budget for AI usage: one shared pool for the test group, an equal allocation per tester.

Every model call is recorded (api.logstore `usage`). A call is billable when a provider charges for it;
OpenRouter-exempt calls (free models, or anything OpenRouter reports at zero cost) are recorded with the
reason and never deducted. GPU/CPU jobs are not part of this pool: they run on Modal's free credit and have
their own dollar limit (api.jobs, MONTHLY_CREDIT_LIMIT_USD).

Before an AI task runs, its estimate is checked against the tester's remaining allocation and the pool; after
it runs, the actual cost (as OpenRouter reports it, else computed from list prices) is what gets deducted.
"""

from __future__ import annotations

from dataclasses import dataclass

from .routing import ModelSpec, call_cost_usd

CATEGORIES = ("coding", "simulation_ai", "debugging", "testing")


class BudgetError(Exception):
    def __init__(self, status: int, message: str, detail: dict | None = None):
        super().__init__(message)
        self.status = status
        self.detail = detail or {}


@dataclass
class CallCost:
    cost_usd: float
    billable: bool
    exempt_reason: str | None
    estimated: bool


def cost_of_call(model: ModelSpec, usage: dict) -> CallCost:
    """What one model call cost, from the provider's own figure when it reports one."""
    tin = int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0)
    tout = int(usage.get("completion_tokens") or usage.get("output_tokens") or 0)
    reported = usage.get("cost")
    if model.free:
        return CallCost(0.0, False, "opencode_zen_free", False)
    if model.id.endswith(":free"):
        return CallCost(0.0, False, "free model (no OpenRouter charge)", False)
    if reported is not None:
        cost = float(reported)
        if cost <= 0:
            return CallCost(0.0, False, "OpenRouter reported no charge", False)
        return CallCost(cost, True, None, False)
    return CallCost(call_cost_usd(model, tin, tout), True, None, True)


def record_call(logs, settings, *, tester, generation_id, category, model: ModelSpec, usage: dict, estimate_inr=None):
    c = cost_of_call(model, usage)
    if tester == "owner":
        c.billable = False
        c.exempt_reason = "owner key usage"
    logs.record_usage(
        tester=tester,
        generation_id=generation_id,
        category=category,
        model=model.id,
        tier=model.tier,
        provider="opencode-zen"
        if model.free
        else ("openrouter" if "openrouter" in settings.llm_base_url else settings.llm_base_url),
        input_tokens=int(usage.get("prompt_tokens") or usage.get("input_tokens") or 0),
        output_tokens=int(usage.get("completion_tokens") or usage.get("output_tokens") or 0),
        cost_usd=round(c.cost_usd, 6),
        cost_inr=round(c.cost_usd * settings.usd_inr, 4),
        billable=int(c.billable),
        exempt_reason=c.exempt_reason,
        estimated=int(c.estimated),
        estimate_inr=estimate_inr,
    )


def summary(logs, settings, tester: str | None = None) -> dict:
    """Used / remaining for one tester (their allocation) and for the whole pool, split by category."""
    rows = logs.usage_rows()
    mine = [r for r in rows if tester is None or r["tester"] == tester]
    billable = [r for r in mine if r["billable"]]
    by_cat = {c: round(sum(r["cost_inr"] for r in billable if r["category"] == c), 2) for c in CATEGORIES}
    used = round(sum(r["cost_inr"] for r in billable), 2)
    pool_used = round(sum(r["cost_inr"] for r in rows if r["billable"]), 2)
    return {
        "currency": "INR",
        "allocation": settings.budget_per_user_inr,
        "used": used,
        "remaining": round(max(settings.budget_per_user_inr - used, 0.0), 2),
        "by_category": by_cat,
        "exempt_calls": sum(1 for r in mine if not r["billable"]),
        "calls": len(mine),
        "pool": {
            "total": settings.budget_total_inr,
            "used": pool_used,
            "remaining": round(max(settings.budget_total_inr - pool_used, 0.0), 2),
        },
    }


def team(logs, settings, testers: list[str]) -> list[dict]:
    rows = logs.usage_rows()
    out = []
    for name in sorted(set(testers) | {r["tester"] for r in rows}):
        if name == "owner":
            continue
        used = round(sum(r["cost_inr"] for r in rows if r["tester"] == name and r["billable"]), 2)
        out.append(
            {
                "tester": name,
                "allocation": settings.budget_per_user_inr,
                "used": used,
                "remaining": round(max(settings.budget_per_user_inr - used, 0.0), 2),
            }
        )
    return out


def check(logs, settings, tester: str, estimate: tuple[float, float], confirmed: bool) -> None:
    """Refuse when the tester's allocation or the pool cannot cover the likely cost; ask above the threshold."""
    if tester == "owner":
        return
    s = summary(logs, settings, tester)
    likely, worst = estimate
    if s["pool"]["remaining"] < likely:
        raise BudgetError(402, "The team's AI budget is used up. Please tell the Orcanion team.", {"budget": s})
    if s["remaining"] < likely:
        raise BudgetError(402, "Your AI budget is used up. Please tell the Orcanion team.", {"budget": s})
    if worst >= settings.confirm_above_inr and not confirmed:
        raise BudgetError(
            409,
            "This task may cost more than usual; confirm to run it.",
            {"needs_confirmation": True, "estimate_inr": list(estimate), "budget": s},
        )