from __future__ import annotations

from collections.abc import Mapping
from typing import Any, TypeVar, BinaryIO, TextIO, TYPE_CHECKING, Generator

from attrs import define as _attrs_define
from attrs import field as _attrs_field

from ..types import UNSET, Unset

from ..types import UNSET, Unset
from typing import cast

if TYPE_CHECKING:
  from ..models.all_llm_calls_aggregate import AllLlmCallsAggregate
  from ..models.token_budget_context_system_detail_type_0 import TokenBudgetContextSystemDetailType0





T = TypeVar("T", bound="TokenBudgetContext")



@_attrs_define
class TokenBudgetContext:
    """ Token budget snapshot from TokenLedger.snapshot().

    Always-present fields track the token breakdown across context categories.
    Optional fields appear when cost/throughput data is available.

        Attributes:
            context_window (int): Total context window size for the model
            budget (int): Usable budget (context_window minus reserved output)
            system (int): Tokens used by system prompt
            history (int): Tokens used by conversation history
            compaction (int): Tokens used by compacted history
            retrieval (int): Tokens used by retrieved document chunks
            query (int): Tokens used by the current user query
            output (int): Output tokens generated
            total_input (int): Sum of system + history + compaction + retrieval + query
            truncated (bool): Whether context was truncated to fit budget
            reasoning (int): Reasoning tokens used by the model
            total_output (int): Sum of output + reasoning tokens
            system_detail (None | TokenBudgetContextSystemDetailType0 | Unset): Breakdown of system prompt token counts by
                section
            input_cost (float | None | Unset): Cost for input tokens (USD)
            output_cost (float | None | Unset): Cost for output tokens (USD)
            total_cost (float | None | Unset): Total cost: input_cost + output_cost (USD)
            elapsed_seconds (float | None | Unset): Wall-clock time for the LLM call(s)
            throughput_tokens (int | None | Unset): Total LLM output tokens (across all calls, for throughput)
            tokens_per_second (float | None | Unset): Output throughput: throughput_tokens / elapsed
            model_name (None | str | Unset): Model used for the LLM call
            all_llm_calls (AllLlmCallsAggregate | None | Unset): Aggregate usage across all LLM calls in the request
     """

    context_window: int
    budget: int
    system: int
    history: int
    compaction: int
    retrieval: int
    query: int
    output: int
    total_input: int
    truncated: bool
    reasoning: int
    total_output: int
    system_detail: None | TokenBudgetContextSystemDetailType0 | Unset = UNSET
    input_cost: float | None | Unset = UNSET
    output_cost: float | None | Unset = UNSET
    total_cost: float | None | Unset = UNSET
    elapsed_seconds: float | None | Unset = UNSET
    throughput_tokens: int | None | Unset = UNSET
    tokens_per_second: float | None | Unset = UNSET
    model_name: None | str | Unset = UNSET
    all_llm_calls: AllLlmCallsAggregate | None | Unset = UNSET
    additional_properties: dict[str, Any] = _attrs_field(init=False, factory=dict)





    def to_dict(self) -> dict[str, Any]:
        from ..models.token_budget_context_system_detail_type_0 import TokenBudgetContextSystemDetailType0
        from ..models.all_llm_calls_aggregate import AllLlmCallsAggregate
        context_window = self.context_window

        budget = self.budget

        system = self.system

        history = self.history

        compaction = self.compaction

        retrieval = self.retrieval

        query = self.query

        output = self.output

        total_input = self.total_input

        truncated = self.truncated

        reasoning = self.reasoning

        total_output = self.total_output

        system_detail: dict[str, Any] | None | Unset
        if isinstance(self.system_detail, Unset):
            system_detail = UNSET
        elif isinstance(self.system_detail, TokenBudgetContextSystemDetailType0):
            system_detail = self.system_detail.to_dict()
        else:
            system_detail = self.system_detail

        input_cost: float | None | Unset
        if isinstance(self.input_cost, Unset):
            input_cost = UNSET
        else:
            input_cost = self.input_cost

        output_cost: float | None | Unset
        if isinstance(self.output_cost, Unset):
            output_cost = UNSET
        else:
            output_cost = self.output_cost

        total_cost: float | None | Unset
        if isinstance(self.total_cost, Unset):
            total_cost = UNSET
        else:
            total_cost = self.total_cost

        elapsed_seconds: float | None | Unset
        if isinstance(self.elapsed_seconds, Unset):
            elapsed_seconds = UNSET
        else:
            elapsed_seconds = self.elapsed_seconds

        throughput_tokens: int | None | Unset
        if isinstance(self.throughput_tokens, Unset):
            throughput_tokens = UNSET
        else:
            throughput_tokens = self.throughput_tokens

        tokens_per_second: float | None | Unset
        if isinstance(self.tokens_per_second, Unset):
            tokens_per_second = UNSET
        else:
            tokens_per_second = self.tokens_per_second

        model_name: None | str | Unset
        if isinstance(self.model_name, Unset):
            model_name = UNSET
        else:
            model_name = self.model_name

        all_llm_calls: dict[str, Any] | None | Unset
        if isinstance(self.all_llm_calls, Unset):
            all_llm_calls = UNSET
        elif isinstance(self.all_llm_calls, AllLlmCallsAggregate):
            all_llm_calls = self.all_llm_calls.to_dict()
        else:
            all_llm_calls = self.all_llm_calls


        field_dict: dict[str, Any] = {}
        field_dict.update(self.additional_properties)
        field_dict.update({
            "context_window": context_window,
            "budget": budget,
            "system": system,
            "history": history,
            "compaction": compaction,
            "retrieval": retrieval,
            "query": query,
            "output": output,
            "total_input": total_input,
            "truncated": truncated,
            "reasoning": reasoning,
            "total_output": total_output,
        })
        if system_detail is not UNSET:
            field_dict["system_detail"] = system_detail
        if input_cost is not UNSET:
            field_dict["input_cost"] = input_cost
        if output_cost is not UNSET:
            field_dict["output_cost"] = output_cost
        if total_cost is not UNSET:
            field_dict["total_cost"] = total_cost
        if elapsed_seconds is not UNSET:
            field_dict["elapsed_seconds"] = elapsed_seconds
        if throughput_tokens is not UNSET:
            field_dict["throughput_tokens"] = throughput_tokens
        if tokens_per_second is not UNSET:
            field_dict["tokens_per_second"] = tokens_per_second
        if model_name is not UNSET:
            field_dict["model_name"] = model_name
        if all_llm_calls is not UNSET:
            field_dict["all_llm_calls"] = all_llm_calls

        return field_dict



    @classmethod
    def from_dict(cls: type[T], src_dict: Mapping[str, Any]) -> T:
        from ..models.all_llm_calls_aggregate import AllLlmCallsAggregate
        from ..models.token_budget_context_system_detail_type_0 import TokenBudgetContextSystemDetailType0
        d = dict(src_dict)
        context_window = d.pop("context_window")

        budget = d.pop("budget")

        system = d.pop("system")

        history = d.pop("history")

        compaction = d.pop("compaction")

        retrieval = d.pop("retrieval")

        query = d.pop("query")

        output = d.pop("output")

        total_input = d.pop("total_input")

        truncated = d.pop("truncated")

        reasoning = d.pop("reasoning")

        total_output = d.pop("total_output")

        def _parse_system_detail(data: object) -> None | TokenBudgetContextSystemDetailType0 | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            try:
                if not isinstance(data, dict):
                    raise TypeError()
                system_detail_type_0 = TokenBudgetContextSystemDetailType0.from_dict(data)



                return system_detail_type_0
            except (TypeError, ValueError, AttributeError, KeyError):
                pass
            return cast(None | TokenBudgetContextSystemDetailType0 | Unset, data)

        system_detail = _parse_system_detail(d.pop("system_detail", UNSET))


        def _parse_input_cost(data: object) -> float | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(float | None | Unset, data)

        input_cost = _parse_input_cost(d.pop("input_cost", UNSET))


        def _parse_output_cost(data: object) -> float | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(float | None | Unset, data)

        output_cost = _parse_output_cost(d.pop("output_cost", UNSET))


        def _parse_total_cost(data: object) -> float | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(float | None | Unset, data)

        total_cost = _parse_total_cost(d.pop("total_cost", UNSET))


        def _parse_elapsed_seconds(data: object) -> float | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(float | None | Unset, data)

        elapsed_seconds = _parse_elapsed_seconds(d.pop("elapsed_seconds", UNSET))


        def _parse_throughput_tokens(data: object) -> int | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(int | None | Unset, data)

        throughput_tokens = _parse_throughput_tokens(d.pop("throughput_tokens", UNSET))


        def _parse_tokens_per_second(data: object) -> float | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(float | None | Unset, data)

        tokens_per_second = _parse_tokens_per_second(d.pop("tokens_per_second", UNSET))


        def _parse_model_name(data: object) -> None | str | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            return cast(None | str | Unset, data)

        model_name = _parse_model_name(d.pop("model_name", UNSET))


        def _parse_all_llm_calls(data: object) -> AllLlmCallsAggregate | None | Unset:
            if data is None:
                return data
            if isinstance(data, Unset):
                return data
            try:
                if not isinstance(data, dict):
                    raise TypeError()
                all_llm_calls_type_0 = AllLlmCallsAggregate.from_dict(data)



                return all_llm_calls_type_0
            except (TypeError, ValueError, AttributeError, KeyError):
                pass
            return cast(AllLlmCallsAggregate | None | Unset, data)

        all_llm_calls = _parse_all_llm_calls(d.pop("all_llm_calls", UNSET))


        token_budget_context = cls(
            context_window=context_window,
            budget=budget,
            system=system,
            history=history,
            compaction=compaction,
            retrieval=retrieval,
            query=query,
            output=output,
            total_input=total_input,
            truncated=truncated,
            reasoning=reasoning,
            total_output=total_output,
            system_detail=system_detail,
            input_cost=input_cost,
            output_cost=output_cost,
            total_cost=total_cost,
            elapsed_seconds=elapsed_seconds,
            throughput_tokens=throughput_tokens,
            tokens_per_second=tokens_per_second,
            model_name=model_name,
            all_llm_calls=all_llm_calls,
        )


        token_budget_context.additional_properties = d
        return token_budget_context

    @property
    def additional_keys(self) -> list[str]:
        return list(self.additional_properties.keys())

    def __getitem__(self, key: str) -> Any:
        return self.additional_properties[key]

    def __setitem__(self, key: str, value: Any) -> None:
        self.additional_properties[key] = value

    def __delitem__(self, key: str) -> None:
        del self.additional_properties[key]

    def __contains__(self, key: str) -> bool:
        return key in self.additional_properties
