//! Agent pool — manages the set of active agents across all pipeline stages.
pub(crate) mod auto_assign;
/// TransitionFired subscriber that snapshots token cost on terminal stage transitions.
pub(crate) mod cost_rollup_subscriber;
mod pipeline;
mod process;
mod query;
mod start;
mod stop;
mod types;
mod wait;
mod worktree;
/// Lifecycle subscribers that create and remove worktrees on pipeline transitions.
pub(crate) mod worktree_lifecycle;

#[cfg(test)]
mod test_helpers;

use crate::io::watcher::WatcherEvent;
use crate::service::status::StatusBroadcaster;
use std::collections::HashMap;
use std::sync::Arc;
use tokio::sync::{Mutex, broadcast};

// Bring pool-internal types into pool's namespace so that sub-modules
// (auto_assign, pipeline, etc.) can access them via `use super::...`.
use types::{StoryAgent, composite_key};

/// Manages concurrent story agents, each in its own worktree.
pub struct AgentPool {
    agents: Arc<Mutex<HashMap<String, StoryAgent>>>,
    port: u16,
    /// Broadcast channel for notifying WebSocket clients of agent state changes.
    /// When an agent transitions state (Pending, Running, Completed, Failed, Stopped),
    /// an `AgentStateChanged` event is emitted so the frontend can refresh the
    /// pipeline board without waiting for a filesystem event.
    watcher_tx: broadcast::Sender<WatcherEvent>,
    /// Project-scoped status broadcaster.  Each agent session creates a
    /// [`crate::service::status::buffer::StatusEventBuffer`] subscribed to this
    /// broadcaster so it passively accumulates pipeline events without side effects.
    ///
    /// Call sites that also build a [`crate::services::Services`] bundle should
    /// obtain the broadcaster via [`status_broadcaster`](AgentPool::status_broadcaster)
    /// and use it as the `status` field so all consumers share a single channel.
    status_broadcaster: Arc<StatusBroadcaster>,
}

impl AgentPool {
    /// Create a new agent pool bound to the given HTTP port and event channel.
    pub fn new(port: u16, watcher_tx: broadcast::Sender<WatcherEvent>) -> Self {
        let pool = Self {
            agents: Arc::new(Mutex::new(HashMap::new())),
            port,
            watcher_tx: watcher_tx.clone(),
            status_broadcaster: Arc::new(StatusBroadcaster::new()),
        };

        // Spawn a background task (only when inside a tokio runtime) that
        // listens for RateLimitWarning and RateLimitHardBlock events and updates the
        // throttled field on the relevant agent so status dots stay current.
        if tokio::runtime::Handle::try_current().is_ok() {
            let agents_clone = Arc::clone(&pool.agents);
            let watcher_tx_clone = watcher_tx.clone();
            let mut rx = watcher_tx.subscribe();
            tokio::spawn(async move {
                loop {
                    let event = match rx.recv().await {
                        Ok(e) => e,
                        Err(broadcast::error::RecvError::Closed) => break,
                        Err(broadcast::error::RecvError::Lagged(_)) => continue,
                    };
                    let (story_id, agent_name, until) = match &event {
                        WatcherEvent::RateLimitWarning {
                            story_id,
                            agent_name,
                        } => (
                            story_id.clone(),
                            agent_name.clone(),
                            // No explicit reset time — use a 15-minute soft timeout.
                            chrono::Utc::now() + chrono::Duration::minutes(15),
                        ),
                        WatcherEvent::RateLimitHardBlock {
                            story_id,
                            agent_name,
                            reset_at,
                        } => (story_id.clone(), agent_name.clone(), *reset_at),
                        _ => continue,
                    };
                    let key = composite_key(&story_id, &agent_name);
                    {
                        let mut agents = agents_clone.lock().await;
                        if let Some(agent) = agents.get_mut(&key) {
                            agent.throttled =
                                Some(crate::agents::AgentExecution::Throttled { until });
                        }
                    }
                    let _ = watcher_tx_clone.send(WatcherEvent::AgentStateChanged);
                }
            });
        }

        pool
    }

    /// Return the HTTP port this pool's agents connect to.
    pub fn port(&self) -> u16 {
        self.port
    }

    /// Return a clone of the broadcaster used for filesystem-watcher events.
    ///
    /// Used by sibling modules (e.g. the merge runner) that need to fire
    /// re-dispatch events without holding a long-lived reference.
    pub(crate) fn watcher_tx(&self) -> broadcast::Sender<WatcherEvent> {
        self.watcher_tx.clone()
    }

    /// Returns the project-scoped [`StatusBroadcaster`] owned by this pool.
    ///
    /// Callers that also construct a [`crate::services::Services`] bundle should
    /// use this as the `status` field so that pipeline events published to the
    /// service are delivered to each session's buffer.
    pub fn status_broadcaster(&self) -> Arc<StatusBroadcaster> {
        Arc::clone(&self.status_broadcaster)
    }

