claudectl 0.22.0

#![allow(dead_code)]

use std::collections::{HashMap, HashSet};
use std::sync::mpsc::{self, Receiver, Sender};
use std::time::Instant;

use crate::config::BrainConfig;
use crate::rules::{self, RuleAction, RuleMatch};
use crate::session::{ClaudeSession, SessionStatus};

use super::client::BrainSuggestion;
use super::context;

/// Result sent back from inference thread.
pub struct BrainResult {
    pub pid: u32,
    pub suggestion: Result<BrainSuggestion, String>,
}

/// The brain inference engine. Manages async inference threads and collects results.
pub struct BrainEngine {
    config: BrainConfig,
    tx: Sender<BrainResult>,
    rx: Receiver<BrainResult>,
    /// PIDs currently being inferred (prevents duplicate requests).
    inflight: HashSet<u32>,
    /// Per-PID cooldown to avoid hammering the LLM.
    cooldown: HashMap<u32, Instant>,
    /// Pending suggestions waiting for user confirmation (advisory mode).
    pub pending: HashMap<u32, BrainSuggestion>,
    /// Last time orchestration evaluation ran.
    last_orchestrate: Option<Instant>,
    /// Whether an orchestration inference is in-flight.
    orchestrate_inflight: bool,
}

const COOLDOWN_SECS: u64 = 10;

impl BrainEngine {
    pub fn new(config: BrainConfig) -> Self {
        let (tx, rx) = mpsc::channel();
        Self {
            config,
            tx,
            rx,
            inflight: HashSet::new(),
            cooldown: HashMap::new(),
            pending: HashMap::new(),
            last_orchestrate: None,
            orchestrate_inflight: false,
        }
    }

    /// Run one tick of the brain engine. Call this from app.tick() after refresh().
    ///
    /// 1. Collect results from completed inference threads
    /// 2. Spawn new inference threads for eligible sessions
    ///
    /// Returns a list of (pid, status_message) for actions taken this tick.
    pub fn tick(
        &mut self,
        sessions: &[ClaudeSession],
        deny_rules: &[crate::rules::AutoRule],
    ) -> Vec<(u32, String)> {
        let mut actions = Vec::new();

        // Phase 1: Collect results from completed inferences
        while let Ok(result) = self.rx.try_recv() {
            // PID 0 = orchestration result
            if result.pid == 0 {
                if let Ok(suggestion) = result.suggestion {
                    let orch_actions = self.handle_orchestration_result(&suggestion, sessions);
                    actions.extend(orch_actions);
                }
                continue;
            }

            self.inflight.remove(&result.pid);
            self.cooldown.insert(result.pid, Instant::now());

            match result.suggestion {
                Ok(suggestion) => {
                    // Check if a deny rule overrides the brain
                    let session = sessions.iter().find(|s| s.pid == result.pid);
                    if let Some(session) = session {
                        let deny_match = rules::evaluate(deny_rules, session);
                        if let Some(dm) = &deny_match {
                            if dm.action == RuleAction::Deny {
                                actions.push((
                                    result.pid,
                                    format!(
                                        "Brain suggested {}, but deny rule '{}' overrides",
                                        suggestion.action.label(),
                                        dm.rule_name,
                                    ),
                                ));
                                continue;
                            }
                        }
                    }

                    if self.config.auto_mode {
                        // Auto mode: execute immediately
                        if let Some(session) = session {
                            match &suggestion.action {
                                RuleAction::Route { target_pid } => {
                                    let target = sessions.iter().find(|s| s.pid == *target_pid);
                                    if let Some(target) = target {
                                        match self.execute_route(session, target) {
                                            Ok(msg) => actions.push((result.pid, msg)),
                                            Err(e) => actions
                                                .push((result.pid, format!("Route error: {e}"))),
                                        }
                                    } else {
                                        actions.push((
                                            result.pid,
                                            format!(
                                                "Route error: target PID {} not found",
                                                target_pid
                                            ),
                                        ));
                                    }
                                }
                                RuleAction::Spawn { .. } => {
                                    // Enforce max_sessions limit
                                    if sessions.len() >= self.config.max_sessions {
                                        actions.push((
                                            result.pid,
                                            format!(
                                                "Spawn blocked: {} sessions active (max {})",
                                                sessions.len(),
                                                self.config.max_sessions
                                            ),
                                        ));
                                    } else {
                                        let rule_match = suggestion_to_rule_match(&suggestion);
                                        match rules::execute(&rule_match, session) {
                                            Ok(msg) => actions.push((result.pid, msg)),
                                            Err(e) => actions
                                                .push((result.pid, format!("Spawn error: {e}"))),
                                        }
                                    }
                                }
                                _ => {
                                    let rule_match = suggestion_to_rule_match(&suggestion);
                                    match rules::execute(&rule_match, session) {
                                        Ok(msg) => actions.push((result.pid, msg)),
                                        Err(e) => {
                                            actions.push((result.pid, format!("Brain error: {e}")))
                                        }
                                    }
                                }
                            }
                        }
                    } else {
                        // Advisory mode: store for user confirmation
                        self.pending.insert(result.pid, suggestion);
                    }
                }
                Err(e) => {
                    crate::logger::log(
                        "BRAIN",
                        &format!("Inference failed for PID {}: {e}", result.pid),
                    );
                }
            }
        }

        // Phase 2: Spawn inference for eligible sessions
        for session in sessions {
            if !matches!(
                session.status,
                SessionStatus::NeedsInput | SessionStatus::WaitingInput
            ) {
                continue;
            }

            if self.inflight.contains(&session.pid) {
                continue;
            }

            if let Some(last) = self.cooldown.get(&session.pid) {
                if last.elapsed().as_secs() < COOLDOWN_SECS {
                    continue;
                }
            }

            // Already have a pending suggestion for this PID
            if self.pending.contains_key(&session.pid) {
                continue;
            }

            self.spawn_inference(session, sessions);
        }

        // Phase 3: Orchestration evaluation (less frequent)
        let orch_actions = self.maybe_orchestrate(sessions);
        actions.extend(orch_actions);

        actions
    }

