//! The subagents panel: every subagent the open thread's agent started, the
//! working ones first with what each is doing right now (its latest tool
//! call, its progress summary, how long it's run, its tool uses and
//! tokens), then the finished ones. A click shows the subagent's call in
//! its conversation.
//!
//! Claude's Task and Agent calls are one subagent each. Pi's `subagent`
//! extension and jcode's swarm spawns can start several, each its own row.
//! A task delegated to the app runs in a thread of its own, which a click
//! opens.

use std::time::Duration;

use gpui_kit::component::{Icon, IconName, h_flex, v_flex};
use gpui_kit::prelude::FluentBuilder as _;
use gpui_kit::*;

use super::{PANEL_WIDTH, Workspace};
use crate::session::{ChatItem, Session, SubagentStop, TaskStatus};
use crate::ui::chat::{describe_call, format_duration};
use crate::ui::icons::Lucide;
use crate::ui::metal::{self, Led, MetalButton, MetalStyled as _, Palette};

/// A subagent as the panel lists it.
struct Subagent {
    session: Entity<Session>,
    /// Its tool call's place in the thread's transcript.
    ix: usize,
    /// Tells apart the rows, several of which can share a call.
    id: usize,
    description: SharedString,
    /// The kind of agent, e.g. `Explore`.
    kind: Option<SharedString>,
    status: TaskStatus,
    background: bool,
    tool_uses: usize,
    tokens: Option<u64>,
    elapsed: Option<Duration>,
    /// What it's doing: its latest tool call while it works.
    activity: Option<(Icon, String)>,
    summary: Option<SharedString>,
    /// The thread a task delegated to the app runs in, by id.
    thread: Option<String>,
    /// How its Stop button stops it, while it runs.
    stop: Option<SubagentStop>,
}

/// A thread's subagents, in the order they started.
fn subagents_of(session: &Entity<Session>, cx: &App) -> Vec<Subagent> {
    let state = session.read(cx);
    let thread = session.entity_id().as_u64() as usize;
    let mut subagents = Vec::new();
    for (ix, item) in state.items.iter().enumerate() {
        let ChatItem::Tool(tool) = item else { continue };
        for (member, entry) in tool.subagents().into_iter().enumerate() {
            subagents.push(Subagent {
                session: session.clone(),
                ix,
                id: thread << 32 | ix << 8 | member,
                description: entry.label.into(),
                // The default kind goes without saying.
                kind: entry.kind.filter(|kind| !kind.is_empty() && kind != "general-purpose").map(Into::into),
                status: entry.status,
                background: entry.background,
                tool_uses: entry.tool_uses,
                tokens: entry.tokens,
                elapsed: entry.elapsed,
                activity: entry.tool.map(|(name, input)| describe_call(&name, &input, &state.cwd)),
                summary: entry.summary.map(Into::into),
                thread: entry.thread,
                stop: entry.stop,
            });
        }
    }
    subagents
}

/// `812 tokens`, `23.4k tokens`, `1.2M tokens`.
fn format_tokens(tokens: u64) -> String {
    match tokens {
        0..1_000 => format!("{tokens} tokens"),
        1_000..1_000_000 => format!("{:.1}k tokens", tokens as f64 / 1_000.).replace(".0k", "k"),
        _ => format!("{:.1}M tokens", tokens as f64 / 1_000_000.).replace(".0M", "M"),
    }
}

impl Workspace {
    /// How many subagents the open thread has working, for the header's light.
    pub(super) fn working_subagents(&self, cx: &App) -> usize {
        self.active.as_ref().map_or(0, |session| {
            subagents_of(session, cx).iter().filter(|subagent| subagent.status == TaskStatus::Running).count()
        })
    }

    /// Opens a subagent's thread, if it isn't open, and shows its call; or
    /// for a task delegated to the app, the thread it runs in.
    fn reveal_subagent(
        &mut self,
        session: Entity<Session>,
        ix: usize,
        thread: Option<String>,
        window: &mut Window,
        cx: &mut Context<Self>,
    ) {
        if let Some((project, child)) = thread.and_then(|thread| self.local_thread(&thread, cx)) {
            self.select(project, child, window, cx);
            return;
        }
        if self.active.as_ref() != Some(&session) {
            let Some(project) = self.project_of(&session) else { return };
            self.select(project, session.clone(), window, cx);
        }
        if let Some(chat) = self.chats.get(&session.entity_id()).cloned() {
            chat.update(cx, |chat, cx| chat.reveal_tool(ix, cx));
        }
    }

