//! Tasks agents delegate to the app (see [`crate::orchestrator`]). Each runs
//! in a subagent thread: a thread of the parent's project, in its folder,
//! that the sidebar leaves out and the subagents panel opens. The parent's
//! `delegate_task` call shows it working. When it finishes, its result goes
//! to the call waiting on it (`mode: "wait"`), or else to the parent's agent
//! as a message the transcript doesn't show, which wakes it.

use std::time::{Duration, SystemTime};

use gpui_kit::*;
use serde_json::{Value, json};

use super::Workspace;
use crate::models::{Catalog, ProbeStatus};
use crate::orchestrator::{self, CAPABILITIES, DELEGATE, TASK_CANCEL, TASK_STATUS, ToolRequest};
use crate::session::{
    ChatItem, Delegation, NewSession, Reasoning, Session, SessionEvent, SessionStatus, SubagentMember, SubagentParent,
    TaskStatus, brief_input, to_millis,
};
use crate::settings;

/// How long `mode: "wait"` waits when the call doesn't say.
const DEFAULT_WAIT: Duration = Duration::from_secs(10 * 60);
/// The most of a subagent's answer that goes back in a message.
const RESULT_LIMIT: usize = 40_000;

/// Whoever is waiting for a tool call's answer.
pub(super) enum Answer {
    /// Claude or Pi, through the MCP server.
    Mcp(smol::channel::Sender<Result<Value, String>>),
    /// jcode, whose session tool call its thread answers.
    Jcode { session: WeakEntity<Session>, call_id: String },
}

impl Answer {
    pub(super) fn send(self, outcome: Result<Value, String>, cx: &mut App) {
        match self {
            Answer::Mcp(reply) => {
                let _ = reply.try_send(outcome);
            }
            Answer::Jcode { session, call_id } => {
                let _ = session.update(cx, |session, _| session.answer_app_tool(&call_id, outcome));
            }
        }
    }
}

/// A `delegate_task` call waiting for its task to finish.
struct Waiter {
    id: u64,
    task: String,
    answer: Answer,
    _timeout: Task<()>,
}

#[derive(Default)]
pub(super) struct Delegations {
    waiting: Vec<Waiter>,
    next_waiter: u64,
    _requests: Option<Task<()>>,
}

/// Where a task stands, as the tools report it.
#[derive(Clone, Copy, PartialEq, Eq)]
enum TaskState {
    Running,
    /// Its agent waits for the user: a permission prompt or a question.
    WaitingForUser,
    Completed,
    Failed,
    Cancelled,
    /// The app stopped while it worked.
    Interrupted,
}

impl TaskState {
    fn of(session: &Session) -> Self {
        let cancelled = session.delegation.as_ref().is_some_and(|delegation| delegation.cancelled);
        let finished = session.items.iter().any(|item| matches!(item, ChatItem::TurnEnd { .. }));
        if cancelled && !session.is_working() {
            TaskState::Cancelled
        } else if session.needs_attention() {
            TaskState::WaitingForUser
        } else if session.is_working() || session.preparing {
            TaskState::Running
        } else if session.status == SessionStatus::Exited || (finished && session.ended_with_error()) {
            TaskState::Failed
        } else if finished {
            TaskState::Completed
        } else if session.interrupted {
            TaskState::Interrupted
        } else {
            // Starting up.
            TaskState::Running
        }
    }

    fn key(self) -> &'static str {
        match self {
            TaskState::Running => "running",
            TaskState::WaitingForUser => "waiting_for_user",
            TaskState::Completed => "completed",
            TaskState::Failed => "failed",
            TaskState::Cancelled => "cancelled",
            TaskState::Interrupted => "interrupted",
        }
    }

    fn is_terminal(self) -> bool {
        !matches!(self, TaskState::Running | TaskState::WaitingForUser)
    }

    fn panel_status(self) -> TaskStatus {
        match self {
            TaskState::Running | TaskState::WaitingForUser => TaskStatus::Running,
            TaskState::Completed => TaskStatus::Completed,
            TaskState::Failed => TaskStatus::Failed,
            TaskState::Cancelled | TaskState::Interrupted => TaskStatus::Stopped,
        }
    }
}

