zesdex/apps/application/src/agent/turn_service.rs

use std::collections::VecDeque;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};
use tracing::{debug, info, warn};

use zesdex_domain::agent::{AgentTurnParams, TurnEvent};
use zesdex_domain::core::{ChatMessage, StreamEvent, ToolDef};
use zesdex_domain::main_agent_prompt;

use super::ToolExecutor;
use crate::ports::ProviderService;

/// Maximum tool-call iterations per agent turn before forcing termination.
const MAX_TURN_ITERATIONS: u32 = 50;

/// Maximum number of consecutive identical tool errors before the loop
/// injects a recovery note and forces a different approach.
const MAX_CONSECUTIVE_TOOL_ERRORS: usize = 3;

/// Total tool-call errors tolerated per turn before the loop is stopped.
const MAX_TOTAL_TOOL_ERRORS: usize = 8;

/// Ceiling for a single tool-result message inserted into context.
///
/// Tool outputs can be huge (read / semantic_search). Truncating keeps the
/// context window from exploding while preserving the important head.
const TOOL_OUTPUT_MAX_CHARS: usize = 12_000;

/// Total conversation characters that trigger auto-compaction before the
/// next LLM call.
const AUTO_COMPACT_CHARS: usize = 60_000;

// ---------------------------------------------------------------------------
// Helper: push a TurnEvent onto the shared queue.
// ---------------------------------------------------------------------------

fn push_event(queue: &Arc<Mutex<VecDeque<TurnEvent>>>, event: TurnEvent) {
    if let Ok(mut q) = queue.lock() {
        q.push_back(event);
    }
}

// ---------------------------------------------------------------------------
// Helper: stream-event callback that forwards tokens to the turn-event queue
// and checks the abort flag on each emission.
// ---------------------------------------------------------------------------

fn make_stream_callback(
    abort: &Arc<AtomicBool>,
    turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
) -> Box<dyn FnMut(&StreamEvent) -> bool + Send> {
    let abort_clone = Arc::clone(abort);
    let events_clone = Arc::clone(turn_events);
    Box::new(move |event: &StreamEvent| -> bool {
        if abort_clone.load(Ordering::SeqCst) {
            return false;
        }
        match event {
            StreamEvent::Token(s) => {
                push_event(&events_clone, TurnEvent::StreamToken(s.clone()));
            }
            StreamEvent::Reasoning(s) => {
                push_event(&events_clone, TurnEvent::StreamReasoning(s.clone()));
            }
            _ => {}
        }
        true
    })
}

// ---------------------------------------------------------------------------
// Helper: truncate a long tool output before it enters the conversation
// context. Preserves the head and appends a clear truncation marker.
// ---------------------------------------------------------------------------

fn truncate_tool_output(output: String) -> String {
    if output.len() <= TOOL_OUTPUT_MAX_CHARS {
        return output;
    }
    let mut result: String = output.chars().take(TOOL_OUTPUT_MAX_CHARS).collect();
    result.push_str(&format!(
        "\n...[truncated {} chars]",
        output.len() - TOOL_OUTPUT_MAX_CHARS
    ));
    result
}

// ---------------------------------------------------------------------------
// Helper: adaptive generation parameters.
// ---------------------------------------------------------------------------

/// Pick a `max_tokens` budget for the turn's next LLM call based on the
/// length of the user's request. Short requests need far fewer tokens than
/// the current hardcoded 4096 — big savings on small tasks.
fn adaptive_max_tokens(request_len: usize) -> u32 {
    if request_len <= 80 {
        800
    } else if request_len <= 400 {
        1600
    } else {
        4096
    }
}

/// Sum the character length of the conversation (user + assistant +
/// tool content) as a cheap proxy for context size.
fn conversation_chars(messages: &[ChatMessage]) -> usize {
    messages
        .iter()
        .map(|m| m.content.as_deref().map(str::len).unwrap_or(0))
        .sum()
}

/// Track repeated tool-call errors so the loop can recover instead of
/// burning iterations retrying the same failing tool.
#[derive(Default)]
struct ErrorTracker {
    consecutive: usize,
    total: usize,
    last_tool: String,
    last_error: String,
}

