magi-code 0.96.1

use std::collections::{HashMap, HashSet};

use crate::providers::{ChatMessage, ProviderConversationItem, ProviderToolResult, ToolCall};

use super::tool_continuation::{ToolCallFingerprint, function_call_item};

#[derive(Debug, Default)]
pub(super) struct AgentTurnState {
    current_turn_items: Vec<ProviderConversationItem>,
    current_assistant_segment: String,
    current_assistant_replay_pending: String,
    seen_tool_calls: HashSet<ToolCallFingerprint>,
    seen_tool_call_ids: HashMap<String, ToolCallFingerprint>,
    function_call_item_ids: HashSet<String>,
    saw_function_call_response_item: bool,
    pending_segment_separator: bool,
    iteration: usize,
    completion_verification_attempts: u8,
    pub(super) direct_tool_usage: crate::code_mode::usage::DirectToolUsage,
}

impl AgentTurnState {
    pub(super) fn request_items_slice(&self) -> &[ProviderConversationItem] {
        &self.current_turn_items
    }

    pub(super) fn has_tool_protocol_activity(&self) -> bool {
        !self.seen_tool_calls.is_empty() || self.saw_function_call_response_item
    }

    pub(super) fn iteration(&self) -> usize {
        self.iteration
    }

    pub(super) fn completion_verification_attempts(&self) -> u8 {
        self.completion_verification_attempts
    }

    pub(super) fn increment_completion_verification_attempts(&mut self) {
        self.completion_verification_attempts =
            self.completion_verification_attempts.saturating_add(1);
    }

    pub(super) fn has_pending_segment_separator(&self) -> bool {
        self.pending_segment_separator
    }

    pub(super) fn commit_pending_segment_separator(&mut self) {
        self.pending_segment_separator = false;
    }

    pub(super) fn push_assistant_delta(&mut self, delta: &str) {
        self.current_assistant_segment.push_str(delta);
        self.current_assistant_replay_pending.push_str(delta);
    }

    pub(super) fn assistant_segment(&self) -> &str {
        &self.current_assistant_segment
    }

    pub(super) fn push_response_item(&mut self, item: serde_json::Value) {
        self.flush_pending_assistant_replay();
        self.push_conversation_item(ProviderConversationItem::ResponseItem(item));
    }

    pub(super) fn prepare_tool_turn(&mut self) {
        if !self.current_assistant_segment.trim().is_empty() {
            self.flush_pending_assistant_replay();
        }
    }

    pub(super) fn register_tool_call(&mut self, call: &ToolCall) -> anyhow::Result<()> {
        let fingerprint = ToolCallFingerprint::from(call);
        if !call.id.trim().is_empty() {
            if let Some(existing) = self.seen_tool_call_ids.get(&call.id) {
                if existing != &fingerprint {
                    anyhow::bail!("conflicting duplicate tool call id: {}", call.id);
                }
            } else {
                self.seen_tool_call_ids
                    .insert(call.id.clone(), fingerprint.clone());
            }
        }
        if !self.seen_tool_calls.insert(fingerprint) {
            anyhow::bail!("duplicate tool call suppressed");
        }
        Ok(())
    }

    pub(super) fn append_function_call_if_missing(&mut self, call: &ToolCall) {
        if !self.has_function_call_item(&call.id) {
            self.push_conversation_item(ProviderConversationItem::ResponseItem(
                function_call_item(call),
            ));
        }
    }

    pub(super) fn append_tool_result(&mut self, provider_result: ProviderToolResult) {
        self.push_conversation_item(ProviderConversationItem::ToolResult(provider_result));
    }

    pub(super) fn append_auto_continue(&mut self, prompt: &str) {
        self.flush_pending_assistant_replay();
        self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::user(prompt)));
    }

    pub(super) fn finish_dangling_tool_intent_recovery(&mut self) {
        self.finish_text_action_for_continuation();
    }

    pub(super) fn append_skill_suggestion_hint(&mut self, hint: String) {
        self.flush_pending_assistant_replay();
        self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::user(hint)));
    }

    pub(super) fn append_provider_context_items(
        &mut self,
        items: impl IntoIterator<Item = ProviderConversationItem>,
    ) {
        self.flush_pending_assistant_replay();
        let items = items.into_iter();
        self.current_turn_items.reserve(items.size_hint().0);
        for item in items {
            self.push_conversation_item(item);
        }
    }

    pub(super) fn finish_text_action_for_continuation(&mut self) {
        self.flush_pending_assistant_replay();
        self.current_assistant_segment.clear();
        self.current_assistant_replay_pending.clear();
        self.pending_segment_separator = true;
        self.iteration += 1;
    }

    pub(super) fn finish_tool_iteration(&mut self) {
        self.current_assistant_segment.clear();
        self.current_assistant_replay_pending.clear();
        self.pending_segment_separator = true;
        self.iteration += 1;
    }

    fn flush_pending_assistant_replay(&mut self) {
        if self.current_assistant_replay_pending.trim().is_empty() {
            self.current_assistant_replay_pending.clear();
            return;
        }
        let text = std::mem::take(&mut self.current_assistant_replay_pending);
        self.push_conversation_item(ProviderConversationItem::Message(ChatMessage::assistant(
            text,
        )));
    }

    fn push_conversation_item(&mut self, item: ProviderConversationItem) {
        self.track_function_call_item(&item);
        self.current_turn_items.push(item);
    }

    fn track_function_call_item(&mut self, item: &ProviderConversationItem) {
        let ProviderConversationItem::ResponseItem(value) = item else {
            return;
        };
        if value.get("type").and_then(serde_json::Value::as_str) != Some("function_call") {
            return;
        }
        self.saw_function_call_response_item = true;
        if let Some(call_id) = value.get("call_id").and_then(serde_json::Value::as_str) {
            self.function_call_item_ids.insert(call_id.to_string());
        }
    }

    fn has_function_call_item(&self, call_id: &str) -> bool {
        self.function_call_item_ids.contains(call_id)
    }
}