from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import AsyncIterator


@dataclass
class CodeActModelUsage:
    """Tracks token usage and costs from interactions with code action models."""

    total_tokens: int = 0
    input_tokens: int = 0
    thinking_tokens: int = 0
    output_tokens: int = 0
    cache_write_tokens: int = 0
    cache_read_tokens: int = 0

    cost: float | None = None
    """Cost of code action model usage in `USD` based on token counts
    or `None` if cost estimation is not available for the used model.
    """

    def update(self, other: "CodeActModelUsage"):
        """Adds token counts and cost of `other` to this instance.
        This is used to accumulate usage across multiple interactions.

        Args:
            other: The usage instance to add to this instance.
        """

        self.total_tokens += other.total_tokens
        self.input_tokens += other.input_tokens
        self.thinking_tokens += other.thinking_tokens
        self.output_tokens += other.output_tokens
        self.cache_write_tokens += other.cache_write_tokens
        self.cache_read_tokens += other.cache_read_tokens

        if self.cost is None and other.cost is not None:
            self.cost = other.cost
        elif self.cost is not None and other.cost is not None:
            self.cost += other.cost


@dataclass
class CodeActModelResponse(ABC):
    """A response from a code action model. If the `code` property is `None` it
    is a final response to the user, otherwise it is a code action."""

    text: str
    """Response text generated by a code action model. Depending on the strategy
    to generate code actions, this may or may not include the generated code. If
    it contains code, it is extracted and available in the `code` property.
    """

    is_error: bool
    """Whether the response `text` contains error information. If `True`, `text`
    contains error information that is NOT related to code execution errors. Not
    handled by applications but rather `freeact`-internally.
    """

    usage: CodeActModelUsage = field(default_factory=CodeActModelUsage)
    """Token usage and costs from the interaction with a code action model."""

    @property
    @abstractmethod
    def tool_use_id(self) -> str | None: ...

    @property
    @abstractmethod
    def tool_use_name(self) -> str | None: ...

    @property
    @abstractmethod
    def code(self) -> str | None:
        """Executable code generated by a code action model. If `None`, this
        response is a final response to the user.
        """


class CodeActModelTurn(ABC):
    """A single interaction with a code action model. This is either initiated
    by a user query or code execution feedback (code action results or execution
    errors).
    """

    @abstractmethod
    async def response(self) -> CodeActModelResponse:
        """Retrieve the complete response from a code action model. Waits until
        the response is available.
        """

    @abstractmethod
    def stream(self) -> AsyncIterator[str]:
        """Stream the code action model's response as it is generated. Once the
        stream is consumed, [`response`][freeact.model.base.CodeActModelTurn.response]
        is immediately available without waiting.
        """


class CodeActModel(ABC):
    """A code action model.

    A code action model is a model that responds with code if wants to perform an
    an action. An action is performed by executing the generated code.

    A code action model responds to user queries and code execution feedback by
    returning a [`CodeActModelTurn`][freeact.model.base.CodeActModelTurn] object
    which is used to retrieve the model response.
    """

    @abstractmethod
    def request(self, user_query: str, **kwargs) -> CodeActModelTurn:
        """Initiates an interaction with this model from a user query.

        Args:
            user_query: The user query (a question, instruction, etc.)
            **kwargs: Additional interaction-specific parameters

        Returns:
            CodeActModelTurn: An object for retrieving the model's response.
        """

    @abstractmethod
    def feedback(
        self,
        feedback: str,
        is_error: bool,
        tool_use_id: str | None,
        tool_use_name: str | None,
        **kwargs,
    ) -> CodeActModelTurn:
        """Initiates an interaction with this model from code execution feedback,
        allowing the model to refine or correct previous responses, or returning
        a final response to the user. A `feedback` call must follow a previous
        `request` or `feedback` call.

        Args:
            feedback: The feedback text from code execution.
            is_error: Whether the `feedback` text contains error information.
            **kwargs: Additional model-specific parameters for the feedback.

        Returns:
            CodeActModelTurn: An object for retrieving the model's response.
        """
