from __future__ import annotations

from dataclasses import dataclass
from typing import Generic, TypeVar

T = TypeVar("T")


@dataclass
class Usage:
    input_tokens: int = 0
    output_tokens: int = 0

    def __add__(self, other: Usage) -> Usage:
        return Usage(
            self.input_tokens + other.input_tokens, self.output_tokens + other.output_tokens
        )


@dataclass
class AgentUsage:
    generator: Usage
    formatter: Usage
    validator: Usage

    @property
    def total(self) -> Usage:
        return self.generator + self.formatter + self.validator


@dataclass
class AgentResult(Generic[T]):
    output: T  # validated Pydantic model instance
    raw_response: str  # raw text from generator (before formatting)
    validation_feedback: str  # validator's feedback on the passing run
    attempts: int  # pipeline runs used (1 = succeeded first try)
    usage: AgentUsage  # token usage broken down by pipeline step