/// A task as `delegate_task` and `task_status` report it.
fn report(session: &Session, wait_timed_out: bool) -> Value {
    let state = TaskState::of(session);
    let summary = state.is_terminal().then(|| session.last_answer()).flatten().map(|answer| truncate(&answer));
    json!({
        "taskId": session.id,
        "childThreadId": session.id,
        "title": session.title.to_string(),
        "status": state.key(),
        "providerInstanceId": session.instance,
        "model": session.active_model.clone().or_else(|| session.model.clone()),
        "summary": summary,
        "waitTimedOut": wait_timed_out,
    })
}

fn truncate(text: &str) -> String {
    if text.len() <= RESULT_LIMIT {
        return text.to_string();
    }
    let mut end = RESULT_LIMIT;
    while !text.is_char_boundary(end) {
        end -= 1;
    }
    format!("{}\n\n[Cut short; task_status has the same text. The subagent's thread has the rest.]", &text[..end])
}

/// The last line of what an agent said.
fn last_line(text: &str) -> Option<String> {
    let line = text.lines().rev().map(str::trim).find(|line| !line.is_empty())?;
    Some(line.chars().take(200).collect())
}

impl Workspace {
    /// Starts answering the calls agents make to the app's tools.
    pub(super) fn listen_for_tools(&mut self, window: &mut Window, cx: &mut Context<Self>) {
        let requests = orchestrator::requests();
        self.delegations._requests = Some(cx.spawn_in(window, async move |this, cx| {
            while let Ok(request) = requests.recv().await {
                let ToolRequest { thread, tool, args, reply } = request;
                let answered = this.update_in(cx, |this, window, cx| {
                    this.answer_tool(&thread, &tool, args, Answer::Mcp(reply), window, cx)
                });
                if answered.is_err() {
                    break;
                }
            }
        }));
    }

    /// jcode asks for the app's tools through its thread.
    pub(super) fn on_session_event(
        &mut self,
        session: &Entity<Session>,
        event: &SessionEvent,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        if let SessionEvent::AppTool { call_id, tool, input } = event {
            let thread = session.read(cx).id.clone();
            let answer = Answer::Jcode { session: session.downgrade(), call_id: call_id.clone() };
            self.answer_tool(&thread, tool, input.clone(), answer, window, cx);
        }
    }

    /// A thread on this device, by id.
    fn thread_by_id(&self, id: &str, cx: &App) -> Option<Entity<Session>> {
        self.local_thread(id, cx).map(|(_, session)| session)
    }

    fn answer_tool(
        &mut self,
        thread: &str,
        tool: &str,
        args: Value,
        answer: Answer,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        let Some(parent) = self.thread_by_id(thread, cx) else {
            answer.send(Err("The thread calling this is gone from Vectrum Code.".into()), cx);
            return;
        };
        if let Some(tool) = orchestrator::device_tool(tool) {
            self.answer_device_tool(thread, tool, args, answer, cx);
            return;
        }
        if parent.read(cx).is_subagent() {
            answer.send(Err("A subagent can't delegate tasks of its own; do the work yourself.".into()), cx);
            return;
        }
        let outcome = match orchestrator::own_tool(tool) {
            Some(CAPABILITIES) => Ok(capabilities(parent.read(cx), cx)),
            Some(DELEGATE) => match self.delegate(&parent, &args, window, cx) {
                Ok(child) if args.get("mode").and_then(Value::as_str) == Some("wait") => {
                    self.wait_for(&parent, child, &args, answer, cx);
                    return;
                }
                Ok(child) => Ok(report(child.read(cx), false)),
                Err(err) => Err(err),
            },
            Some(TASK_STATUS) => self.task_of(&parent, &args, cx).map(|child| {
                // Reading a finished task's result counts as hearing of it.
                if TaskState::of(child.read(cx)).is_terminal() {
                    self.mark_delivered(&child, cx);
                }
                report(child.read(cx), false)
            }),
            Some(TASK_CANCEL) => self.task_of(&parent, &args, cx).map(|child| {
                if !TaskState::of(child.read(cx)).is_terminal() {
                    // The parent knows: no message about it.
                    child.update(cx, |child, cx| {
                        if let Some(delegation) = child.delegation.as_mut() {
                            delegation.cancelled = true;
                            delegation.delivered = true;
                        }
                        child.interrupt(cx);
                    });
                    self.dirty.insert(child.entity_id());
                }
                let mut report = report(child.read(cx), false);
                if !TaskState::of(child.read(cx)).is_terminal() {
                    report["status"] = "cancel_requested".into();
                }
                report
            }),
            _ => Err(format!("Vectrum Code has no tool named {tool}.")),
        };
        answer.send(outcome, cx);
    }

