procyon 0.1.2

use std::time::Duration;

use color_eyre::{eyre::bail, Result};
use serde::Deserialize;

use super::{ContentPart, Message, ToolDefinition};
use crate::config::AppConfig;
use crate::llm::LlmClient;
use crate::tools::ToolRegistry;

/// Configuration for spawning a sub-agent.
#[derive(Deserialize)]
pub struct SubAgentConfig {
    /// System prompt for the sub-agent (persona + context).
    pub system_prompt: String,
    /// The user message to process.
    pub message: String,
    /// Model override (None = use parent config).
    pub model: Option<String>,
    /// Max tokens override (None = use parent config).
    pub max_tokens: Option<u32>,
    /// Maximum number of tool round-trips (None = 20).
    pub max_rounds: Option<usize>,
    /// Tool names the sub-agent may use. None = all tools except spawn_agent.
    pub allowed_tools: Option<Vec<String>>,
    /// Deadline for a single LLM request, in seconds (None = 120).
    ///
    /// Party mode wraps the whole sub-agent in a timeout, but every other caller had none: a
    /// provider that accepts the connection and then stops sending bytes left `spawn_agent` and
    /// `talk_to` blocked with no upper bound and nothing to report.
    pub timeout_secs: Option<u64>,
}

/// Per-request deadline when the caller states none.
const DEFAULT_REQUEST_TIMEOUT: Duration = Duration::from_secs(120);

/// The response from a sub-agent.
pub struct SubAgentResponse {
    /// Final text response from the sub-agent.
    pub text: String,
    /// Number of LLM round trips made.
    pub round_trips: usize,
}