    fn stop_subagent(&mut self, session: &Entity<Session>, stop: &SubagentStop, cx: &mut Context<Self>) {
        match stop {
            SubagentStop::Thread(task) => self.stop_delegated(task, cx),
            _ => session.update(cx, |session, cx| session.stop_subagent(stop, cx)),
        }
        cx.notify();
    }

    pub(super) fn render_subagents_panel(&self, cx: &mut Context<Self>) -> impl IntoElement {
        let p = metal::palette(cx);
        let mine = self.active.as_ref().map(|session| subagents_of(session, cx)).unwrap_or_default();
        let (working, mut finished): (Vec<Subagent>, Vec<Subagent>) =
            mine.into_iter().partition(|subagent| subagent.status == TaskStatus::Running);
        finished.reverse();
        let empty = working.is_empty() && finished.is_empty();

        let header = h_flex()
            .h(px(36.))
            .pl(px(10.))
            .gap(px(10.))
            .child(Icon::new(Lucide::Bot).size(px(16.)).text_color(p.icon))
            .child(div().flex_1().text_size(px(13.5)).font_weight(FontWeight::MEDIUM).child("Subagents"))
            .when(!working.is_empty(), |this| this.child(metal::tag(format!("{} working", working.len()), p)))
            .child(
                MetalButton::new("close-subagents")
                    .ghost()
                    .mini()
                    .icon(IconName::Close)
                    .tooltip("Close subagents")
                    .on_click(cx.listener(|this, _, _, cx| {
                        this.right_panel = None;
                        cx.notify();
                    })),
            );

        let section = |title: String, rows: Vec<AnyElement>| {
            v_flex()
                .gap(px(2.))
                .child(div().px(px(10.)).pt(px(8.)).pb(px(4.)).child(metal::engraved(title, p)))
                .children(rows)
        };
        let rows = |subagents: Vec<Subagent>, cx: &mut Context<Self>| -> Vec<AnyElement> {
            subagents.into_iter().map(|subagent| self.render_subagent(subagent, p, cx)).collect()
        };
        let working_count = working.len();
        let finished_count = finished.len();
        let working_rows = rows(working, cx);
        let finished_rows = rows(finished, cx);

        v_flex()
            .w(px(PANEL_WIDTH))
            .flex_none()
            .h_full()
            .p(px(12.))
            .pl(px(0.))
            .child(
                v_flex()
                    .max_h_full()
                    .p(px(6.))
                    .rounded(px(18.))
                    .raised(p)
                    .child(header)
                    .child(div().px(px(8.)).py(px(4.)).child(metal::rule(p)))
                    .child(
                        v_flex()
                            .id("subagents-list")
                            .flex_shrink(1.)
                            .min_h_0()
                            .overflow_y_scroll()
                            .when(empty, |this| {
                                this.child(
                                    v_flex()
                                        .px(px(10.))
                                        .py(px(10.))
                                        .gap(px(4.))
                                        .child(div().text_size(px(13.)).text_color(p.ink2).child("No subagents in this thread"))
                                        .child(div().text_size(px(12.)).text_color(p.ink3).child(
                                            "When the agent hands work to a subagent, it shows up here with what it's doing.",
                                        )),
                                )
                            })
                            .when(working_count > 0, |this| this.child(section("Working".into(), working_rows)))
                            .when(finished_count > 0, |this| this.child(section(format!("Finished · {finished_count}"), finished_rows))),
                    ),
            )
    }