    /// Starts a subagent thread on a task `parent` delegated.
    fn delegate(
        &mut self,
        parent: &Entity<Session>,
        args: &Value,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) -> Result<Entity<Session>, String> {
        fn text(value: Option<&Value>) -> Option<&str> {
            value.and_then(Value::as_str).map(str::trim).filter(|text| !text.is_empty())
        }
        let task = text(args.get("task")).ok_or("delegate_task needs a task.")?.to_string();
        let title = text(args.get("title")).map(str::to_string);
        let target = args.get("target");
        let state = parent.read(cx);
        let instance_id = text(target.and_then(|target| target.get("providerInstanceId"))).unwrap_or(&state.instance);
        let instance = settings::instances()
            .into_iter()
            .find(|instance| instance.id == instance_id)
            .ok_or_else(|| format!("There's no provider instance \"{instance_id}\"; orchestrator_capabilities lists them."))?;
        if !instance.enabled {
            return Err(format!("The provider instance \"{instance_id}\" is turned off in Vectrum Code's settings."));
        }
        // What the parent runs carries over to a child on the same agent.
        let same = instance.id == state.instance;
        let model = text(target.and_then(|target| target.get("model")))
            .map(str::to_string)
            .or_else(|| state.model.clone().filter(|_| same));
        let new = NewSession {
            instance: instance.id.clone(),
            provider: instance.kind,
            project: state.project.clone(),
            model: model.clone(),
            permission_mode: if instance.kind == state.provider { state.permission_mode } else { Default::default() },
            reasoning: if same { state.reasoning } else { Reasoning::Default },
            fast_mode: same && state.fast_mode,
            worktree: None,
        };
        let origin = SubagentParent { id: state.id.clone(), cwd: state.cwd.clone(), branch: state.branch.clone() };
        // A different agent or model is worth naming in the panel.
        let kind = (!same || model != state.model).then(|| {
            let agent = if instance.name.trim().is_empty() { instance.kind.label().to_string() } else { instance.name.clone() };
            match &model {
                Some(model) => format!("{agent} · {model}"),
                None => agent,
            }
        });
        let project_ix = self.project_of(parent).ok_or("The thread's project is gone.")?;
        let label = title.clone().unwrap_or_else(|| {
            task.lines().map(str::trim).find(|line| !line.is_empty()).unwrap_or_default().chars().take(200).collect()
        });
        let child = cx.new(|cx| Session::new_subagent(new, &origin, title, task, cx));
        self.attach(project_ix, child.clone(), window, cx);
        let member = SubagentMember {
            id: child.read(cx).id.clone(),
            label,
            kind,
            status: "running".into(),
            started: to_millis(SystemTime::now()),
            ..Default::default()
        };
        parent.update(cx, |parent, cx| parent.link_delegation(args, member, cx));
        self.dirty.insert(child.entity_id());
        cx.notify();
        Ok(child)
    }

    /// Holds a `mode: "wait"` call until its task finishes, or until it's
    /// waited as long as it may.
    fn wait_for(&mut self, parent: &Entity<Session>, child: Entity<Session>, args: &Value, answer: Answer, cx: &mut Context<Self>) {
        let asked = args.get("timeoutMs").and_then(Value::as_f64).filter(|ms| *ms > 0.).map(|ms| Duration::from_millis(ms as u64));
        let timeout = asked.unwrap_or(DEFAULT_WAIT).min(orchestrator::wait_limit(parent.read(cx).provider));
        self.delegations.next_waiter += 1;
        let id = self.delegations.next_waiter;
        let task_id = child.read(cx).id.clone();
        let timer = cx.spawn(async move |this, cx| {
            cx.background_executor().timer(timeout).await;
            let _ = this.update(cx, |this, cx| {
                let Some(ix) = this.delegations.waiting.iter().position(|waiter| waiter.id == id) else { return };
                let waiter = this.delegations.waiting.remove(ix);
                let report = report(child.read(cx), true);
                waiter.answer.send(Ok(report), cx);
            });
        });
        let task = task_id;
        self.delegations.waiting.push(Waiter { id, task, answer, _timeout: timer });
    }