impl ErrorTracker {
    fn record(&mut self, tool_name: &str, error: &str, messages: &mut Vec<ChatMessage>) {
        if self.last_tool == tool_name {
            self.consecutive += 1;
        } else {
            self.consecutive = 1;
        }
        self.last_tool = tool_name.to_string();
        self.last_error = error.to_string();
        self.total += 1;

        // Inject a recovery note once the same tool keeps failing.
        if self.consecutive >= MAX_CONSECUTIVE_TOOL_ERRORS
            && !messages.iter().any(|m| {
                m.content
                    .as_deref()
                    .is_some_and(|c| c.contains("[System note]"))
            })
        {
            messages.push(ChatMessage::system(
                zesdex_domain::agent::prompt::error_recovery_note(tool_name, error),
            ));
        }
    }

    fn should_stop(&self) -> bool {
        self.consecutive >= MAX_CONSECUTIVE_TOOL_ERRORS * 2 || self.total >= MAX_TOTAL_TOOL_ERRORS
    }
}

// ---------------------------------------------------------------------------
// Helper: execute a single tool call, push events, return the result string.
// ---------------------------------------------------------------------------

async fn execute_tool_call<T: ToolExecutor>(
    tool_executor: &T,
    turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
    tc: &zesdex_domain::core::ToolCall,
) -> String {
    let name = &tc.function.name;
    let args = zesdex_domain::core::tool_call::sanitize_tool_arguments(&tc.function.arguments);

    debug!("executing tool: {name}");

    let output = match tool_executor.execute(name, &args).await {
        Ok(o) => o,
        Err(e) => format!("Error: {e}"),
    };

    let is_error = output.starts_with("Error:");
    let output = truncate_tool_output(output);

    push_event(
        turn_events,
        TurnEvent::ToolResult {
            tool_call_id: tc.id.clone(),
            tool_name: name.clone(),
            output: output.clone(),
            is_error,
            path: None,
        },
    );

    output
}

// ---------------------------------------------------------------------------
// Helper: emit usage event from optional LLM response metadata.
// ---------------------------------------------------------------------------

fn emit_usage(turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>, usage: Option<(u64, u64)>) {
    if let Some((tokens_in, tokens_out)) = usage {
        push_event(
            turn_events,
            TurnEvent::Usage {
                tokens_in,
                tokens_out,
            },
        );
    }
}

// ---------------------------------------------------------------------------
// Service implementation
// ---------------------------------------------------------------------------

/// Service implementation for executing an agent turn asynchronously.
///
/// The turn loop is adaptive and token-aware:
/// - No mandatory explore phase — the *agent* decides when to call the
///   `explore_codebase` tool (see the main prompt), so simple queries skip
///   exploration entirely.
/// - `max_tokens` / `temperature` adapt to the request length and phase.
/// - Repeated tool errors trigger a system recovery note and eventually
///   stop the loop instead of burning iterations.
/// - Tool outputs are truncated before entering context.
/// - Oversized histories are auto-compacted before the next LLM call.
pub struct AgentTurnServiceImpl<P: ProviderService, T: ToolExecutor> {
    provider: Arc<P>,
    tool_executor: Arc<T>,
    tool_defs: Vec<ToolDef>,
}

impl<P: ProviderService, T: ToolExecutor> AgentTurnServiceImpl<P, T> {
    pub fn new(provider: Arc<P>, tool_executor: Arc<T>, tool_defs: Vec<ToolDef>) -> Self {
        Self {
            provider,
            tool_executor,
            tool_defs,
        }
    }

    /// Execute a single LLM call with the current message list, handling
    /// streaming events and error reporting.
    async fn call_llm(
        &self,
        messages: &[ChatMessage],
        abort: &Arc<AtomicBool>,
        turn_events: &Arc<Mutex<VecDeque<TurnEvent>>>,
        max_tokens: u32,
        temperature: f32,
    ) -> Result<(ChatMessage, Option<(u64, u64)>), String> {
        let on_event = make_stream_callback(abort, turn_events);

        self.provider
            .chat_stream(
                messages,
                Some(self.tool_defs.clone()),
                Some(max_tokens),
                Some(temperature),
                on_event,
            )
            .await
            .map_err(|e| format!("LLM error: {e}"))
    }