    fn spawn_inference(&mut self, session: &ClaudeSession, all_sessions: &[ClaudeSession]) {
        let pid = session.pid;
        let config = self.config.clone();
        let tx = self.tx.clone();

        // Build context on the main thread (reads JSONL files)
        let mut brain_ctx =
            context::build_context(session, all_sessions, config.max_context_tokens);

        // Inject few-shot examples from past decisions
        if config.few_shot_count > 0 {
            let similar = super::decisions::retrieve_similar(
                session.pending_tool_name.as_deref(),
                session.display_name(),
                config.few_shot_count,
            );
            brain_ctx.few_shot_examples = super::decisions::format_few_shot_examples(&similar);
        }

        let prompt = context::format_brain_prompt(&brain_ctx);

        self.inflight.insert(pid);

        std::thread::spawn(move || {
            let suggestion = super::client::infer(&config, &prompt);
            let _ = tx.send(BrainResult { pid, suggestion });
        });
    }

    /// Execute a route: read source's recent transcript, summarize via LLM,
    /// and either send directly (if target is waiting) or queue in mailbox.
    fn execute_route(
        &self,
        source: &ClaudeSession,
        target: &ClaudeSession,
    ) -> Result<String, String> {
        // Build source context to get recent transcript
        let source_ctx = context::build_context(
            source,
            std::slice::from_ref(source),
            self.config.max_context_tokens,
        );

        // Summarize for target's task
        let summary = super::client::summarize_for_routing(
            &self.config,
            &source_ctx.recent_transcript,
            source.display_name(),
            target.display_name(),
        )?;

        // If target is waiting for input, deliver directly; otherwise queue in mailbox
        if target.status == SessionStatus::WaitingInput {
            rules::execute_route(source, target, &summary, "brain")
        } else {
            super::mailbox::enqueue(source.pid, source.display_name(), target.pid, &summary);
            Ok(format!(
                "Brain: queued message from {} → {} (mailbox, target is {})",
                source.display_name(),
                target.display_name(),
                target.status,
            ))
        }
    }