    /// A task `parent` delegated, by the `taskId` in a call.
    fn task_of(&self, parent: &Entity<Session>, args: &Value, cx: &App) -> Result<Entity<Session>, String> {
        let task = args.get("taskId").and_then(Value::as_str).ok_or("Name the task with taskId.")?;
        let parent_id = &parent.read(cx).id;
        self.thread_by_id(task, cx)
            .filter(|child| child.read(cx).delegation.as_ref().is_some_and(|delegation| &delegation.parent == parent_id))
            .ok_or_else(|| format!("This thread delegated no task {task}."))
    }

    fn mark_delivered(&mut self, child: &Entity<Session>, cx: &mut Context<Self>) {
        child.update(cx, |child, _| {
            if let Some(delegation) = child.delegation.as_mut() {
                delegation.delivered = true;
            }
        });
        self.dirty.insert(child.entity_id());
    }

    /// Keeps a subagent's row in its parent current, and once its task is
    /// done, hands the result over.
    pub(super) fn track_delegation(&mut self, child: &Entity<Session>, cx: &mut Context<Self>) {
        let state = child.read(cx);
        // A mirrored subagent's host keeps its parent current.
        if state.remote.is_some() {
            return;
        }
        let Some(delegation) = state.delegation.clone() else { return };
        let task_state = TaskState::of(state);
        let task = state.id.clone();
        let tools: Vec<_> = state
            .items
            .iter()
            .filter_map(|item| match item {
                ChatItem::Tool(tool) => Some(tool),
                _ => None,
            })
            .collect();
        let tool_uses = tools.len();
        let latest_tool = tools.last().map(|tool| (tool.name.clone(), brief_input(&tool.input)));
        let summary = state.last_answer().as_deref().and_then(last_line);
        let tokens = state.context_tokens;
        let Some(parent) = self.thread_by_id(&delegation.parent, cx) else { return };
        parent.update(cx, |parent, cx| {
            parent.update_delegation(
                &task,
                |member| {
                    let status = task_state.panel_status();
                    if member.status != status.key() {
                        member.status = status.key().to_string();
                        member.finished = (status != TaskStatus::Running).then(|| to_millis(SystemTime::now()));
                    }
                    member.tool_uses = tool_uses;
                    member.tool = latest_tool;
                    member.summary = summary.or(member.summary.take());
                    member.tokens = tokens.or(member.tokens);
                },
                cx,
            )
        });
        if !task_state.is_terminal() || delegation.delivered {
            return;
        }
        self.mark_delivered(child, cx);
        let (waiting, rest): (Vec<Waiter>, Vec<Waiter>) =
            std::mem::take(&mut self.delegations.waiting).into_iter().partition(|waiter| waiter.task == task);
        self.delegations.waiting = rest;
        let result = report(child.read(cx), false);
        if !waiting.is_empty() {
            for waiter in waiting {
                waiter.answer.send(Ok(result.clone()), cx);
            }
            return;
        }
        // Nobody's waiting: tell the parent's agent, which wakes it.
        let title = child.read(cx).title.to_string();
        let mut message = format!(
            "{} {task} (\"{title}\") {}.",
            orchestrator::RESULT_PREFIX,
            match task_state {
                TaskState::Completed => "completed",
                TaskState::Failed => "failed",
                TaskState::Cancelled => "was stopped by the user",
                _ => "stopped",
            }
        );
        match result.get("summary").and_then(Value::as_str) {
            Some(summary) => message.push_str(&format!(" Its final answer:\n\n{summary}")),
            None => message.push_str(" It gave no answer."),
        }
        parent.update(cx, |parent, cx| parent.send_hidden(message, cx));
    }

    /// Stops a delegated task from the subagents panel. Its parent's agent
    /// hears that the user stopped it.
    pub(super) fn stop_delegated(&mut self, task: &str, cx: &mut Context<Self>) {
        let Some(child) = self.thread_by_id(task, cx) else { return };
        child.update(cx, |child, cx| {
            if let Some(delegation) = child.delegation.as_mut() {
                delegation.cancelled = true;
            }
            child.interrupt(cx);
            cx.notify();
        });
        self.dirty.insert(child.entity_id());
    }
}

