// SPDX-License-Identifier: Apache-2.0
// Copyright 2026 Mason Stallmo

//! Delegation of an investigation to a nested agent.
//!
//! A question like "which of these forty files defines the retry policy" is
//! answered by reading a great deal and keeping almost none of it. In the
//! session's own context that reading fills the window and ages out the
//! results an in-progress edit depends on.
//!
//! [`TaskTool`] runs a nested [`Agent`] over the same workspace with a
//! conversation of its own, and only its final message comes back.

use std::sync::Arc;
use std::sync::atomic::AtomicBool;

use async_trait::async_trait;
use ein_agent::model_clients::ModelClient;
use ein_agent::tools::{Param, ReadOnlyToolSet, Tool, ToolDef, ToolEffect, ToolResult, ToolSet};
use ein_agent::{Agent, budget_for_context_window};
use ein_plugin::model_client::{Message, Role};
use ein_wasm::{ModelClientSpec, PluginRuntime};
use serde::Deserialize;
use tracing::{info, warn};

/// The share of the session's context window a delegated task may use.
///
/// A task given the whole window can read until it triggers its own eviction,
/// after which it reports from a summary of what it read rather than from the
/// reading.
const WINDOW_FRACTION: f64 = 0.5;

/// The `Task` tool's arguments.
///
/// `deny_unknown_fields` enforces the schema's `additionalProperties: false`,
/// which a provider may not.
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
struct TaskArgs {
    description: String,
    prompt: String,
}

/// The `Task` tool, which runs a read-only nested agent and returns its
/// answer.
pub struct TaskTool {
    runtime: PluginRuntime,
    model_spec: ModelClientSpec,
    /// Loaded once and shared by every call, since loading the plugin
    /// directory instantiates every plugin in it. This is a separate load from
    /// the session's own set, which holds `Task`.
    tools: Arc<dyn ToolSet>,
    /// The session's cancel flag, so stopping the turn also stops a task
    /// running inside it.
    cancel: Arc<AtomicBool>,
    context_window: Option<u32>,
}

impl TaskTool {
    pub fn new(
        runtime: PluginRuntime,
        model_spec: ModelClientSpec,
        tools: Arc<dyn ToolSet>,
        cancel: Arc<AtomicBool>,
        context_window: Option<u32>,
    ) -> Self {
        Self {
            runtime,
            model_spec,
            tools,
            cancel,
            context_window,
        }
    }
}

#[async_trait]
impl Tool for TaskTool {
    fn name(&self) -> &str {
        "Task"
    }

    fn schema(&self) -> ToolDef {
        ToolDef::function(
            self.name(),
            "Delegate an investigation to a separate agent that reads the \
             workspace and reports back. Use it when answering something means \
             reading a lot and keeping little: locating where a behaviour is \
             implemented, working out how a subsystem fits together, checking \
             whether a pattern appears anywhere else. The agent's reading stays \
             in its own context and never enters yours; only its written answer \
             comes back.\n\n\
             Do not use it for work you can do in one or two calls — a single \
             known file, a single search — where delegating costs a whole model \
             turn to save nothing. Do not use it to make changes: the agent can \
             only read, and cannot write, edit, run commands, or reply to the \
             user. Act on what it reports yourself.\n\n\
             It cannot ask you anything and you cannot steer it once it starts, \
             so `prompt` has to carry everything it needs: what to look for, \
             where to start, and what a useful answer contains.",
        )
        .with_param(Param::string(
            "description",
            "A few words naming the task, shown to the user while it runs.",
        ))
        .with_param(Param::string(
            "prompt",
            "The complete instruction for the agent: what to investigate, where \
             to begin, and what its answer should tell you.",
        ))
        .build()
    }

    fn primary_arg(&self) -> Option<&str> {
        Some("description")
    }

    /// Returns `ReadOnly`, which its tool set enforces, so a batch of
    /// delegations runs concurrently.
    fn effect(&self) -> ToolEffect {
        ToolEffect::ReadOnly
    }

    async fn call(&self, id: &str, args: &str) -> anyhow::Result<ToolResult> {
        let args: TaskArgs = serde_json::from_str(args)?;

        info!("[task] delegating: {}", args.description);

        let model = self
            .runtime
            .model_clients()
            .new_session(&self.model_spec)
            .await?;

        let answer = delegate(
            model,
            self.tools.clone(),
            self.cancel.clone(),
            delegate_window(self.context_window),
            args.prompt,
        )
        .await?;

        Ok(report(id, &args.description, &answer))
    }
}