    /// Accept a pending brain suggestion (user pressed 'b').
    pub fn accept(&mut self, pid: u32, session: &ClaudeSession) -> Option<String> {
        let suggestion = self.pending.remove(&pid)?;
        let rule_match = suggestion_to_rule_match(&suggestion);
        match rules::execute(&rule_match, session) {
            Ok(msg) => Some(msg),
            Err(e) => Some(format!("Brain execute error: {e}")),
        }
    }

    /// Reject a pending brain suggestion (user pressed 'B').
    pub fn reject(&mut self, pid: u32) -> Option<BrainSuggestion> {
        self.pending.remove(&pid)
    }

    /// Clear pending suggestions for PIDs that are no longer in NeedsInput/WaitingInput.
    pub fn cleanup(&mut self, sessions: &[ClaudeSession]) {
        let active_pids: HashSet<u32> = sessions.iter().map(|s| s.pid).collect();
        self.pending.retain(|pid, _| {
            active_pids.contains(pid)
                && sessions.iter().any(|s| {
                    s.pid == *pid
                        && matches!(
                            s.status,
                            SessionStatus::NeedsInput | SessionStatus::WaitingInput
                        )
                })
        });
        self.inflight.retain(|pid| active_pids.contains(pid));
    }

    /// Run orchestration evaluation: ask the brain if any cross-session actions
    /// should be taken (spawn, route, terminate). Runs less frequently than
    /// per-session advisory (every orchestrate_interval_secs).
    pub fn maybe_orchestrate(&mut self, sessions: &[ClaudeSession]) -> Vec<(u32, String)> {
        if !self.config.orchestrate || !self.config.auto_mode {
            return Vec::new();
        }

        if sessions.len() < 2 {
            return Vec::new();
        }

        // Check interval
        let interval = std::time::Duration::from_secs(self.config.orchestrate_interval_secs);
        if let Some(last) = self.last_orchestrate {
            if last.elapsed() < interval {
                return Vec::new();
            }
        }

        if self.orchestrate_inflight {
            return Vec::new();
        }

        self.last_orchestrate = Some(Instant::now());
        self.orchestrate_inflight = true;

        // Build orchestration prompt with all sessions
        let prompt = build_orchestration_prompt(sessions, &self.config);
        let config = self.config.clone();
        let tx = self.tx.clone();

        // Use PID 0 as sentinel for orchestration results
        std::thread::spawn(move || {
            let suggestion = super::client::infer(&config, &prompt);
            let _ = tx.send(BrainResult { pid: 0, suggestion });
        });

        Vec::new()
    }

    /// Check if a result is an orchestration response (pid == 0).
    pub fn handle_orchestration_result(
        &mut self,
        suggestion: &BrainSuggestion,
        sessions: &[ClaudeSession],
    ) -> Vec<(u32, String)> {
        self.orchestrate_inflight = false;
        let mut actions = Vec::new();

        // The orchestration response may suggest multiple actions.
        // For now, handle the primary action.
        match &suggestion.action {
            RuleAction::Spawn { .. } => {
                if sessions.len() >= self.config.max_sessions {
                    actions.push((
                        0,
                        format!(
                            "Orchestrate: spawn blocked ({} sessions, max {})",
                            sessions.len(),
                            self.config.max_sessions
                        ),
                    ));
                } else {
                    let rule_match = suggestion_to_rule_match(suggestion);
                    // Need a dummy session for execute — use first available
                    if let Some(session) = sessions.first() {
                        match rules::execute(&rule_match, session) {
                            Ok(msg) => actions.push((0, format!("Orchestrate: {msg}"))),
                            Err(e) => actions.push((0, format!("Orchestrate error: {e}"))),
                        }
                    }
                }
            }
            RuleAction::Route { target_pid } => {
                // Find source (most recently active) and target
                if let Some(target) = sessions.iter().find(|s| s.pid == *target_pid) {
                    if let Some(source) = sessions
                        .iter()
                        .find(|s| s.pid != *target_pid && s.status == SessionStatus::WaitingInput)
                    {
                        match self.execute_route(source, target) {
                            Ok(msg) => actions.push((0, format!("Orchestrate: {msg}"))),
                            Err(e) => actions.push((0, format!("Orchestrate error: {e}"))),
                        }
                    }
                }
            }
            RuleAction::Terminate => {
                // Orchestration terminate — brain should include which PID in reasoning
                actions.push((
                    0,
                    format!(
                        "Orchestrate: terminate suggested — {}",
                        suggestion.reasoning
                    ),
                ));
            }
            _ => {
                // approve/deny/send don't make sense at the orchestration level
                actions.push((
                    0,
                    format!(
                        "Orchestrate: {} — {}",
                        suggestion.action.label(),
                        suggestion.reasoning
                    ),
                ));
            }
        }

        actions
    }
}