    /// Auto-compact the history in place if it exceeds the threshold.
    ///
    /// Runs at most once per turn. Skips the synthetic system prompt that
    /// this service inserts at index 0.
    async fn auto_compact_if_needed(&self, messages: &mut Vec<ChatMessage>) {
        if conversation_chars(messages) <= AUTO_COMPACT_CHARS {
            return;
        }
        // Keep the system prompt (index 0) out of compaction.
        let sys = messages[0].clone();
        let mut rest: Vec<ChatMessage> = messages.drain(1..).collect();
        let before = rest.len();
        if let Err(e) = super::compact_messages_with_ai(&mut rest, self.provider.as_ref()).await {
            warn!("auto-compact failed (non-fatal): {e}");
        }
        info!(
            "auto-compacted history: {} messages -> {}",
            before,
            rest.len()
        );
        let mut rebuilt = Vec::with_capacity(rest.len() + 1);
        rebuilt.push(sys);
        rebuilt.extend(rest);
        *messages = rebuilt;
    }
}

impl<P: ProviderService, T: ToolExecutor> super::AgentTurnService for AgentTurnServiceImpl<P, T> {
    async fn run_turn(&self, mut params: AgentTurnParams) -> anyhow::Result<()> {
        info!(
            "Starting async agent turn with {} messages (model: {})",
            params.messages.len(),
            params.model
        );

        // Insert system prompt at position 0 once and keep it there for the
        // entire turn, avoiding per-iteration clones of the full message list.
        params
            .messages
            .insert(0, ChatMessage::system(main_agent_prompt()));
        let original_count = params.messages.len();

        // Estimate request complexity from the last user message.
        let request_len = params
            .messages
            .last()
            .and_then(|m| m.content.as_deref())
            .map(str::len)
            .unwrap_or(0);

        let mut errors = ErrorTracker::default();
        // Track whether the previous call produced tool calls — used to
        // lower temperature once the agent starts producing a final answer.
        let mut saw_tool_calls = false;

        for iteration in 0..MAX_TURN_ITERATIONS {
            // ── Check abort flag ────────────────────────────────────────
            if params.abort.load(Ordering::SeqCst) {
                params.abort.store(false, Ordering::SeqCst);
                push_event(
                    &params.turn_events,
                    TurnEvent::SystemNote {
                        kind: "info".into(),
                        message: "Turn aborted by user".into(),
                    },
                );
                break;
            }

            if errors.should_stop() {
                push_event(
                    &params.turn_events,
                    TurnEvent::SystemNote {
                        kind: "warn".into(),
                        message: "Stopping: repeated tool errors without progress".into(),
                    },
                );
                break;
            }

            debug!("agent turn iteration {iteration}");

            // ── Auto-compact oversized history before the LLM call ─────
            self.auto_compact_if_needed(&mut params.messages).await;

            // ── Adaptive generation parameters ─────────────────────────
            let max_tokens = adaptive_max_tokens(request_len);
            // Lower temperature while the agent is still choosing tools to
            // keep tool selection deterministic; raise it for the final
            // free-form answer.
            let temperature = if saw_tool_calls { 0.2 } else { 0.7 };

            // ── Stream start + call LLM ─────────────────────────────────
            push_event(&params.turn_events, TurnEvent::StreamStart);

            let result = self
                .call_llm(
                    &params.messages,
                    &params.abort,
                    &params.turn_events,
                    max_tokens,
                    temperature,
                )
                .await;

            match result {
                Ok((assistant_msg, usage)) => {
                    let content = assistant_msg.content.clone().unwrap_or_default();
                    let tool_calls = assistant_msg.tool_calls.clone().unwrap_or_default();

                    push_event(
                        &params.turn_events,
                        TurnEvent::StreamDone(assistant_msg.clone()),
                    );

                    emit_usage(&params.turn_events, usage);

                    // ── No tool calls → assistant is done ──────────────
                    if tool_calls.is_empty() {
                        params.messages.push(ChatMessage::assistant(Some(content)));
                        break;
                    }

                    saw_tool_calls = true;
                    params.messages.push(assistant_msg);

                    // ── Execute each tool call ──────────────────────────
                    for tc in &tool_calls {
                        let output =
                            execute_tool_call(self.tool_executor.as_ref(), &params.turn_events, tc)
                                .await;
                        if output.starts_with("Error:") {
                            errors.record(&tc.function.name, &output, &mut params.messages);
                        }
                        params
                            .messages
                            .push(ChatMessage::tool(tc.id.clone(), output));
                    }
                }
                Err(e) => {
                    warn!("{e}");
                    push_event(&params.turn_events, TurnEvent::Error(e));
                    break;
                }
            }
        }

        // Remove the synthetic sys_msg before shipping events to the TUI
        // so the transcript shows only the actual user/assistant/tool exchange.
        let compacted: Vec<ChatMessage> = params.messages.drain(original_count - 1..).collect();
        push_event(&params.turn_events, TurnEvent::Compacted(compacted));
        push_event(&params.turn_events, TurnEvent::Done);
        params.in_flight.store(false, Ordering::SeqCst);

        Ok(())
    }
}