/// Runs a sub-agent with its own LLM client and message history.
///
/// The sub-agent gets a simplified loop: no streaming to UI, no session logging,
/// no compaction. It runs until the LLM produces a text-only response (no tool calls).
pub async fn run_subagent(
    config: &AppConfig,
    subagent_config: SubAgentConfig,
    registry: &ToolRegistry,
) -> Result<SubAgentResponse> {
    let client = build_client(
        config,
        subagent_config.model.as_deref(),
        subagent_config.max_tokens,
    )?;

    // Build tool definitions, filtering by allowed_tools
    let all_defs = registry.definitions();
    let tool_defs: Vec<ToolDefinition> = match &subagent_config.allowed_tools {
        Some(allowed) => all_defs
            .into_iter()
            .filter(|d| allowed.contains(&d.name))
            .collect(),
        None => all_defs
            .into_iter()
            .filter(|d| d.name != "spawn_agent")
            .collect(),
    };

    // A sub-agent has no compaction, so the window is a hard ceiling rather than a threshold to
    // relieve. Without this it simply ran until the provider refused the request, and the refusal
    // surfaced as a generic failure that discarded every round of work already paid for.
    let model = subagent_config
        .model
        .as_deref()
        .unwrap_or(&config.default_model);
    let max_output = subagent_config.max_tokens.unwrap_or(config.max_tokens) as usize;
    let envelope = crate::budget::price_envelope(Some(&subagent_config.system_prompt), &tool_defs);
    // The same detour the interactive turn loop takes: Ollama's real window is a server setting
    // `/api/ps` reports, not the hardcoded conservative fallback `context_window` assumes for it.
    // Without this every `--bench`/`--exec` run against Ollama budgeted against 4,096 even when
    // `OLLAMA_CONTEXT_LENGTH` was raised, and stopped after round one blaming a ceiling that no
    // longer existed.
    let window = crate::llm::ollama_context_length(config, model)
        .await
        .unwrap_or_else(|| crate::budget::context_window(config.provider, model));
    let ceiling = crate::budget::threshold_tokens(crate::budget::usable_window(window, max_output));

    let mut history = vec![Message::user(&subagent_config.message)];
    let mut round_trips = 0;
    // Zero is filtered rather than honoured: it arrives from a model-supplied `spawn_agent` call,
    // and `Duration::from_secs(0)` expires on the first poll, so every round would fail instantly
    // with a message blaming a timeout the caller never meant to set.
    let request_timeout = subagent_config
        .timeout_secs
        .filter(|secs| *secs > 0)
        .map(Duration::from_secs)
        .unwrap_or(DEFAULT_REQUEST_TIMEOUT);

    // Simplified agent loop: run until text response or max rounds. The guard is the loop
    // condition rather than a break after the increment, so `round_trips` never reports a request
    // that was not made.
    let max_rounds = subagent_config.max_rounds.unwrap_or(20);
    let mut retried_empty = false;
    let mut hit_ceiling = false;
    while round_trips < max_rounds {
        // Checked before the request rather than after the tool results are appended, so the round
        // that would have overflowed is never sent.
        if should_stop(
            envelope + crate::budget::price_history(&history),
            ceiling,
            last_assistant_text(&history).is_some(),
        ) {
            hit_ceiling = true;
            break;
        }
        round_trips += 1;

        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel::<String>();
        // The timeout wraps one request, not the whole loop: a sub-agent legitimately doing many
        // rounds of tool work must not be cut off for taking a while overall.
        let outcome = match tokio::time::timeout(
            request_timeout,
            client.send_message_streaming(
                &history,
                Some(&tool_defs),
                Some(&subagent_config.system_prompt),
                &tx,
            ),
        )
        .await
        {
            Ok(result) => result?,
            Err(_) => bail!(
                "Sub-agent LLM request timed out after {}s (round {})",
                request_timeout.as_secs(),
                round_trips
            ),
        };

        // Drain any remaining chunks
        drop(tx);
        while rx.try_recv().is_ok() {}

        let blocks = outcome.blocks;
        if blocks.is_empty() {
            // One more chance, and only one: nothing is appended to the history before retrying,
            // so the next request is byte-identical. A provider that keeps returning an empty
            // choice would otherwise be paid for `max_rounds` copies of the same request.
            let last_was_tool_results = history.last().is_some_and(|m| {
                m.content
                    .iter()
                    .any(|p| matches!(p, ContentPart::ToolResult { .. }))
            });
            if last_was_tool_results && !retried_empty {
                retried_empty = true;
                continue;
            }
            break;
        }
        retried_empty = false;

        history.push(Message::assistant(blocks.clone()));

        // Check for tool calls
        let tool_uses: Vec<_> = blocks
            .iter()
            .filter_map(|b| match b {
                ContentPart::ToolUse { id, name, input } => {
                    Some((id.clone(), name.clone(), input.clone()))
                }
                _ => None,
            })
            .collect();

        if tool_uses.is_empty() {
            // No tool calls — extract text and return
            let text = blocks
                .iter()
                .filter_map(|b| match b {
                    ContentPart::Text { text } => Some(text.as_str()),
                    _ => None,
                })
                .collect::<Vec<_>>()
                .join("");

            return Ok(SubAgentResponse { text, round_trips });
        }

        // Execute tool calls
        let mut results = Vec::with_capacity(tool_uses.len());
        for (id, name, input) in tool_uses {
            let outcome = registry.execute(&name, input).await;
            let result_str = super::clamp_tool_result(match outcome {
                Ok(r) => r,
                Err(e) => {
                    let err_str = e.to_string();
                    // Sanitize error: remove internal details like file paths and stack traces
                    let sanitized = err_str
                        .lines()
                        .next()
                        .unwrap_or("Unknown error")
                        .to_string();
                    format!("Error: {}", sanitized)
                }
            });
            results.push((id, result_str));
        }

        history.push(Message::tool_results(results));
    }

    // If we exhausted rounds — or ran out of context — extract whatever text we have
    let text = last_assistant_text(&history)
        .unwrap_or_else(|| "Sub-agent did not produce a response".to_string());

    // Stated rather than silent: the caller is a model deciding what to do next, and partial work
    // reported as if it were finished is worse than partial work labelled as partial.
    let text = if hit_ceiling {
        format!(
            "{}\n\n[Sub-agent stopped after {} round(s): its context window was full. The answer \
             above is partial. Split the task or narrow the tools it may use.]",
            text, round_trips
        )
    } else {
        text
    };

    Ok(SubAgentResponse { text, round_trips })
}

/// Whether to stop rather than send the next round.
///
/// A sub-agent has no compaction, so stopping is its only lever — and stopping is worth pulling
/// only when it preserves something. That is the condition here: an answer already exists in the
/// history, and continuing would risk the provider refusing the request and taking it with them.
///
/// Being over the ceiling with *nothing* to return is not a reason to stop, which is the bug this
/// replaces. The envelope — a system prompt plus every tool schema — is irreducible; when it alone
/// exceeds the ceiling, every round is over budget from the first one, and the loop exited having
/// answered "context window full" with no answer in it. That is not a rare configuration: a
/// 4,096-token Ollama server with `max_tokens = 4096` leaves `usable_window` a quarter of the
/// window, less than the tool definitions cost, so `spawn_agent`, `talk_to`, party mode and
/// `--exec` all failed identically and before doing any work.
///
/// Sending anyway puts the decision where the information is. A provider that cannot take the
/// request says so, and Ollama's silent truncation is a risk the interactive path already names out
/// loud rather than pre-empting.
fn should_stop(priced: usize, ceiling: usize, has_answer: bool) -> bool {
    has_answer && priced > ceiling
}