    /// Create a pool with a dummy watcher channel for unit tests.
    #[cfg(test)]
    pub fn new_test(port: u16) -> Self {
        let (watcher_tx, _) = broadcast::channel(16);
        Self::new(port, watcher_tx)
    }

    /// Notify WebSocket clients that agent state has changed, so the pipeline
    /// board and agent panel can refresh.
    fn notify_agent_state_changed(watcher_tx: &broadcast::Sender<WatcherEvent>) {
        let _ = watcher_tx.send(WatcherEvent::AgentStateChanged);
    }
}

#[cfg(test)]
mod tests {
    use super::AgentPool;
    use crate::agents::AgentStatus;
    use crate::service::status::StatusEvent;
    use crate::service::status::buffer::BufferedItem;
    use tokio::time::{Duration, sleep};

    fn stage_transition_event(story_id: &str) -> StatusEvent {
        StatusEvent::StageTransition {
            story_id: story_id.to_string(),
            story_name: String::new(),
            from_stage: "1_backlog".to_string(),
            to_stage: "2_current".to_string(),
        }
    }

    /// Criterion 3: a session for project A only buffers events from project A;
    /// a session for project B sees nothing from A.
    ///
    /// Each `AgentPool` owns its own `StatusBroadcaster`.  Publishing to pool A's
    /// broadcaster must not reach pool B's agents, and vice versa.
    #[tokio::test]
    async fn agent_buffer_is_scoped_to_its_pool_broadcaster() {
        let pool_a = AgentPool::new_test(3201);
        let pool_b = AgentPool::new_test(3202);

        // Each agent subscribes to its own pool's broadcaster at inject time.
        pool_a.inject_test_agent_with_live_buffer("story-a", "coder", AgentStatus::Running);
        pool_b.inject_test_agent_with_live_buffer("story-b", "coder", AgentStatus::Running);

        // Publish an event only to project A.
        pool_a
            .status_broadcaster()
            .publish(stage_transition_event("story-a"));

        // Yield so background receive tasks can process the event.
        sleep(Duration::from_millis(20)).await;

        // Pool A's agent buffer must contain the event.
        let items_a = pool_a
            .drain_agent_status_buffer("story-a", "coder")
            .expect("pool A agent must have a buffer");
        assert_eq!(
            items_a.len(),
            1,
            "pool A agent should have buffered 1 event"
        );
        assert!(
            matches!(
                items_a[0],
                BufferedItem::Event(StatusEvent::StageTransition { .. })
            ),
            "pool A agent should see the StageTransition event"
        );

        // Pool B's agent buffer must be empty — it uses a different broadcaster.
        let items_b = pool_b
            .drain_agent_status_buffer("story-b", "coder")
            .expect("pool B agent must have a buffer");
        assert!(
            items_b.is_empty(),
            "pool B agent must not see events from pool A (got {} items)",
            items_b.len()
        );
    }

    /// Criterion 3 (reverse direction): events published to project B must not
    /// reach project A's agents.
    #[tokio::test]
    async fn agent_buffer_project_b_events_do_not_reach_project_a() {
        let pool_a = AgentPool::new_test(3203);
        let pool_b = AgentPool::new_test(3204);

        pool_a.inject_test_agent_with_live_buffer("story-x", "coder", AgentStatus::Running);
        pool_b.inject_test_agent_with_live_buffer("story-y", "coder", AgentStatus::Running);

        // Publish only to project B.
        pool_b
            .status_broadcaster()
            .publish(stage_transition_event("story-y"));

        sleep(Duration::from_millis(20)).await;

        // Pool B agent sees the event.
        let items_b = pool_b
            .drain_agent_status_buffer("story-y", "coder")
            .expect("pool B agent must have a buffer");
        assert_eq!(items_b.len(), 1, "pool B agent should see 1 event");

        // Pool A agent sees nothing.
        let items_a = pool_a
            .drain_agent_status_buffer("story-x", "coder")
            .expect("pool A agent must have a buffer");
        assert!(
            items_a.is_empty(),
            "pool A agent must not see project B events"
        );
    }