/// Runs one delegated turn and returns what it wrote.
///
/// Generic rather than taking a [`TaskTool`], so it can be tested without a
/// WASM runtime or a provider.
///
/// The agent sees only the read-only tools in `tools` and has no event handler
/// or message sink. Its events would otherwise reach the session's stream
/// looking like the session's own: its text drawn as the session's reply, and
/// its usage overwriting the parent's count.
async fn delegate<MC: ModelClient, TS: ToolSet>(
    model: MC,
    tools: TS,
    cancel: Arc<AtomicBool>,
    window: Option<u32>,
    prompt: String,
) -> anyhow::Result<String> {
    let mut builder = Agent::builder_with_tool_set(model, ReadOnlyToolSet::new(tools))
        .with_system_preamble(vec![Message {
            role: Role::System,
            content: Some(preamble()),
            tool_calls: None,
            tool_call_id: None,
        }])
        .with_cancel_handle(cancel);

    if let Some(window) = window {
        builder = builder
            .with_context_window(window)
            .tool_result_budget(budget_for_context_window(window));
    }

    Ok(builder
        .build()
        .chat(prompt)
        .await?
        .content
        .unwrap_or_default())
}

/// Returns the window a delegated task gets, [`WINDOW_FRACTION`] of the
/// session's and at least 1.
fn delegate_window(session: Option<u32>) -> Option<u32> {
    session.map(|window| ((f64::from(window) * WINDOW_FRACTION) as u32).max(1))
}

/// Returns the task's answer as the delegating agent reads it.
///
/// A task that wrote nothing is reported as that. An empty result reads as a
/// tool that silently did nothing, which suggests calling it again.
fn report(id: &str, description: &str, answer: &str) -> ToolResult {
    if !answer.trim().is_empty() {
        return ToolResult::new(id, answer.to_owned());
    }

    warn!("[task] '{description}' produced no answer");

    ToolResult::new(
        id,
        format!(
            "The delegated agent for `{description}` ended without writing an \
             answer. It may have spent its tool calls before concluding, or the \
             turn may have been cancelled. Narrow the task, or investigate it \
             directly."
        ),
    )
}

