wsb 0.155.49

use grep::regex::RegexMatcher;
use grep::searcher::{sinks::UTF8, SearcherBuilder};
use ignore::WalkBuilder;
use serde::Serialize;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::mpsc;

use crate::web::db::SessionDb;

// ── Public types ─────────────────────────────────────────────────────────────

/// A single result item, tagged by its origin.
#[derive(Serialize, Clone, Debug)]
#[serde(tag = "kind", rename_all = "snake_case")]
pub enum SearchHit {
    Session {
        session_id: String,
        project: String,
        agent: String,
        name: String,
        ordinal: i64,
        session_ended_at: Option<String>,
        snippet: String,
    },
    File {
        /// Path relative to the workspace root.
        path: String,
        line_number: u64,
        snippet: String,
    },
}

/// Events sent from a source to the SSE handler.
#[derive(Serialize, Clone, Debug)]
#[serde(tag = "type", rename_all = "snake_case")]
pub enum SearchEvent {
    Hits { source: String, hits: Vec<SearchHit> },
    Done { source: String },
    Error { source: String, message: String },
}

// ── Cancellation guard ────────────────────────────────────────────────────────

/// Sets the cancel flag when dropped.  Embed in the SSE stream generator so
/// that when the client disconnects (stream dropped) all running sources see
/// the flag and stop at their next checkpoint.
pub struct CancelOnDrop(pub Arc<AtomicBool>);

impl Drop for CancelOnDrop {
    fn drop(&mut self) {
        self.0.store(true, Ordering::Relaxed);
    }
}

// ── SessionSource ─────────────────────────────────────────────────────────────

/// Searches indexed session messages via SQLite FTS5.
/// Must be run as a tokio::spawn (async task), not spawn_blocking.
pub struct SessionSource {
    db: SessionDb,
    sort: Option<String>,
    project: Option<String>,
}

impl SessionSource {
    pub fn new(db: SessionDb, sort: Option<String>, project: Option<String>) -> Self {
        Self { db, sort, project }
    }

    pub async fn run(
        self,
        query: String,
        tx: mpsc::Sender<SearchEvent>,
        cancelled: Arc<AtomicBool>,
    ) {
        if cancelled.load(Ordering::Relaxed) {
            return;
        }

        let sort = self.sort.as_deref().unwrap_or("relevance");
        let project = self.project.as_deref();
        let result = self.db.global_search(&query, 50, 0, sort, project).await;

        if cancelled.load(Ordering::Relaxed) {
            return;
        }

        match result {
            Ok(rows) => {
                let hits: Vec<SearchHit> = rows
                    .into_iter()
                    .map(|r| SearchHit::Session {
                        session_id: r.session_id,
                        project: r.project,
                        agent: r.agent,
                        name: r.name,
                        ordinal: r.ordinal,
                        session_ended_at: r.session_ended_at,
                        snippet: r.snippet,
                    })
                    .collect();

                if !hits.is_empty() {
                    let _ = tx.send(SearchEvent::Hits {
                        source: "sessions".into(),
                        hits,
                    }).await;
                }
                let _ = tx.send(SearchEvent::Done {
                    source: "sessions".into(),
                }).await;
            }
            Err(e) => {
                let _ = tx.send(SearchEvent::Error {
                    source: "sessions".into(),
                    message: e.to_string(),
                }).await;
            }
        }
    }
}

// ── FileSource ────────────────────────────────────────────────────────────────

/// Searches source files in the workspace using ripgrep internals.
/// Must be run via tokio::task::spawn_blocking — it performs blocking I/O
/// throughout.  Does not need the async runtime inside the blocking thread.
///
/// The query is treated as a literal string (not a regex) so users don't need
/// to escape special characters.  Results stream file-by-file.
pub struct FileSource {
    workspace_root: PathBuf,
}

impl FileSource {
    pub fn new(workspace_root: PathBuf) -> Self {
        Self { workspace_root }
    }

    pub fn run(
        self,
        query: String,
        tx: mpsc::Sender<SearchEvent>,
        cancelled: Arc<AtomicBool>,
    ) {
        if cancelled.load(Ordering::Relaxed) {
            return;
        }

        // Escape the query so it matches literally, not as a regex.
        let pattern = regex::escape(&query);
        let matcher = match RegexMatcher::new_line_matcher(&pattern) {
            Ok(m) => m,
            Err(e) => {
                let _ = tx.blocking_send(SearchEvent::Error {
                    source: "files".into(),
                    message: format!("invalid search pattern: {}", e),
                });
                return;
            }
        };

        let walker = WalkBuilder::new(&self.workspace_root)
            .standard_filters(true)
            .build();

        let mut searcher = SearcherBuilder::new().line_number(true).build();

        for entry in walker {
            if cancelled.load(Ordering::Relaxed) {
                break;
            }

            let entry = match entry {
                Ok(e) => e,
                Err(_) => continue,
            };

            if !entry.file_type().map(|ft| ft.is_file()).unwrap_or(false) {
                continue;
            }

            let path = entry.path();
            let rel_path = path
                .strip_prefix(&self.workspace_root)
                .unwrap_or(path)
                .to_string_lossy()
                .into_owned();

            let mut file_hits: Vec<SearchHit> = Vec::new();

            // search_path errors (binary files, permission denied) are silently
            // skipped — they are not actionable at this level.
            let _ = searcher.search_path(
                &matcher,
                path,
                UTF8(|line_number, line| {
                    file_hits.push(SearchHit::File {
                        path: rel_path.clone(),
                        line_number,
                        snippet: line.trim_end_matches('\n').to_string(),
                    });
                    Ok(true)
                }),
            );

            if !file_hits.is_empty() {
                if tx
                    .blocking_send(SearchEvent::Hits {
                        source: "files".into(),
                        hits: file_hits,
                    })
                    .is_err()
                {
                    // Consumer is gone (client disconnected), no point continuing.
                    return;
                }
            }
        }

        let _ = tx.blocking_send(SearchEvent::Done {
            source: "files".into(),
        });
    }
}