    /// Criterion 4: events arriving on the buffer do not invoke the agent's model.
    ///
    /// The buffer's background task calls only `push()` on the inner deque.
    /// We verify this at the pool level: publish N events and confirm that
    /// (a) they all accumulate silently, and (b) no external state is mutated
    /// beyond the buffer itself.
    #[tokio::test]
    async fn buffer_events_accumulate_without_side_effects() {
        use std::sync::Arc;
        use std::sync::atomic::{AtomicUsize, Ordering};

        let pool = AgentPool::new_test(3205);
        pool.inject_test_agent_with_live_buffer("story-z", "coder", AgentStatus::Running);

        // A counter that represents "model call count" — should stay 0.
        let model_call_count = Arc::new(AtomicUsize::new(0));

        // Publish 5 events to the pool's broadcaster.
        let bc = pool.status_broadcaster();
        for i in 0..5u32 {
            bc.publish(StatusEvent::MergeFailure {
                story_id: format!("story_{i}"),
                story_name: String::new(),
                reason: "test".to_string(),
            });
        }

        // Yield so the background receive task can process all events.
        sleep(Duration::from_millis(30)).await;

        // Model call counter must remain 0.
        assert_eq!(
            model_call_count.load(Ordering::Relaxed),
            0,
            "no model calls should occur when events arrive on the buffer"
        );

        // All 5 events must be in the buffer.
        let items = pool
            .drain_agent_status_buffer("story-z", "coder")
            .expect("agent must have a buffer");
        assert_eq!(items.len(), 5, "all 5 events should be buffered");
        for item in &items {
            assert!(
                matches!(item, BufferedItem::Event(_)),
                "every item must be a buffered event, not a side-effect artifact"
            );
        }
    }

    /// Story 736 — Criterion 3 end-to-end: fire N status events at an idle
    /// agent session; verify the agent is not invoked; then drain and format the
    /// buffer to simulate what `start_agent` does on the next user message;
    /// assert all N events appear in the `<recent-events>` block and the buffer
    /// is empty afterwards.
    #[tokio::test]
    async fn drain_and_format_prepends_all_buffered_events_and_empties_buffer() {
        use crate::service::status::buffer::BufferedItem;
        use crate::service::status::buffer::format_buffered_items;

        let pool = AgentPool::new_test(3207);
        // Idle agent with a live buffer — simulates a previous session that is
        // now completed but has been accumulating pipeline events.
        pool.inject_test_agent_with_live_buffer("story-736", "coder", AgentStatus::Completed);

        // Fire N status events at the pool's broadcaster while the agent is idle.
        // The agent must NOT be invoked — events only accumulate in the buffer.
        let bc = pool.status_broadcaster();
        let n = 4_usize;
        for i in 0..n {
            bc.publish(StatusEvent::MergeFailure {
                story_id: format!("other-{i}"),
                story_name: String::new(),
                reason: format!("reason-{i}"),
            });
        }
        // Yield so the background receive tasks can process all events.
        sleep(Duration::from_millis(30)).await;

        // Drain the buffer — this is what start_agent does before spawning the
        // next agent turn (story 736, criterion 1).
        let items = pool
            .drain_agent_status_buffer("story-736", "coder")
            .expect("agent must have a buffer");
        assert_eq!(items.len(), n, "all {n} events must be buffered");
        assert!(
            items.iter().all(|i| matches!(i, BufferedItem::Event(_))),
            "all items must be events (no truncation)"
        );

        // Format into a <recent-events> block (criterion 1).
        let block = format_buffered_items(&items).expect("non-empty slice must yield a block");
        assert!(
            block.starts_with("<recent-events>\n"),
            "block must open with <recent-events>: {block}"
        );
        assert!(
            block.ends_with("\n</recent-events>"),
            "block must close with </recent-events>: {block}"
        );
        for i in 0..n {
            assert!(
                block.contains(&format!("reason-{i}")),
                "block must include content from event {i}: {block}"
            );
        }

        // After drain the buffer must be empty (criterion 2).
        let after = pool
            .drain_agent_status_buffer("story-736", "coder")
            .expect("agent must still have a buffer");
        assert!(
            after.is_empty(),
            "buffer must be empty after drain (got {} items)",
            after.len()
        );
    }

    /// Criterion 2: on session restart a fresh buffer is created, clearing any
    /// events accumulated in the previous session.
    ///
    /// We simulate restart by replacing the agent entry (as `start_agent` does
    /// for a Completed/Failed agent) and confirming the new buffer starts empty
    /// even though the old buffer had events.
    #[tokio::test]
    async fn restart_replaces_buffer_with_fresh_subscription() {
        let pool = AgentPool::new_test(3206);

        // First session: agent with a live buffer that accumulates an event.
        pool.inject_test_agent_with_live_buffer("story-restart", "coder", AgentStatus::Completed);
        pool.status_broadcaster()
            .publish(stage_transition_event("story-restart"));
        sleep(Duration::from_millis(20)).await;

        // Confirm first session's buffer has the event.
        let first_items = pool
            .drain_agent_status_buffer("story-restart", "coder")
            .expect("first session must have buffer");
        assert_eq!(
            first_items.len(),
            1,
            "first session buffer should have 1 event"
        );

        // Simulate restart: re-inject the agent (as start_agent overwrites the
        // entry with a fresh StoryAgent including a new StatusEventBuffer).
        pool.inject_test_agent_with_live_buffer("story-restart", "coder", AgentStatus::Running);

        // The new session's buffer must be empty (no carry-over from first session).
        let second_items = pool
            .drain_agent_status_buffer("story-restart", "coder")
            .expect("second session must have buffer");
        assert!(
            second_items.is_empty(),
            "restarted session buffer must start empty (got {} items)",
            second_items.len()
        );
    }
}
