"""Per-task token accumulator using contextvars.

Usage:
    acc = TokenAccumulator()
    set_accumulator(acc)
    # ... all LLM calls within this context auto-record ...
    summary = acc.summary()
    clear_accumulator()
"""

from __future__ import annotations

import threading
from contextvars import ContextVar

_accumulator_var: ContextVar[TokenAccumulator | None] = ContextVar("token_accumulator", default=None)
_stage_var: ContextVar[str] = ContextVar("token_stage", default="execution")

_INPUT_RATE = 3.0e-6
_OUTPUT_RATE = 15.0e-6


def get_accumulator() -> TokenAccumulator | None:
    return _accumulator_var.get()


def set_accumulator(acc: TokenAccumulator) -> None:
    _accumulator_var.set(acc)


def clear_accumulator() -> None:
    _accumulator_var.set(None)


def get_stage() -> str:
    return _stage_var.get()


def set_stage(stage: str) -> None:
    _stage_var.set(stage)


class TokenAccumulator:
    """Thread-safe token counter for a single task execution."""

    def __init__(self) -> None:
        self._calls: list[dict] = []
        self._lock = threading.Lock()

    def record(
        self,
        method: str,
        input_tokens: int,
        output_tokens: int,
        purpose: str = "",
        cache_read_input_tokens: int = 0,
        cache_creation_input_tokens: int = 0,
        stage: str | None = None,
        provider: str | None = None,
        model: str | None = None,
        cost_usd: float | None = None,
    ) -> None:
        with self._lock:
            self._calls.append(
                {
                    "method": method,
                    "input_tokens": input_tokens,
                    "output_tokens": output_tokens,
                    "purpose": purpose,
                    "cache_read_input_tokens": cache_read_input_tokens,
                    "cache_creation_input_tokens": cache_creation_input_tokens,
                    "stage": stage if stage is not None else get_stage(),
                    "provider": provider,
                    "model": model,
                    "cost_usd": cost_usd,
                }
            )

    def total_input(self) -> int:
        with self._lock:
            return sum(c["input_tokens"] for c in self._calls)

    def total_output(self) -> int:
        with self._lock:
            return sum(c["output_tokens"] for c in self._calls)

    def total(self) -> int:
        with self._lock:
            return sum(c["input_tokens"] + c["output_tokens"] for c in self._calls)

    def summary(self) -> dict:
        with self._lock:
            calls = list(self._calls)

        # Build per-stage breakdown
        stages: dict[str, dict] = {}
        for call in calls:
            s = call.get("stage", "execution")
            if s not in stages:
                stages[s] = {"input_tokens": 0, "output_tokens": 0, "calls": 0}
            stages[s]["input_tokens"] += call["input_tokens"]
            stages[s]["output_tokens"] += call["output_tokens"]
            stages[s]["calls"] += 1

        total_input = sum(c["input_tokens"] for c in calls)
        total_output = sum(c["output_tokens"] for c in calls)
        cost_usd = sum(
            c["cost_usd"]
            if c.get("cost_usd") is not None
            else (c["input_tokens"] * _INPUT_RATE) + (c["output_tokens"] * _OUTPUT_RATE)
            for c in calls
        )
        providers = sorted({c["provider"] for c in calls if c.get("provider")})
        models = sorted({c["model"] for c in calls if c.get("model")})

        return {
            "calls": calls,
            "input_tokens": total_input,
            "output_tokens": total_output,
            "total_tokens": total_input + total_output,
            "cache_read_input_tokens": sum(c.get("cache_read_input_tokens", 0) for c in calls),
            "cache_creation_input_tokens": sum(c.get("cache_creation_input_tokens", 0) for c in calls),
            "cost_usd": round(cost_usd, 6),
            "providers": providers,
            "models": models,
            "stages": stages,
        }
