use bamboo_agent_core::tools::ToolCallAccumulator;

use super::StreamHandlingOutput;

pub(super) struct StreamAccumulationState {
    response_id: Option<String>,
    content: String,
    reasoning_content: String,
    token_count: usize,
    tool_calls: ToolCallAccumulator,
    output_tokens: u64,
    thinking_tokens: u64,
    cache_creation_input_tokens: u64,
    cache_read_input_tokens: u64,
}

impl StreamAccumulationState {
    pub(super) fn new() -> Self {
        Self {
            response_id: None,
            content: String::new(),
            reasoning_content: String::new(),
            token_count: 0,
            tool_calls: ToolCallAccumulator::new(),
            output_tokens: 0,
            thinking_tokens: 0,
            cache_creation_input_tokens: 0,
            cache_read_input_tokens: 0,
        }
    }

    pub(super) fn append_token(&mut self, token: &str) {
        self.token_count += token.len();
        self.content.push_str(token);
    }

    pub(super) fn append_reasoning_token(&mut self, token: &str) {
        self.reasoning_content.push_str(token);
    }

    pub(super) fn set_response_id(&mut self, response_id: String) {
        self.response_id = Some(response_id);
    }

    pub(super) fn extend_tool_calls(
        &mut self,
        partial_calls: Vec<bamboo_agent_core::tools::ToolCall>,
    ) {
        self.tool_calls.extend(partial_calls);
    }

    pub(super) fn record_usage(&mut self, output_tokens: u64, thinking_tokens: u64) {
        self.output_tokens = output_tokens;
        self.thinking_tokens = thinking_tokens;
    }

    pub(super) fn record_cache(&mut self, creation: u64, read: u64) {
        self.cache_creation_input_tokens =
            self.cache_creation_input_tokens.saturating_add(creation);
        self.cache_read_input_tokens = self.cache_read_input_tokens.saturating_add(read);
    }

    pub(super) fn into_output(self) -> StreamHandlingOutput {
        StreamHandlingOutput {
            response_id: self.response_id,
            content: self.content,
            reasoning_content: self.reasoning_content,
            token_count: self.token_count,
            tool_calls: self.tool_calls.finalize(),
            output_tokens: self.output_tokens,
            thinking_tokens: self.thinking_tokens,
            cache_creation_input_tokens: self.cache_creation_input_tokens,
            cache_read_input_tokens: self.cache_read_input_tokens,
        }
    }
}