/// What `orchestrator_capabilities` reports: every agent that's on, with
/// its models, and what the calling thread runs.
fn capabilities(thread: &Session, cx: &App) -> Value {
    let instances: Vec<Value> = Catalog::global(cx)
        .instances
        .iter()
        .filter(|state| state.config.enabled && !matches!(state.status, ProbeStatus::Disabled))
        .map(|state| {
            let models: Vec<Value> = state
                .models
                .iter()
                .map(|model| json!({"id": model.id, "name": model.name, "detail": model.detail}))
                .collect();
            json!({
                "providerInstanceId": state.config.id,
                "kind": state.config.kind.key(),
                "name": if state.config.name.trim().is_empty() { state.config.kind.label().to_string() } else { state.config.name.clone() },
                "models": models,
            })
        })
        .collect();
    json!({
        "current": {"providerInstanceId": thread.instance, "model": thread.model, "activeModel": thread.active_model},
        "providerInstances": instances,
        "tools": [DELEGATE, TASK_STATUS, TASK_CANCEL, CAPABILITIES],
    })
}

/// The thread a `delegate_task` call started, from its result: the reply's
/// `childThreadId`, however the agent quoted it.
fn delegated_thread(result: &str) -> Option<String> {
    let (_, rest) = result.split_once("childThreadId")?;
    let id: String = rest
        .trim_start_matches(|c: char| matches!(c, '"' | '\\' | ':') || c.is_whitespace())
        .chars()
        .take_while(|c| c.is_ascii_alphanumeric() || *c == '-')
        .collect();
    (!id.is_empty()).then_some(id)
}

impl Workspace {
    /// Marks subagents' threads again when a build that didn't know them
    /// saved them as ordinary threads, from the `delegate_task` calls that
    /// started them, and gives their parents back their rows.
    pub(super) fn recover_delegations(&mut self, cx: &mut Context<Self>) {
        let sessions: Vec<Entity<Session>> = self
            .projects
            .iter()
            .filter(|project| project.environment.is_none())
            .flat_map(|project| project.sessions.clone())
            .collect();
        let mut started = Vec::new();
        for parent in &sessions {
            for (ix, item) in parent.read(cx).items.iter().enumerate() {
                if let ChatItem::Tool(tool) = item
                    && tool.is_delegation()
                    && let Some(child) = tool.result.as_deref().and_then(delegated_thread)
                {
                    started.push((parent.clone(), ix, child));
                }
            }
        }
        for (parent, ix, child_id) in started {
            let Some(child) = sessions.iter().find(|session| session.read(cx).id == child_id) else { continue };
            if child.read(cx).delegation.is_some() || child == &parent {
                continue;
            }
            let parent_id = parent.read(cx).id.clone();
            child.update(cx, |child, _| {
                child.delegation = Some(Delegation { parent: parent_id, delivered: true, cancelled: false });
            });
            let state = child.read(cx);
            let status = TaskState::of(state).panel_status();
            let member = SubagentMember {
                id: child_id.clone(),
                label: state.title.to_string(),
                status: status.key().to_string(),
                started: to_millis(state.created),
                finished: (status != TaskStatus::Running).then(|| to_millis(state.updated)),
                tokens: state.context_tokens,
                summary: state.last_answer().as_deref().and_then(last_line),
                ..Default::default()
            };
            parent.update(cx, |parent, _| {
                if let Some(ChatItem::Tool(tool)) = parent.items.get_mut(ix) {
                    let progress = tool.subagent.get_or_insert_with(Default::default);
                    if !progress.members.iter().any(|member| member.id == child_id) {
                        progress.members.push(member);
                    }
                }
            });
            self.dirty.insert(child.entity_id());
            self.dirty.insert(parent.entity_id());
        }
    }
}

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

    #[test]
    fn finds_the_child_thread_however_the_result_is_quoted() {
        let id = Some("1a127-118828-4".to_string());
        assert_eq!(delegated_thread("{\n  \"taskId\": \"1a127-118828-4\",\n  \"childThreadId\": \"1a127-118828-4\"\n}"), id);
        assert_eq!(delegated_thread(r#"[{"type":"text","text":"{\"childThreadId\":\"1a127-118828-4\"}"}]"#), id);
        assert_eq!(delegated_thread("The task failed to start."), None);
    }
}