"""
Streaming utilities for accumulating content during streaming responses.
"""
from __future__ import annotations

import time
from typing import Any


class StreamingAccumulator:
    """
    Accumulates content and metadata during streaming responses.

    Used by provider wrappers to collect streaming chunks and emit
    final tracking events when the stream completes.
    """

    def __init__(self) -> None:
        self.content_chunks: list[str] = []
        self.tool_calls: list[dict[str, Any]] = []
        self.token_usage: dict[str, int] = {}
        self.finish_reason: str | None = None
        self.start_time = time.time()
        self.first_byte_time: float | None = None
        self.is_error = False
        self.error_message: str | None = None

    def add_content(self, content: str) -> None:
        """Add a content chunk from the stream."""
        if content and not self.first_byte_time:
            self.first_byte_time = time.time()

        if content:
            self.content_chunks.append(content)

    def add_tool_call(self, tool_call: dict[str, Any]) -> None:
        """Add a tool call from the stream."""
        self.tool_calls.append(tool_call)

    def update_tokens(self, token_update: dict[str, int]) -> None:
        """Update token usage with new data."""
        for key, value in token_update.items():
            if key in self.token_usage:
                # For incremental updates, add to existing
                self.token_usage[key] += value
            else:
                # For cumulative updates, replace
                self.token_usage[key] = value

    def set_finish_reason(self, reason: str) -> None:
        """Set the finish reason."""
        self.finish_reason = reason

    def set_error(self, error: str) -> None:
        """Mark the stream as having an error."""
        self.is_error = True
        self.error_message = error

    def get_full_content(self) -> str:
        """Get the accumulated content as a single string."""
        return "".join(self.content_chunks)

    def get_latency_ms(self) -> float:
        """Get total latency in milliseconds."""
        return (time.time() - self.start_time) * 1000

    def get_ttfb_ms(self) -> float | None:
        """Get time to first byte in milliseconds."""
        if self.first_byte_time:
            return (self.first_byte_time - self.start_time) * 1000
        return None