// ── Tests ─────────────────────────────────────────────────────────────────────

#[cfg(test)]
mod tests {
    use super::*;
    use tempfile::TempDir;
    use tokio::sync::mpsc;

    fn run_file_source(
        source: FileSource,
        query: &str,
        cancelled: Arc<AtomicBool>,
    ) -> Vec<SearchEvent> {
        let (tx, mut rx) = mpsc::channel(256);
        let q = query.to_string();
        // FileSource::run is blocking — run it on a dedicated thread with no
        // async runtime involvement.
        let handle = std::thread::spawn(move || source.run(q, tx, cancelled));
        handle.join().unwrap();
        let mut events = Vec::new();
        while let Ok(e) = rx.try_recv() {
            events.push(e);
        }
        events
    }

    // ── FileSource tests ──────────────────────────────────────────────────────

    #[test]
    fn file_source_finds_literal_matches() {
        let dir = TempDir::new().unwrap();
        let file = dir.path().join("hello.txt");
        std::fs::write(&file, "foo bar\nbaz qux\nfoo again\n").unwrap();

        let source = FileSource::new(dir.path().to_path_buf());
        let cancelled = Arc::new(AtomicBool::new(false));
        let events = run_file_source(source, "foo", cancelled);

        let hits: Vec<&SearchHit> = events
            .iter()
            .filter_map(|e| match e {
                SearchEvent::Hits { hits, .. } => Some(hits.iter()),
                _ => None,
            })
            .flatten()
            .collect();

        assert_eq!(hits.len(), 2, "expected two 'foo' matches");

        let done_count = events
            .iter()
            .filter(|e| matches!(e, SearchEvent::Done { source } if source == "files"))
            .count();
        assert_eq!(done_count, 1, "expected exactly one Done event");
    }

    #[test]
    fn file_source_escapes_regex_metacharacters() {
        let dir = TempDir::new().unwrap();
        // File contains a literal dot — should be matched literally, not as
        // a regex wildcard (which would match any character).
        let file = dir.path().join("dots.txt");
        std::fs::write(&file, "1.2.3\nxyz\n").unwrap();

        let source = FileSource::new(dir.path().to_path_buf());
        let cancelled = Arc::new(AtomicBool::new(false));
        let events = run_file_source(source, "1.2.3", cancelled);

        let hits: Vec<&SearchHit> = events
            .iter()
            .filter_map(|e| match e {
                SearchEvent::Hits { hits, .. } => Some(hits.iter()),
                _ => None,
            })
            .flatten()
            .collect();

        // "xyz" does NOT match "1.2.3" literally, so only one hit.
        assert_eq!(hits.len(), 1);
        match &hits[0] {
            SearchHit::File { snippet, .. } => assert!(snippet.contains("1.2.3")),
            other => panic!("expected File hit, got {:?}", other),
        }
    }

    #[test]
    fn file_source_respects_cancellation() {
        let dir = TempDir::new().unwrap();
        // Write enough files that we'd normally get many hits.
        for i in 0..20 {
            let path = dir.path().join(format!("file{}.txt", i));
            std::fs::write(&path, "needle\nneedle\nneedle\n").unwrap();
        }

        let cancelled = Arc::new(AtomicBool::new(true)); // pre-cancelled
        let source = FileSource::new(dir.path().to_path_buf());
        let events = run_file_source(source, "needle", cancelled);

        // No Hits should have been sent because the source checks cancellation
        // before the walk loop.
        let hit_count: usize = events
            .iter()
            .filter(|e| matches!(e, SearchEvent::Hits { .. }))
            .count();
        assert_eq!(hit_count, 0, "pre-cancelled source must send no hits");
    }

    #[test]
    fn file_source_sends_done_even_with_no_matches() {
        let dir = TempDir::new().unwrap();
        std::fs::write(dir.path().join("a.txt"), "nothing here\n").unwrap();

        let source = FileSource::new(dir.path().to_path_buf());
        let cancelled = Arc::new(AtomicBool::new(false));
        let events = run_file_source(source, "UNLIKELY_TOKEN_XYZ", cancelled);

        let done_count = events
            .iter()
            .filter(|e| matches!(e, SearchEvent::Done { .. }))
            .count();
        assert_eq!(done_count, 1, "Done must be sent even when no matches found");

        let hit_count = events
            .iter()
            .filter(|e| matches!(e, SearchEvent::Hits { .. }))
            .count();
        assert_eq!(hit_count, 0);
    }

    #[test]
    fn file_source_relative_paths() {
        let dir = TempDir::new().unwrap();
        let sub = dir.path().join("subdir");
        std::fs::create_dir(&sub).unwrap();
        std::fs::write(sub.join("code.rs"), "let x = target;\n").unwrap();

        let source = FileSource::new(dir.path().to_path_buf());
        let cancelled = Arc::new(AtomicBool::new(false));
        let events = run_file_source(source, "target", cancelled);

        let paths: Vec<String> = events
            .iter()
            .filter_map(|e| match e {
                SearchEvent::Hits { hits, .. } => Some(hits.iter()),
                _ => None,
            })
            .flatten()
            .filter_map(|h| match h {
                SearchHit::File { path, .. } => Some(path.clone()),
                _ => None,
            })
            .collect();

        assert_eq!(paths.len(), 1);
        // Path must be relative, not absolute.
        assert!(
            !paths[0].starts_with('/'),
            "path must be relative, got: {}",
            paths[0]
        );
        assert!(paths[0].contains("subdir"), "path must include subdir");
    }
}