    /// One subagent: its light, what it was asked to do and for how long,
    /// its kind and totals, and while it works, what it's doing now.
    fn render_subagent(&self, subagent: Subagent, p: &Palette, cx: &mut Context<Self>) -> AnyElement {
        let running = subagent.status == TaskStatus::Running;
        let (led, status) = match subagent.status {
            TaskStatus::Running => (Led::Live, None),
            TaskStatus::Completed => (Led::Off, None),
            TaskStatus::Failed => (Led::Failed, Some("Failed")),
            TaskStatus::Stopped => (Led::Off, Some("Stopped")),
        };
        let id = subagent.id;
        // A delegated task's agent can stop to ask the user something.
        let waiting = running
            && subagent
                .thread
                .as_ref()
                .and_then(|thread| self.local_thread(thread, cx))
                .is_some_and(|(_, child)| child.read(cx).needs_attention());
        let status = if waiting { Some("Waiting for you") } else { status };
        let meta: Vec<String> = [
            subagent.kind.as_ref().map(ToString::to_string),
            (subagent.tool_uses > 0)
                .then(|| format!("{} {}", subagent.tool_uses, if subagent.tool_uses == 1 { "tool" } else { "tools" })),
            subagent.tokens.filter(|tokens| *tokens > 0).map(format_tokens),
            subagent.background.then(|| "background".to_string()),
        ]
        .into_iter()
        .flatten()
        .collect();
        let (session, ix, child) = (subagent.session.clone(), subagent.ix, subagent.thread.clone());
        let (stop, stop_session) = (subagent.stop.clone(), subagent.session.clone());
        h_flex()
            .id(("subagent", id))
            .items_start()
            .gap(px(10.))
            .px(px(10.))
            .py(px(7.))
            .rounded(px(12.))
            .cursor_pointer()
            .hover(|style| style.bg(p.row_hover))
            .on_click(cx.listener(move |this, _, window, cx| this.reveal_subagent(session.clone(), ix, child.clone(), window, cx)))
            .child(div().flex_none().h(px(18.)).flex().items_center().child(metal::led(("subagent-led", id), led, running, p)))
            .child(
                v_flex()
                    .flex_1()
                    .min_w_0()
                    .gap(px(2.))
                    .child(
                        h_flex()
                            .gap(px(8.))
                            .child(
                                div()
                                    .flex_1()
                                    .min_w_0()
                                    .truncate()
                                    .text_size(px(13.))
                                    .text_color(if running { p.ink } else { p.ink2 })
                                    .child(subagent.description.clone()),
                            )
                            .when_some(status, |this, status| {
                                this.child(
                                    div()
                                        .flex_none()
                                        .text_size(px(11.5))
                                        .text_color(if subagent.status == TaskStatus::Failed { p.red } else { p.ink3 })
                                        .child(status),
                                )
                            })
                            .when_some(subagent.elapsed, |this, elapsed| {
                                this.child(
                                    div()
                                        .flex_none()
                                        .font_family(metal::MONO)
                                        .text_size(px(11.))
                                        .text_color(p.ink3)
                                        .child(format_duration(elapsed.as_millis() as u64)),
                                )
                            })
                            .when_some(stop, |this, stop| {
                                let tip = if stop == SubagentStop::Turn {
                                    "Stop (stops the whole turn: Pi runs these inside one tool call)"
                                } else {
                                    "Stop this subagent"
                                };
                                let session = stop_session.clone();
                                this.child(
                                    MetalButton::new(("stop-subagent", id))
                                        .ghost()
                                        .mini()
                                        .icon(Lucide::Square)
                                        .tooltip(tip)
                                        .on_click(cx.listener(move |this, _, _, cx| {
                                            cx.stop_propagation();
                                            this.stop_subagent(&session, &stop, cx);
                                        })),
                                )
                            }),
                    )
                    .when_some(subagent.activity, |this, (icon, label)| {
                        this.child(
                            h_flex()
                                .gap(px(6.))
                                .text_size(px(12.))
                                .text_color(p.ink2)
                                .child(icon.size(px(13.)).text_color(p.ink3))
                                .child(div().flex_1().min_w_0().truncate().child(label)),
                        )
                    })
                    .when_some(subagent.summary.filter(|_| running), |this, summary| {
                        this.child(div().text_size(px(12.)).text_color(p.ink3).line_clamp(2).child(summary))
                    })
                    .when(!meta.is_empty(), |this| {
                        this.child(
                            div()
                                .font_family(metal::MONO)
                                .text_size(px(10.5))
                                .text_color(p.ink3)
                                .child(meta.join(" · ")),
                        )
                    }),
            )
            .into_any_element()
    }
}

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

    #[::core::prelude::v1::test]
    fn token_counts() {
        assert_eq!(format_tokens(812), "812 tokens");
        assert_eq!(format_tokens(23_412), "23.4k tokens");
        assert_eq!(format_tokens(40_000), "40k tokens");
        assert_eq!(format_tokens(1_340_000), "1.3M tokens");
    }
}