/// Build the orchestration prompt from the prompt library.
fn build_orchestration_prompt(sessions: &[ClaudeSession], _config: &BrainConfig) -> String {
    let session_map = context::format_global_session_map_public(sessions);
    let template = super::prompts::load(super::prompts::ORCHESTRATION);
    super::prompts::expand(
        &template,
        &[
            ("session_count", &sessions.len().to_string()),
            ("session_map", &session_map),
        ],
    )
}

fn suggestion_to_rule_match(suggestion: &BrainSuggestion) -> RuleMatch {
    RuleMatch {
        rule_name: format!(
            "brain ({}% confidence)",
            (suggestion.confidence * 100.0) as u32
        ),
        action: suggestion.action.clone(),
        message: suggestion.message.clone(),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::session::{RawSession, TelemetryStatus};

    fn make_config() -> BrainConfig {
        BrainConfig {
            enabled: true,
            endpoint: "http://localhost:11434/api/generate".into(),
            model: "test".into(),
            auto_mode: false,
            timeout_ms: 1000,
            max_context_tokens: 1000,
            few_shot_count: 5,
            max_sessions: 10,
            orchestrate: false,
            orchestrate_interval_secs: 30,
        }
    }

    fn make_session(pid: u32, status: SessionStatus) -> ClaudeSession {
        let raw = RawSession {
            pid,
            session_id: "test".into(),
            cwd: "/tmp/test".into(),
            started_at: 0,
        };
        let mut s = ClaudeSession::from_raw(raw);
        s.status = status;
        s.telemetry_status = TelemetryStatus::Available;
        s.pending_tool_name = Some("Bash".into());
        s
    }

    #[test]
    fn engine_creates_without_panic() {
        let _engine = BrainEngine::new(make_config());
    }

    #[test]
    fn suggestion_to_rule_match_format() {
        let suggestion = BrainSuggestion {
            action: RuleAction::Approve,
            message: None,
            reasoning: "safe".into(),
            confidence: 0.95,
        };
        let rm = suggestion_to_rule_match(&suggestion);
        assert_eq!(rm.action, RuleAction::Approve);
        assert!(rm.rule_name.contains("95%"));
    }

    #[test]
    fn cleanup_removes_stale_pending() {
        let mut engine = BrainEngine::new(make_config());
        engine.pending.insert(
            999,
            BrainSuggestion {
                action: RuleAction::Approve,
                message: None,
                reasoning: "test".into(),
                confidence: 0.9,
            },
        );

        // PID 999 not in sessions list → should be cleaned up
        engine.cleanup(&[]);
        assert!(engine.pending.is_empty());
    }

    #[test]
    fn cleanup_keeps_active_pending() {
        let mut engine = BrainEngine::new(make_config());
        let session = make_session(100, SessionStatus::NeedsInput);
        engine.pending.insert(
            100,
            BrainSuggestion {
                action: RuleAction::Approve,
                message: None,
                reasoning: "test".into(),
                confidence: 0.9,
            },
        );

        engine.cleanup(&[session]);
        assert!(engine.pending.contains_key(&100));
    }

    #[test]
    fn reject_removes_and_returns_suggestion() {
        let mut engine = BrainEngine::new(make_config());
        engine.pending.insert(
            100,
            BrainSuggestion {
                action: RuleAction::Approve,
                message: None,
                reasoning: "test".into(),
                confidence: 0.9,
            },
        );

        let rejected = engine.reject(100);
        assert!(rejected.is_some());
        assert!(engine.pending.is_empty());
    }
}