/// Returns the delegated agent's system preamble.
///
/// The session's preamble does not apply: the task has no core memory to write
/// and no user to address. It is told what it cannot infer, that its context
/// is discarded and only its reply survives.
fn preamble() -> String {
    "You are a sub-agent. Another agent has delegated one investigation to you \
     and is waiting on your answer.\n\
     \n\
     You can only read: the file and search tools, and nothing that changes \
     anything. You cannot run commands, edit files, or reply to the user \
     directly. Do not try to; report what you found instead.\n\
     \n\
     Your reply is the only thing that survives. Everything you read is \
     discarded when you finish, and the agent waiting on you sees none of it — \
     so a reply saying where to look, without saying what is there, is of no \
     use to it. Quote the file paths, line numbers, symbol names and error \
     text you found, exactly as they appear. An approximation of a path is \
     worse than leaving it out, because it will be acted on.\n\
     \n\
     Answer in one message when you are done reading. Say plainly what you did \
     not manage to establish rather than filling it in."
        .to_string()
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::Mutex;
    use std::sync::atomic::Ordering;

    use ein_agent::tools::NativeToolSet;
    use ein_plugin::model_client::{
        Choice, CompletionResponse, FinishReason, FunctionCall, ToolCall,
    };

    /// A tool that counts its calls, so a test can assert it never ran.
    struct Recorder {
        name: &'static str,
        effect: ToolEffect,
        calls: Arc<Mutex<usize>>,
    }

    impl Recorder {
        fn new(name: &'static str, effect: ToolEffect) -> Self {
            Self {
                name,
                effect,
                calls: Arc::new(Mutex::new(0)),
            }
        }
    }

    #[async_trait]
    impl Tool for Recorder {
        fn name(&self) -> &str {
            self.name
        }

        fn schema(&self) -> ToolDef {
            ToolDef::function(self.name, "a tool").build()
        }

        fn effect(&self) -> ToolEffect {
            self.effect
        }

        async fn call(&self, id: &str, _args: &str) -> anyhow::Result<ToolResult> {
            *self.calls.lock().unwrap() += 1;

            Ok(ToolResult::new(id, format!("{} ran", self.name)))
        }
    }

    /// A model client that plays a fixed script of turns and keeps the
    /// requests it was sent.
    struct Script {
        turns: Mutex<Vec<CompletionResponse>>,
        seen: Arc<Mutex<Vec<Vec<Message>>>>,
    }

    impl Script {
        fn new(turns: Vec<CompletionResponse>) -> Self {
            Self {
                turns: Mutex::new(turns),
                seen: Arc::new(Mutex::new(Vec::new())),
            }
        }
    }

    #[async_trait]
    impl ModelClient for Script {
        async fn complete(
            &mut self,
            messages: &[Message],
            _tools: &[ToolDef],
        ) -> anyhow::Result<CompletionResponse> {
            self.seen.lock().unwrap().push(messages.to_vec());

            let mut turns = self.turns.lock().unwrap();

            Ok(match turns.is_empty() {
                true => says("done"),
                false => turns.remove(0),
            })
        }
    }

    fn says(text: &str) -> CompletionResponse {
        CompletionResponse {
            choices: vec![Choice {
                index: None,
                finish_reason: FinishReason::Stop,
                message: Message {
                    role: Role::Assistant,
                    content: Some(text.to_owned()),
                    tool_calls: None,
                    tool_call_id: None,
                },
            }],
            usage: None,
            error: None,
            provenance: None,
        }
    }

    fn calls(name: &str) -> CompletionResponse {
        CompletionResponse {
            choices: vec![Choice {
                index: None,
                finish_reason: FinishReason::ToolCalls,
                message: Message {
                    role: Role::Assistant,
                    content: None,
                    tool_calls: Some(vec![ToolCall::Function {
                        id: "1".to_owned(),
                        index: 0,
                        function: FunctionCall {
                            name: name.to_owned(),
                            arguments: "{}".to_owned(),
                        },
                    }]),
                    tool_call_id: None,
                },
            }],
            usage: None,
            error: None,
            provenance: None,
        }
    }

    fn quiet() -> Arc<AtomicBool> {
        Arc::new(AtomicBool::new(false))
    }

    /// A delegated agent given a mutating tool cannot call it.
    ///
    /// Its changes would be ones the delegating agent never saw.
    #[tokio::test]
    async fn a_delegated_agent_cannot_reach_a_tool_that_changes_anything() {
        let editor = Recorder::new("Edit", ToolEffect::Mutating);
        let edits = editor.calls.clone();
        let reader = Recorder::new("Read", ToolEffect::ReadOnly);
        let reads = reader.calls.clone();

        let mut tools = NativeToolSet::default();
        tools.insert(editor);
        tools.insert(reader);

        let answer = delegate(
            Script::new(vec![calls("Read"), calls("Edit"), says("what I found")]),
            tools,
            quiet(),
            None,
            "look into it".to_owned(),
        )
        .await
        .unwrap();

        assert_eq!(*reads.lock().unwrap(), 1);
        assert_eq!(*edits.lock().unwrap(), 0, "a delegated task edited a file");
        assert_eq!(answer, "what I found");
    }

    /// Only the task's final message comes back, not its tool results.
    #[tokio::test]
    async fn only_the_answer_crosses_back_not_the_reading() {
        let reader = Recorder::new("Read", ToolEffect::ReadOnly);

        let mut tools = NativeToolSet::default();
        tools.insert(reader);

        let answer = delegate(
            Script::new(vec![
                calls("Read"),
                calls("Read"),
                says("the retry policy is in retry.rs"),
            ]),
            tools,
            quiet(),
            None,
            "where is the retry policy".to_owned(),
        )
        .await
        .unwrap();

        assert_eq!(answer, "the retry policy is in retry.rs");
        assert!(!answer.contains("ran"), "a tool result came back with it");
    }

    /// A tool that answers by delegating, so the parent's events can be
    /// watched.
    struct Delegating {
        answer: &'static str,
    }

    #[async_trait]
    impl Tool for Delegating {
        fn name(&self) -> &str {
            "Task"
        }

        fn schema(&self) -> ToolDef {
            ToolDef::function(self.name(), "delegate").build()
        }

        fn effect(&self) -> ToolEffect {
            ToolEffect::Unknown
        }

        async fn call(&self, id: &str, _args: &str) -> anyhow::Result<ToolResult> {
            let answer = delegate(
                Script::new(vec![says(self.answer)]),
                NativeToolSet::default(),
                quiet(),
                None,
                "look into it".to_owned(),
            )
            .await?;

            Ok(ToolResult::new(id, answer))
        }
    }

    /// Nothing a delegated agent writes reaches the session's event stream.
    ///
    /// `delegate` installing no event handler is all that keeps a task's text
    /// from being drawn as the session's reply.
    #[tokio::test]
    async fn a_delegated_turn_says_nothing_to_the_session() {
        const NESTED: &str = "the retry policy is in retry.rs";

        let mut tools = NativeToolSet::default();
        tools.insert(Delegating { answer: NESTED });

        let parent = Script::new(vec![calls("Task"), says("done")]);
        let seen = parent.seen.clone();

        let heard = Arc::new(Mutex::new(Vec::new()));
        let collect = heard.clone();

        let answer = Agent::builder_with_tool_set(parent, tools)
            .with_event_handler(move |event| {
                let heard = collect.clone();

                async move {
                    heard.lock().unwrap().push(event);
                }
            })
            .build()
            .chat("where is the retry policy".to_owned())
            .await
            .unwrap()
            .content
            .unwrap_or_default();

        for event in heard.lock().unwrap().iter() {
            if let ein_agent::AgentEvent::ContentDelta(text)
            | ein_agent::AgentEvent::ReasoningDelta(text) = event
            {
                assert!(
                    !text.contains(NESTED),
                    "a delegated turn's text reached the session's stream: {text}"
                );
            }
        }

        assert_eq!(answer, "done", "the session answered for itself");

        // The answer did reach the parent, so the check above did not pass on
        // a task that produced nothing.
        let last = seen.lock().unwrap().last().cloned().unwrap_or_default();
        assert!(
            last.iter()
                .any(|m| m.content.as_deref().unwrap_or_default().contains(NESTED)),
            "the task's answer never reached the agent that asked for it",
        );
    }

    /// A cancel raised before the task starts stops it, and stays raised.
    ///
    /// An agent clearing a flag it was given would swallow a cancel raised
    /// while the task started, and neither agent would stop.
    #[tokio::test]
    async fn a_cancel_raised_before_the_task_starts_still_stops_it() {
        let reader = Recorder::new("Read", ToolEffect::ReadOnly);
        let reads = reader.calls.clone();

        let mut tools = NativeToolSet::default();
        tools.insert(reader);

        let cancel = Arc::new(AtomicBool::new(true));

        delegate(
            Script::new(vec![calls("Read"), says("done")]),
            tools,
            cancel.clone(),
            None,
            "look into it".to_owned(),
        )
        .await
        .unwrap();

        assert_eq!(*reads.lock().unwrap(), 0, "the task ran after a cancel");
        assert!(
            cancel.load(Ordering::Relaxed),
            "the delegated agent cleared the session's cancel flag"
        );
    }

    /// The delegated agent's preamble tells it its context is discarded.
    ///
    /// Otherwise it answers by pointing at what it read, which the delegating
    /// agent cannot see.
    #[tokio::test]
    async fn the_delegated_agent_is_told_its_context_is_discarded() {
        let script = Script::new(vec![says("ok")]);
        let seen = script.seen.clone();

        delegate(
            script,
            NativeToolSet::default(),
            quiet(),
            None,
            "go".to_owned(),
        )
        .await
        .unwrap();

        let first = seen.lock().unwrap()[0][0].clone();

        assert_eq!(first.role, Role::System);
        assert!(
            first.content.as_deref().unwrap().contains("discarded"),
            "{first:?}"
        );
    }

    /// A delegated task gets half the session's window.
    #[test]
    fn a_delegated_task_gets_less_of_the_window_than_the_session() {
        assert_eq!(delegate_window(Some(200_000)), Some(100_000));
        assert_eq!(delegate_window(None), None);
    }

    /// A window that would round to zero is kept at 1.
    ///
    /// Zero disables pressure management.
    #[test]
    fn a_tiny_window_still_leaves_something_to_manage() {
        assert_eq!(delegate_window(Some(1)), Some(1));
    }

    /// A task that wrote nothing is reported as that, not as an empty result.
    ///
    /// An empty result suggests calling the tool again, which costs another
    /// delegated turn.
    #[test]
    fn a_task_that_wrote_no_answer_says_so_rather_than_coming_back_empty() {
        for answer in ["", "   \n "] {
            let result = report("1", "find the retry policy", answer);

            assert!(
                result.content.contains("without writing an answer"),
                "{:?}",
                result.content
            );
            assert!(result.content.contains("find the retry policy"));
        }
    }

    #[test]
    fn an_answer_comes_back_as_it_was_written() {
        assert_eq!(report("1", "d", "the answer").content, "the answer");
    }
}