/// The last assistant text in the history, if there is any.
///
/// Both the stop condition and the return value need it: "is there something to keep" and "what do
/// we keep" have to be the same question, or the loop can stop to preserve an answer it then fails
/// to find.
fn last_assistant_text(history: &[Message]) -> Option<String> {
    history.iter().rev().find_map(|message| {
        if message.role != super::Role::Assistant {
            return None;
        }
        let text: String = message
            .content
            .iter()
            .filter_map(|part| match part {
                ContentPart::Text { text } => Some(text.as_str()),
                _ => None,
            })
            .collect();
        (!text.is_empty()).then_some(text)
    })
}

/// Builds an LlmClient from config, optionally overriding the model.
fn build_client(
    config: &AppConfig,
    model_override: Option<&str>,
    max_tokens: Option<u32>,
) -> Result<LlmClient> {
    let mut cfg = config.clone();
    if let Some(model) = model_override {
        cfg.default_model = model.to_string();
    }
    if let Some(tokens) = max_tokens {
        cfg.max_tokens = tokens;
    }
    LlmClient::from_config(&cfg)
}

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

    #[test]
    fn subagent_config_holds_all_fields() {
        let config = SubAgentConfig {
            system_prompt: "You are Tyler.".to_string(),
            message: "Design the contract.".to_string(),
            model: Some("claude-haiku".to_string()),
            max_tokens: Some(2048),
            max_rounds: Some(10),
            allowed_tools: Some(vec!["read_file".to_string()]),
            timeout_secs: Some(30),
        };

        assert_eq!(config.system_prompt, "You are Tyler.");
        assert_eq!(config.message, "Design the contract.");
        assert_eq!(config.model.as_deref(), Some("claude-haiku"));
        assert_eq!(config.max_tokens, Some(2048));
        assert_eq!(config.max_rounds, Some(10));
        assert_eq!(config.allowed_tools.as_ref().unwrap().len(), 1);
        assert_eq!(config.timeout_secs, Some(30));
    }

    // The default matters: every caller but party mode passes None, and before this there was no
    // upper bound at all on a request that never returns.
    #[test]
    fn the_default_request_timeout_is_bounded() {
        assert_eq!(DEFAULT_REQUEST_TIMEOUT, Duration::from_secs(120));
    }

    /// The resolution `run_subagent` performs, extracted so it can be checked without a provider.
    fn resolve_timeout(configured: Option<u64>) -> Duration {
        configured
            .filter(|secs| *secs > 0)
            .map(Duration::from_secs)
            .unwrap_or(DEFAULT_REQUEST_TIMEOUT)
    }

    // `timeout_secs` arrives from a model-supplied spawn_agent call. Zero would expire on the
    // first poll, so every round would fail instantly blaming a limit nobody set.
    #[test]
    fn a_zero_timeout_falls_back_to_the_default() {
        assert_eq!(resolve_timeout(Some(0)), DEFAULT_REQUEST_TIMEOUT);
        assert_eq!(resolve_timeout(None), DEFAULT_REQUEST_TIMEOUT);
        assert_eq!(resolve_timeout(Some(30)), Duration::from_secs(30));
    }

    // Regression: a 4,096-token Ollama window with `max_tokens = 4096` leaves a ceiling smaller
    // than the tool schemas, so every sub-agent stopped before its first request and answered
    // "context window full" with nothing in it.
    #[test]
    fn being_over_budget_with_nothing_to_show_is_not_a_reason_to_stop() {
        assert!(!should_stop(100_000, 800, false));
    }

    #[test]
    fn a_round_that_would_risk_an_existing_answer_stops() {
        assert!(should_stop(801, 800, true));
        assert!(!should_stop(800, 800, true));
    }

    #[test]
    fn the_answer_kept_is_the_last_thing_the_model_said() {
        let history = vec![
            Message::user("go"),
            Message::assistant(vec![ContentPart::Text {
                text: "first".to_string(),
            }]),
            Message::assistant(vec![ContentPart::Text {
                text: "second".to_string(),
            }]),
        ];
        assert_eq!(last_assistant_text(&history).as_deref(), Some("second"));
    }

    #[test]
    fn a_history_with_no_answer_in_it_reports_none() {
        assert!(last_assistant_text(&[Message::user("go")]).is_none());
        // A turn that only called tools has no text to keep, which is what the stop condition
        // above has to distinguish from a real answer.
        let tools_only = vec![Message::assistant(vec![ContentPart::ToolUse {
            id: "1".to_string(),
            name: "read_file".to_string(),
            input: serde_json::json!({}),
        }])];
        assert!(last_assistant_text(&tools_only).is_none());
    }

    #[test]
    fn subagent_response_holds_text_and_round_trips() {
        let response = SubAgentResponse {
            text: "Done.".to_string(),
            round_trips: 3,
        };
        assert_eq!(response.text, "Done.");
        assert_eq!(response.round_trips, 3);
    }
}