// ---------------------------------------------------------------------------
// Conversation compaction
// ---------------------------------------------------------------------------

/// Maximum number of recent messages to preserve during compaction.
const COMPACT_KEEP_TAIL: usize = 6;

/// Compacts conversation history using AI summarisation.
///
/// Flow: if the message count exceeds `KEEP_TAIL + 2`, the oldest messages
/// are drained and summarised by the LLM. The summary is inserted as a
/// system message at the head of the remaining history.
pub async fn compact_messages_with_ai<P: ProviderService>(
    messages: &mut Vec<ChatMessage>,
    provider: &P,
) -> anyhow::Result<()> {
    if messages.len() <= COMPACT_KEEP_TAIL + 2 {
        return Ok(()); // Not enough messages to compact
    }

    let split_idx = messages.len() - COMPACT_KEEP_TAIL;
    let evicted: Vec<_> = messages.drain(..split_idx).collect();

    let mut summary_prompt = vec![ChatMessage::system(zesdex_domain::compaction_prompt())];
    summary_prompt.extend(evicted);
    summary_prompt.push(ChatMessage::user(
        "Please summarise our previous conversation above for context continuity.".to_string(),
    ));

    match provider
        .chat(&summary_prompt, None, Some(1024), Some(0.3))
        .await
    {
        Ok((summary_msg, _)) => {
            let summary_text = summary_msg
                .content
                .unwrap_or_else(|| "Previous context summarised.".to_string());
            let summary_node = ChatMessage::system(format!(
                "[AI Summary of Previous Conversation]\n{}",
                summary_text.trim()
            ));
            messages.insert(0, summary_node);
            Ok(())
        }
        Err(e) => {
            warn!("AI summarisation failed during compact, falling back to simple notice: {e}");
            messages.insert(
                0,
                ChatMessage::system(
                    "[Earlier conversation messages compacted to save context window]".to_string(),
                ),
            );
            Ok(())
        }
    }
}

// ---------------------------------------------------------------------------
// Tests
// ---------------------------------------------------------------------------

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn truncate_short_output_is_unchanged() {
        let out = "short".to_string();
        assert_eq!(truncate_tool_output(out.clone()), out);
    }

    #[test]
    fn truncate_long_output_preserves_head_and_marks_cut() {
        let long = "x".repeat(TOOL_OUTPUT_MAX_CHARS + 500);
        let truncated = truncate_tool_output(long.clone());
        assert!(truncated.len() < long.len());
        assert!(truncated.contains("...[truncated"));
        assert!(truncated.starts_with("xxx"));
    }

    #[test]
    fn adaptive_max_tokens_scales_with_request_len() {
        assert_eq!(adaptive_max_tokens(10), 800);
        assert_eq!(adaptive_max_tokens(200), 1600);
        assert_eq!(adaptive_max_tokens(5000), 4096);
    }

    #[test]
    fn error_tracker_injects_recovery_note_after_repeats() {
        let mut tracker = ErrorTracker::default();
        let mut messages: Vec<ChatMessage> = Vec::new();
        tracker.record("read", "Error: File not found", &mut messages);
        tracker.record("read", "Error: File not found", &mut messages);
        assert!(!tracker.should_stop());
        // Third consecutive failure → recovery note injected.
        tracker.record("read", "Error: File not found", &mut messages);
        assert!(messages.iter().any(|m| m
            .content
            .as_deref()
            .is_some_and(|c| c.contains("[System note]"))));
    }

    #[test]
    fn error_tracker_stops_after_too_many_errors() {
        let mut tracker = ErrorTracker::default();
        let mut messages: Vec<ChatMessage> = Vec::new();
        for i in 0..MAX_TOTAL_TOOL_ERRORS {
            tracker.record("bash", &format!("Error: boom {i}"), &mut messages);
        }
        assert!(tracker.should_stop());
    }

    #[test]
    fn conversation_chars_sums_content_only() {
        let messages = vec![
            ChatMessage::system("sys".to_string()),
            ChatMessage::user("hello world".to_string()),
            ChatMessage::tool("id".to_string(), "output".to_string()),
        ];
        assert_eq!(conversation_chars(&messages), 3 + 11 + 6);
    }
}